chore: mypy compliance with 3.11 syntax

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

.pre-commit-config.yaml

@@ -38,6 +38,7 @@ repos:
     rev: 'v1.4.1'
     hooks:
     -   id: mypy
+        exclude: "^py/torch_tensorrt/fx"
   - repo: local
     hooks:
     -   id: dont-commit-upstream

py/torch_tensorrt/_Device.py

@@ -1,4 +1,11 @@
-from typing import Self, Optional, Any, Tuple
+from typing import TypeVar, Optional, Any, Tuple
+import sys
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
 import torch
 
 # from torch_tensorrt import _enums
@@ -25,7 +32,9 @@ class Device(object):
         allow_gpu_fallback (bool): Whether falling back to GPU if DLA cannot support an op should be allowed
     """
 
-    device_type: Optional[trt.DeviceType] = None  #: Target device type (GPU or DLA). Set implicitly based on if dla_core is specified.
+    device_type: Optional[
+        trt.DeviceType
+    ] = None  #: Target device type (GPU or DLA). Set implicitly based on if dla_core is specified.
     gpu_id: int = -1  #: Device ID for target GPU
     dla_core: int = -1  #: Core ID for target DLA core
     allow_gpu_fallback: bool = False  #: Whether falling back to GPU if DLA cannot support an op should be allowed
@@ -140,11 +149,7 @@ def _from_torch_device(cls, torch_dev: torch.device) -> Self:
 
     @classmethod
     def _current_device(cls) -> Self:
-        try:
-            dev = _C._get_current_device()
-        except RuntimeError:
-            logging.log(logging.Level.Error, "Cannot get current device")
-            return None
+        dev = _C._get_current_device()
         return cls(gpu_id=dev.gpu_id)
 
     @staticmethod

py/torch_tensorrt/_Input.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import List, Dict, Any, Tuple, Optional, Union
+from typing import List, Dict, Any, Tuple, Optional, Union, Sequence
 
 import torch
 
@@ -27,13 +27,17 @@ class _ShapeMode(Enum):
         STATIC = 0
         DYNAMIC = 1
 
-    shape_mode: Optional[_ShapeMode] = None  #: Is input statically or dynamically shaped
-    shape: Optional[Union[Tuple[int, ...], Dict[str, Tuple[int, ...]]]] = None  #: Either a single Tuple or a dict of tuples defining the input shape. Static shaped inputs will have a single tuple. Dynamic inputs will have a dict of the form ``{ "min_shape": Tuple, "opt_shape": Tuple, "max_shape": Tuple }``
-    dtype: _enums.dtype = ( # type: ignore[name-defined]
+    shape_mode: Optional[
+        _ShapeMode
+    ] = None  #: Is input statically or dynamically shaped
+    shape: Optional[
+        Tuple[int, ...] | Dict[str, Tuple[int, ...]]
+    ] = None  #: Either a single Tuple or a dict of tuples defining the input shape. Static shaped inputs will have a single tuple. Dynamic inputs will have a dict of the form ``{ "min_shape": Tuple, "opt_shape": Tuple, "max_shape": Tuple }``
+    dtype: _enums.dtype = (  # type: ignore[name-defined]
         _enums.dtype.unknown
     )  #: The expected data type of the input tensor (default: torch_tensorrt.dtype.float32)
     _explicit_set_dtype: bool = False
-    format: _enums.TensorFormat = ( # type: ignore[name-defined]
+    format: _enums.TensorFormat = (  # type: ignore[name-defined]
         _enums.TensorFormat.contiguous
     )  #: The expected format of the input tensor (default: torch_tensorrt.TensorFormat.NCHW)
 
@@ -187,7 +191,9 @@ def __str__(self) -> str:
                     str(self.tensor_domain[1]),
                 )
             else:
-                raise RuntimeError(f"Input shape is dynamic but shapes are not provided as dictionary (found: {self.shape})")
+                raise RuntimeError(
+                    f"Input shape is dynamic but shapes are not provided as dictionary (found: {self.shape})"
+                )
         else:
             raise RuntimeError("Unknown input shape mode")
 
@@ -203,7 +209,7 @@ def _supported_input_size_type(input_size: Any) -> bool:
             return False
 
     @staticmethod
-    def _parse_dtype(dtype: Any) -> _enums.dtype: # type: ignore[name-defined]
+    def _parse_dtype(dtype: Any) -> _enums.dtype:  # type: ignore[name-defined]
         if isinstance(dtype, torch.dtype):
             if dtype == torch.long:
                 return _enums.dtype.long
@@ -231,7 +237,7 @@ def _parse_dtype(dtype: Any) -> _enums.dtype: # type: ignore[name-defined]
             )
 
     @staticmethod
-    def _to_torch_dtype(dtype: _enums.dtype) -> torch.dtype: # type: ignore[name-defined]
+    def _to_torch_dtype(dtype: _enums.dtype) -> torch.dtype:  # type: ignore[name-defined]
         if dtype == _enums.dtype.long:
             return torch.long
         elif dtype == _enums.dtype.int32:
@@ -250,7 +256,7 @@ def is_trt_dtype(self) -> bool:
         return bool(self.dtype != _enums.dtype.long)
 
     @staticmethod
-    def _parse_format(format: Any) -> _enums.TensorFormat: # type: ignore[name-defined]
+    def _parse_format(format: Any) -> _enums.TensorFormat:  # type: ignore[name-defined]
         if isinstance(format, torch.memory_format):
             if format == torch.contiguous_format:
                 return _enums.TensorFormat.contiguous
@@ -270,7 +276,9 @@ def _parse_format(format: Any) -> _enums.TensorFormat: # type: ignore[name-defin
             )
 
     @staticmethod
-    def _parse_tensor_domain(domain: Optional[Tuple[float, float]]) -> Tuple[float, float]:
+    def _parse_tensor_domain(
+        domain: Optional[Tuple[float, float]]
+    ) -> Tuple[float, float]:
         """
         Produce a tuple of integers which specifies a tensor domain in the interval format: [lo, hi)
 
@@ -349,7 +357,7 @@ def from_tensor(
 
     @classmethod
     def from_tensors(
-        cls, ts: torch.Tensor, disable_memory_format_check: bool = False
+        cls, ts: Sequence[torch.Tensor], disable_memory_format_check: bool = False
     ) -> List["Input"]:
         """
         Produce a list of Inputs which contain
@@ -369,7 +377,9 @@ def from_tensors(
             for t in ts
         ]
 
-    def example_tensor(self, optimization_profile_field: Optional[str] = None) -> Optional[torch.Tensor]:
+    def example_tensor(
+        self, optimization_profile_field: Optional[str] = None
+    ) -> torch.Tensor:
         """
         Get an example tensor of the shape specified by the Input object
 
@@ -388,7 +398,9 @@ def example_tensor(self, optimization_profile_field: Optional[str] = None) -> Op
                 if isinstance(self.shape, tuple):
                     return torch.rand(self.shape).to(dtype=self.torch_dtype)
                 else:
-                    RuntimeError(f"Input shape is dynamic but shapes are not provided as sequence (found: {self.shape})")
+                    RuntimeError(
+                        f"Input shape is dynamic but shapes are not provided as sequence (found: {self.shape})"
+                    )
         else:
             if optimization_profile_field is not None:
                 try:
@@ -408,11 +420,12 @@ def example_tensor(self, optimization_profile_field: Optional[str] = None) -> Op
                         dtype=self.torch_dtype
                     )
                 else:
-                    raise RuntimeError(f"Input shape is dynamic but shapes are not provided as dictionary (found: {self.shape})")
+                    raise RuntimeError(
+                        f"Input shape is dynamic but shapes are not provided as dictionary (found: {self.shape})"
+                    )
 
             else:
                 raise ValueError(
                     "Requested an example tensor from a dynamic shaped input but did not specific which profile field to use."
                 )
-        return None
-
+        raise

py/torch_tensorrt/_compile.py

@@ -2,7 +2,9 @@
 
 import torch_tensorrt.ts
 
-from torch_tensorrt import logging, Input, dtype
+from torch_tensorrt import logging
+from torch_tensorrt._Input import Input
+from torch_tensorrt._enums import dtype
 import torch
 import torch.fx
 from enum import Enum
@@ -12,12 +14,18 @@
 from torch_tensorrt.fx.utils import LowerPrecision
 
 
-def _non_fx_input_interface(inputs: List[Input | torch.Tensor | InputTensorSpec]) -> TypeGuard[List[Input | torch.Tensor]]:
+def _non_fx_input_interface(
+    inputs: List[Input | torch.Tensor | InputTensorSpec],
+) -> TypeGuard[List[Input | torch.Tensor]]:
     return all([isinstance(i, torch.Tensor | Input) for i in inputs])
 
-def _fx_input_interface(inputs: List[Input | torch.Tensor | InputTensorSpec]) -> TypeGuard[List[InputTensorSpec | torch.Tensor]]:
+
+def _fx_input_interface(
+    inputs: List[Input | torch.Tensor | InputTensorSpec],
+) -> TypeGuard[List[InputTensorSpec | torch.Tensor]]:
     return all([isinstance(i, torch.Tensor | InputTensorSpec) for i in inputs])
 
+
 class _IRType(Enum):
     """Enum to set the minimum required logging level to print a message to stdout"""
 
@@ -89,10 +97,12 @@ def _get_target_ir(module_type: _ModuleType, ir: str) -> _IRType:
 def compile(
     module: Any,
     ir: str = "default",
-    inputs: List[Union[Input, torch.Tensor, InputTensorSpec]] = [],
-    enabled_precisions: Set[Union[torch.dtype, dtype]] = set([torch.float]),
+    inputs: List[Input | torch.Tensor | InputTensorSpec] = [],
+    enabled_precisions: Set[torch.dtype | dtype] = set([torch.float]),
     **kwargs: Any,
-) -> Union[torch.nn.Module, torch.jit.ScriptModule, torch.fx.GraphModule, Callable[[Any], Any]]:
+) -> (
+    torch.nn.Module | torch.jit.ScriptModule | torch.fx.GraphModule | Callable[..., Any]
+):
     """Compile a PyTorch module for NVIDIA GPUs using TensorRT
 
     Takes a existing PyTorch module and a set of settings to configure the compiler
@@ -168,7 +178,7 @@ def compile(
         )
         return compiled_fx_module
     elif target_ir == _IRType.dynamo:
-        return torch_tensorrt.dynamo.compile(
+        compiled_aten_module: torch.fx.GraphModule = torch_tensorrt.dynamo.compile(
             module,
             inputs=inputs,
             enabled_precisions=enabled_precisions,
@@ -198,8 +208,8 @@ def convert_method_to_trt_engine(
     module: Any,
     method_name: str,
     ir: str = "default",
-    inputs: List[Union[Input, torch.Tensor]] = [],
-    enabled_precisions: Set[Union[torch.dtype, dtype]] = set([torch.float]),
+    inputs: List[Input | torch.Tensor] = [],
+    enabled_precisions: Set[torch.dtype | dtype] = set([torch.float]),
     **kwargs: Any,
 ) -> bytes:
     """Convert a TorchScript module method to a serialized TensorRT engine

py/torch_tensorrt/dynamo/_defaults.py

@@ -8,5 +8,5 @@
 MAX_AUX_STREAMS = None
 VERSION_COMPATIBLE = False
 OPTIMIZATION_LEVEL = None
-USE_PYTHON_RUNTIME = None
 TRUNCATE_LONG_AND_DOUBLE = False
+USE_PYTHON_RUNTIME = False

py/torch_tensorrt/dynamo/aten_tracer.py

@@ -45,7 +45,6 @@ def __init__(
         specialize_int: bool = True,
         verbose: bool = True,
     ) -> None:
-
         self.capture_scalar_outputs = capture_scalar_outputs
         self.guard_nn_modules = guard_nn_modules
         self.dynamic_shapes = dynamic_shapes
@@ -128,7 +127,11 @@ def dynamo_trace(
 
 
 @req_torch_version("2.dev")
-def trace(model: Union[torch.nn.Module, torch.fx.GraphModule], inputs: Tuple[Any, ...], **kwargs: Any) -> torch.fx.GraphModule:
+def trace(
+    model: torch.nn.Module | torch.fx.GraphModule,
+    inputs: Tuple[Any, ...],
+    **kwargs: Any,
+) -> torch.fx.GraphModule:
     """
     Optimized trace with necessary passes which re-compose some ops or replace some ops
     These passes should be general and functional purpose
@@ -149,11 +152,11 @@ def trace(model: Union[torch.nn.Module, torch.fx.GraphModule], inputs: Tuple[Any
     fx_module, __package__ = dynamo_trace(model, inputs, True, "symbolic")
     print(fx_module.graph)
     for passes in passes_list:
-        pr: PassResult = passes(fx_module) #type: ignore[assignment] #The type hints in fx are wrong
+        pr: PassResult = passes(fx_module)  # type: ignore[assignment] #The type hints in fx are wrong
         fx_module = pr.graph_module
 
     fx_module(*inputs)
 
     fx_module = run_const_fold(fx_module)
     print(fx_module.graph)
-    return fx_module #type: ignore[no-any-return]
+    return fx_module  # type: ignore[no-any-return]

py/torch_tensorrt/dynamo/backend/backends.py

@@ -27,22 +27,19 @@
 logger = logging.getLogger(__name__)
 
 
-@td.register_backend(name="torch_tensorrt") # type: ignore[misc]
+@td.register_backend(name="torch_tensorrt")  # type: ignore[misc]
 def torch_tensorrt_backend(
-    gm: torch.fx.GraphModule,
-    sample_inputs: Sequence[torch.Tensor],
-    **kwargs: Any
+    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor], **kwargs: Any
 ) -> torch.nn.Module:
     DEFAULT_BACKEND = aot_torch_tensorrt_aten_backend
 
     compiled_mod: torch.nn.Module = DEFAULT_BACKEND(gm, sample_inputs, **kwargs)
     return compiled_mod
 
-@td.register_backend(name="aot_torch_tensorrt_aten") # type: ignore[misc]
+
+@td.register_backend(name="aot_torch_tensorrt_aten")  # type: ignore[misc]
 def aot_torch_tensorrt_aten_backend(
-    gm: torch.fx.GraphModule,
-    sample_inputs: Sequence[torch.Tensor],
-    **kwargs: Any
+    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor], **kwargs: Any
 ) -> torch.nn.Module:
     settings = parse_dynamo_kwargs(kwargs)
 
@@ -58,7 +55,7 @@ def aot_torch_tensorrt_aten_backend(
     return aot_module_simplified(
         gm,
         sample_inputs,
-        fw_compiler=make_boxed_compiler(custom_backend), # type: ignore[no-untyped-call]
+        fw_compiler=make_boxed_compiler(custom_backend),  # type: ignore[no-untyped-call]
         decompositions=get_decompositions(),
     )