chore(//py/torch_tensorrt/ts): Make compile_spec conform to mypy

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

Author: narendasan

py/torch_tensorrt/ts/_compile_spec.py

@@ -1,4 +1,4 @@
-from typing import List, Dict, Any, Set
+from typing import List, Dict, Any, Set, Union, Optional
 import torch
 from torch_tensorrt import _C
 import torch_tensorrt._C.ts as _ts_C
@@ -13,7 +13,7 @@
 import tensorrt as trt
 
 
-def _internal_input_to_torch_class_input(i: _C.Input) -> torch.classes.tensorrt._Input:
+def _internal_input_to_torch_class_input(i: _C.Input) -> torch.classes.tensorrt._Input: # type: ignore[name-defined]
     clone = torch.classes.tensorrt._Input()
     clone._set_min(i.min)
     clone._set_opt(i.opt)
@@ -40,7 +40,7 @@ def _supported_input_size_type(input_size: Any) -> bool:
         )
 
 
-def _parse_op_precision(precision: Any) -> _enums.dtype:
+def _parse_op_precision(precision: Any) -> _enums.dtype:  # type: ignore[name-defined]
     if isinstance(precision, torch.dtype):
         if precision == torch.int8:
             return _enums.dtype.int8
@@ -64,7 +64,7 @@ def _parse_op_precision(precision: Any) -> _enums.dtype:
         )
 
 
-def _parse_enabled_precisions(precisions: Any) -> Set:
+def _parse_enabled_precisions(precisions: Any) -> Set[_enums.dtype]: # type: ignore[name-defined]
     parsed_precisions = set()
     if any([isinstance(precisions, type) for type in [list, tuple, set]]):
         for p in precisions:
@@ -74,7 +74,7 @@ def _parse_enabled_precisions(precisions: Any) -> Set:
     return parsed_precisions
 
 
-def _parse_device_type(device: Any) -> _enums.DeviceType:
+def _parse_device_type(device: Any) -> _enums.DeviceType:  # type: ignore[name-defined]
     if isinstance(device, torch.device):
         if device.type == "cuda":
             return _C.DeviceType.gpu
@@ -159,7 +159,7 @@ def _parse_torch_fallback(fallback_info: Dict[str, Any]) -> _ts_C.TorchFallback:
     return info
 
 
-def _parse_input_signature(input_signature: Any, depth: int = 0):
+def _parse_input_signature(input_signature: Any, depth: int = 0) -> Any:
     if depth > 2:
         raise AssertionError(
             "Input nesting depth exceeds max supported depth, use 1 level: [A, B], or 2 level: [A, (B, C)]"
@@ -197,13 +197,16 @@ def _parse_input_signature(input_signature: Any, depth: int = 0):
         if i.shape_mode == Input._ShapeMode.STATIC:
             ts_i = TorchScriptInput(shape=i.shape, dtype=i.dtype, format=i.format)
         elif i.shape_mode == Input._ShapeMode.DYNAMIC:
-            ts_i = TorchScriptInput(
-                min_shape=i.shape["min_shape"],
-                opt_shape=i.shape["opt_shape"],
-                max_shape=i.shape["max_shape"],
-                dtype=i.dtype,
-                format=i.format,
-            )
+            if isinstance(i.shape, dict):
+                ts_i = TorchScriptInput(
+                    min_shape=i.shape["min_shape"],
+                    opt_shape=i.shape["opt_shape"],
+                    max_shape=i.shape["max_shape"],
+                    dtype=i.dtype,
+                    format=i.format,
+                )
+            else:
+                raise ValueError(f"Input set as dynamic, expected dictionary of shapes but found {i.shape}")
         else:
             raise ValueError(
                 "Invalid shape mode detected for input while parsing the input_signature"
@@ -342,24 +345,24 @@ def _parse_compile_spec(compile_spec_: Dict[str, Any]) -> _ts_C.CompileSpec:
 
 
 def TensorRTCompileSpec(
-    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,
-    truncate_long_and_double=False,
-    calibrator=None,
-    allow_shape_tensors=False,
-) -> torch.classes.tensorrt.CompileSpec:
+    inputs: List[Union[torch.Tensor, Input]] = [],
+    input_signature: Optional[Any] = None,
+    device: Union[torch.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: bool = False,
+    calibrator: object = None,
+    allow_shape_tensors: bool = False,
+) -> torch.classes.tensorrt.CompileSpec: # type: ignore[name-defined]
     """Utility to create a formated spec dictionary for using the PyTorch TensorRT backend
 
     Keyword Args: