chore(//py/torch_tensorrt/ts): Making compile mypy compliant

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

Author: narendasan

py/torch_tensorrt/_Device.py

@@ -1,3 +1,4 @@
+from typing import Self
 import torch
 
 # from torch_tensorrt import _enums
@@ -137,7 +138,7 @@ def _from_torch_device(cls, torch_dev: torch.device):
         return cls(gpu_id=gpu_id)
 
     @classmethod
-    def _current_device(cls):
+    def _current_device(cls) -> Self:
         try:
             dev = _C._get_current_device()
         except RuntimeError:

py/torch_tensorrt/ts/_compiler.py

@@ -1,4 +1,4 @@
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Tuple, Union, Set, Optional
 import torch
 from torch import nn
 
@@ -7,31 +7,32 @@
 from torch_tensorrt.ts._compile_spec import _parse_compile_spec, _parse_device
 from torch_tensorrt._Device import Device
 from types import FunctionType
+from torch_tensorrt._Input import Input
 
 
 def compile(
     module: torch.jit.ScriptModule,
-    inputs=[],
-    input_signature=None,
-    device=Device._current_device(),
-    disable_tf32=False,
-    sparse_weights=False,
-    enabled_precisions=set(),
-    refit=False,
-    debug=False,
-    capability=_enums.EngineCapability.default,
-    num_avg_timing_iters=1,
-    workspace_size=0,
-    dla_sram_size=1048576,
-    dla_local_dram_size=1073741824,
-    dla_global_dram_size=536870912,
-    calibrator=None,
-    truncate_long_and_double=False,
-    require_full_compilation=False,
-    min_block_size=3,
-    torch_executed_ops=[],
-    torch_executed_modules=[],
-    allow_shape_tensors=False,
+    inputs: List[Union[Input, torch.Tensor]] = [],
+    input_signature: Optional[Tuple[Union[Input, torch.Tensor]]] = None,
+    device: Device = Device._current_device(),
+    disable_tf32: bool = False,
+    sparse_weights: bool = False,
+    enabled_precisions: Set[Union[torch.dtype, _enums.dtype]] = set(), # type: ignore[name-defined]
+    refit: bool = False,
+    debug: bool = False,
+    capability: _enums.EngineCapability = _enums.EngineCapability.default, # type: ignore[name-defined]
+    num_avg_timing_iters: int = 1,
+    workspace_size: int = 0,
+    dla_sram_size: int = 1048576,
+    dla_local_dram_size: int = 1073741824,
+    dla_global_dram_size: int = 536870912,
+    calibrator: object = None,
+    truncate_long_and_double: bool = False,
+    require_full_compilation: bool = False,
+    min_block_size: int = 3,
+    torch_executed_ops: List[str] = [],
+    torch_executed_modules: List[str] = [],
+    allow_shape_tensors: bool = False,
 ) -> torch.jit.ScriptModule:
     """Compile a TorchScript module for NVIDIA GPUs using TensorRT
 
@@ -137,30 +138,30 @@ def compile(
     }
 
     compiled_cpp_mod = _C.compile_graph(module._c, _parse_compile_spec(spec))
-    compiled_module = torch.jit._recursive.wrap_cpp_module(compiled_cpp_mod)
+    compiled_module: torch.jit.ScriptModule = torch.jit._recursive.wrap_cpp_module(compiled_cpp_mod) # type: ignore[no-untyped-call]
     return compiled_module
 
 
 def convert_method_to_trt_engine(
     module: torch.jit.ScriptModule,
     method_name: str = "forward",
-    inputs=[],
-    device=Device._current_device(),
-    disable_tf32=False,
-    sparse_weights=False,
-    enabled_precisions=set(),
-    refit=False,
-    debug=False,
-    capability=_enums.EngineCapability.default,
-    num_avg_timing_iters=1,
-    workspace_size=0,
-    dla_sram_size=1048576,
-    dla_local_dram_size=1073741824,
-    dla_global_dram_size=536870912,
-    truncate_long_and_double=False,
-    calibrator=None,
-    allow_shape_tensors=False,
-) -> bytearray:
+    inputs: List[Union[Input, torch.Tensor]] = [],
+    device: Device = Device._current_device(),
+    disable_tf32: bool = False,
+    sparse_weights: bool = False,
+    enabled_precisions: Set[Union[torch.dtype, _enums.dtype]] = set(), # type: ignore[name-defined]
+    refit: bool = False,
+    debug: bool =False,
+    capability: _enums.EngineCapability = _enums.EngineCapability.default, # type: ignore[name-defined]
+    num_avg_timing_iters: int = 1,
+    workspace_size: int = 0,
+    dla_sram_size: int = 1048576,
+    dla_local_dram_size: int = 1073741824,
+    dla_global_dram_size: int = 536870912,
+    truncate_long_and_double: int = False,
+    calibrator: object = None,
+    allow_shape_tensors: bool = False,
+) -> bytes:
     """Convert a TorchScript module method to a serialized TensorRT engine
 
     Converts a specified method of a module to a serialized TensorRT engine given a dictionary of conversion settings
@@ -221,7 +222,7 @@ def convert_method_to_trt_engine(
         allow_shape_tensors: (Experimental) Allow aten::size to output shape tensors using IShapeLayer in TensorRT
 
     Returns:
-        bytearray: Serialized TensorRT engine, can either be saved to a file or deserialized via TensorRT APIs
+        bytes: Serialized TensorRT engine, can either be saved to a file or deserialized via TensorRT APIs
     """
     if isinstance(module, torch.jit.ScriptFunction):
         raise TypeError(
@@ -293,8 +294,8 @@ def embed_engine_in_new_module(
         input_binding_names,
         output_binding_names,
     )
-    return torch.jit._recursive.wrap_cpp_module(cpp_mod)
-
+    wrapped_mod: torch.jit.ScriptModule = torch.jit._recursive.wrap_cpp_module(cpp_mod) # type: ignore[no-untyped-call]
+    return wrapped_mod
 
 def check_method_op_support(
     module: torch.jit.ScriptModule, method_name: str = "forward"
@@ -312,4 +313,5 @@ def check_method_op_support(
     Returns:
         bool: True if supported Method
     """
-    return _C.check_method_op_support(module._c, method_name)
+    supported: bool =  _C.check_method_op_support(module._c, method_name)
+    return supported