chore(//py/torch_tensorrt): Making Input mypy compilaint

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

Author: narendasan

py/torch_tensorrt/_Input.py

@@ -27,20 +27,20 @@ class _ShapeMode(Enum):
         STATIC = 0
         DYNAMIC = 1
 
-    shape_mode = None  #: (torch_tensorrt.Input._ShapeMode): Is input statically or dynamically shaped
-    shape: Union[Tuple, Dict] = None  #: (Tuple or Dict): 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 = (
+    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]
         _enums.dtype.unknown
     )  #: The expected data type of the input tensor (default: torch_tensorrt.dtype.float32)
-    _explicit_set_dtype = False
-    format = (
+    _explicit_set_dtype: bool = False
+    format: _enums.TensorFormat = ( # type: ignore[name-defined]
         _enums.TensorFormat.contiguous
     )  #: The expected format of the input tensor (default: torch_tensorrt.TensorFormat.NCHW)
 
-    DOMAIN_OFFSET = 2.0
-    low_tensor_domain_incl = 0.0
-    high_tensor_domain_excl = low_tensor_domain_incl + DOMAIN_OFFSET
-    torch_dtype = torch.float32
+    DOMAIN_OFFSET: float = 2.0
+    low_tensor_domain_incl: float = 0.0
+    high_tensor_domain_excl: float = low_tensor_domain_incl + DOMAIN_OFFSET
+    torch_dtype: torch.dtype = torch.float32
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         """__init__ Method for torch_tensorrt.Input
@@ -165,15 +165,18 @@ def __str__(self) -> str:
                 str(self.tensor_domain[1]),
             )
         elif self.shape_mode == Input._ShapeMode.DYNAMIC:
-            return "Input(min_shape={}, opt_shape={}, max_shape={}, dtype={}, format={}, domain=[{}, {}))".format(
-                self.shape["min_shape"],
-                self.shape["opt_shape"],
-                self.shape["max_shape"],
-                str(self.dtype),
-                str(self.format),
-                str(self.tensor_domain[0]),
-                str(self.tensor_domain[1]),
-            )
+            if isinstance(self.shape, dict):
+                return "Input(min_shape={}, opt_shape={}, max_shape={}, dtype={}, format={}, domain=[{}, {}))".format(
+                    self.shape["min_shape"],
+                    self.shape["opt_shape"],
+                    self.shape["max_shape"],
+                    str(self.dtype),
+                    str(self.format),
+                    str(self.tensor_domain[0]),
+                    str(self.tensor_domain[1]),
+                )
+            else:
+                raise RuntimeError(f"Input shape is dynamic but shapes are not provided as dictionary (found: {self.shape})")
         else:
             raise RuntimeError("Unknown input shape mode")
 
@@ -189,7 +192,7 @@ def _supported_input_size_type(input_size: Any) -> bool:
             return False
 
     @staticmethod
-    def _parse_dtype(dtype: Any) -> _enums.dtype:
+    def _parse_dtype(dtype: Any) -> _enums.dtype: # type: ignore[name-defined]
         if isinstance(dtype, torch.dtype):
             if dtype == torch.long:
                 return _enums.dtype.long
@@ -217,7 +220,7 @@ def _parse_dtype(dtype: Any) -> _enums.dtype:
             )
 
     @staticmethod
-    def _to_torch_dtype(dtype: _enums.dtype) -> torch.dtype:
+    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:
@@ -233,10 +236,10 @@ def _to_torch_dtype(dtype: _enums.dtype) -> torch.dtype:
             return torch.float32
 
     def is_trt_dtype(self) -> bool:
-        return self.dtype != _enums.dtype.long
+        return bool(self.dtype != _enums.dtype.long)
 
     @staticmethod
-    def _parse_format(format: Any) -> _enums.TensorFormat:
+    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
@@ -256,7 +259,7 @@ def _parse_format(format: Any) -> _enums.TensorFormat:
             )
 
     @staticmethod
-    def _parse_tensor_domain(domain: Optional[Tuple[float, float]]) -> Tuple:
+    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)
 
@@ -355,7 +358,7 @@ def from_tensors(
             for t in ts
         ]
 
-    def example_tensor(self, optimization_profile_field: str = None) -> torch.Tensor:
+    def example_tensor(self, optimization_profile_field: Optional[str] = None) -> Optional[torch.Tensor]:
         """
         Get an example tensor of the shape specified by the Input object
 
@@ -365,38 +368,40 @@ def example_tensor(self, optimization_profile_field: str = None) -> torch.Tensor
         Returns:
             A PyTorch Tensor
         """
-        if optimization_profile_field is not None:
-            try:
-                assert any(
-                    [
-                        optimization_profile_field == field_name
-                        for field_name in ["min_shape", "opt_shape", "max_shape"]
-                    ]
-                )
-            except:
+        if self.shape_mode == Input._ShapeMode.STATIC:
+            if optimization_profile_field is not None:
                 raise ValueError(
-                    "Invalid field name, expected one of min_shape, opt_shape, max_shape"
+                    "Specified a optimization profile field but the input is static"
                 )
+            else:
+                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})")
+        else:
+            if optimization_profile_field is not None:
+                try:
+                    assert any(
+                        [
+                            optimization_profile_field == field_name
+                            for field_name in ["min_shape", "opt_shape", "max_shape"]
+                        ]
+                    )
+                except:
+                    raise ValueError(
+                        "Invalid field name, expected one of min_shape, opt_shape, max_shape"
+                    )
 
-        if (
-            optimization_profile_field is not None
-            and self.shape_mode == Input._ShapeMode.STATIC
-        ):
-            raise ValueError(
-                "Specified a optimization profile field but the input is static"
-            )
+                if isinstance(self.shape, dict):
+                    return torch.rand(self.shape[optimization_profile_field]).to(
+                        dtype=self.torch_dtype
+                    )
+                else:
+                    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."
+                )
 
-        if (
-            optimization_profile_field is None
-            and self.shape_mode == Input._ShapeMode.DYNAMIC
-        ):
-            raise ValueError(
-                "Requested an example tensor from a dynamic shaped input but did not specific which profile field to use."
-            )
 
-        if self.shape_mode == Input._ShapeMode.STATIC:
-            return torch.rand(self.shape).to(dtype=self.torch_dtype)
-        else:
-            return torch.rand(self.shape[optimization_profile_field]).to(
-                dtype=self.torch_dtype
-            )