chore(//py/torch_tensorrt/ts): Make torch_tensorrt.ts.TorchScriptInput

mypy compliant

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

Author: narendasan

py/torch_tensorrt/_Input.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import List, Dict, Any, Tuple, Optional
+from typing import List, Dict, Any, Tuple, Optional, Union
 
 import torch
 
@@ -28,7 +28,7 @@ class _ShapeMode(Enum):
         DYNAMIC = 1
 
     shape_mode = None  #: (torch_tensorrt.Input._ShapeMode): Is input statically or dynamically shaped
-    shape = 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 }``
+    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 = (
         _enums.dtype.unknown
     )  #: The expected data type of the input tensor (default: torch_tensorrt.dtype.float32)
@@ -42,7 +42,7 @@ class _ShapeMode(Enum):
     high_tensor_domain_excl = low_tensor_domain_incl + DOMAIN_OFFSET
     torch_dtype = torch.float32
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         """__init__ Method for torch_tensorrt.Input
 
         Input accepts one of a few construction patterns

py/torch_tensorrt/ts/ts_input.py renamed to py/torch_tensorrt/ts/_Input.py

@@ -9,7 +9,7 @@
 from torch_tensorrt._Input import Input
 
 
-class TSInput(Input):
+class TorchScriptInput(Input):
     """
     Defines an input to a module in terms of expected shape, data type and tensor format.
 
@@ -26,7 +26,7 @@ class TSInput(Input):
         format (torch_tensorrt.TensorFormat): The expected format of the input tensor (default: torch_tensorrt.TensorFormat.NCHW)
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         """__init__ Method for torch_tensorrt.Input
 
         Input accepts one of a few construction patterns
@@ -52,38 +52,39 @@ def __init__(self, *args, **kwargs):
             - Input(shape=(1,3,32,32), dtype=torch_tensorrt.dtype.int32, format=torch_tensorrt.TensorFormat.NCHW)
             - Input(min_shape=(1,3,32,32), opt_shape=[2,3,32,32], max_shape=(3,3,32,32)) #Implicitly dtype=torch_tensorrt.dtype.float32, format=torch_tensorrt.TensorFormat.NCHW
         """
-        super(TSInput, self).__init__(*args, **kwargs)
+        super().__init__(*args, **kwargs)
 
     def _to_internal(self) -> _C.Input:
         internal_in = _C.Input()
         if self.shape_mode == Input._ShapeMode.DYNAMIC:
-            if not Input._supported_input_size_type(self.shape["min_shape"]):
-                raise TypeError(
-                    "Input shape specifications for inputs are required to be a List, tuple or torch.Size, found type: "
-                    + str(type(self.shape["min_shape"]))
-                    + " for min_shape"
-                )
-            else:
-                internal_in.min = self.shape["min_shape"]
-
-            if not Input._supported_input_size_type(self.shape["opt_shape"]):
-                raise TypeError(
-                    "Input shape specifications for inputs are required to be a List, tuple or torch.Size, found type: "
-                    + str(type(self.shape["opt_shape"]))
-                    + " for opt_shape"
-                )
-            else:
-                internal_in.opt = self.shape["opt_shape"]
-
-            if not Input._supported_input_size_type(self.shape["max_shape"]):
-                raise TypeError(
-                    "Input shape specifications for inputs are required to be a List, tuple or torch.Size, found type: "
-                    + str(type(self.shape["max_shape"]))
-                    + " for max_shape"
-                )
-            else:
-                internal_in.max = self.shape["max_shape"]
-            internal_in.input_is_dynamic = True
+            if isinstance(self.shape, dict):
+                if not Input._supported_input_size_type(self.shape["min_shape"]):
+                    raise TypeError(
+                        "Input shape specifications for inputs are required to be a List, tuple or torch.Size, found type: "
+                        + str(type(self.shape["min_shape"]))
+                        + " for min_shape"
+                    )
+                else:
+                    internal_in.min = self.shape["min_shape"]
+
+                if not Input._supported_input_size_type(self.shape["opt_shape"]):
+                    raise TypeError(
+                        "Input shape specifications for inputs are required to be a List, tuple or torch.Size, found type: "
+                        + str(type(self.shape["opt_shape"]))
+                        + " for opt_shape"
+                    )
+                else:
+                    internal_in.opt = self.shape["opt_shape"]
+
+                if not Input._supported_input_size_type(self.shape["max_shape"]):
+                    raise TypeError(
+                        "Input shape specifications for inputs are required to be a List, tuple or torch.Size, found type: "
+                        + str(type(self.shape["max_shape"]))
+                        + " for max_shape"
+                    )
+                else:
+                    internal_in.max = self.shape["max_shape"]
+                internal_in.input_is_dynamic = True
         else:
             if not Input._supported_input_size_type(self.shape):
                 raise TypeError(

py/torch_tensorrt/ts/__init__.py

@@ -1,3 +1,3 @@
 from torch_tensorrt.ts._compiler import *
 from torch_tensorrt.ts._compile_spec import TensorRTCompileSpec
-from torch_tensorrt.ts.ts_input import TSInput
+from torch_tensorrt.ts._Input import TorchScriptInput

py/torch_tensorrt/ts/_compile_spec.py

@@ -9,7 +9,7 @@
 from typing import Tuple, List, Dict
 import warnings
 from copy import deepcopy
-from torch_tensorrt.ts.ts_input import TSInput
+from torch_tensorrt.ts._Input import TorchScriptInput
 import tensorrt as trt
 
 
@@ -195,9 +195,9 @@ def _parse_input_signature(input_signature: Any, depth: int = 0):
 
         ts_i = i
         if i.shape_mode == Input._ShapeMode.STATIC:
-            ts_i = TSInput(shape=i.shape, dtype=i.dtype, format=i.format)
+            ts_i = TorchScriptInput(shape=i.shape, dtype=i.dtype, format=i.format)
         elif i.shape_mode == Input._ShapeMode.DYNAMIC:
-            ts_i = TSInput(
+            ts_i = TorchScriptInput(
                 min_shape=i.shape["min_shape"],
                 opt_shape=i.shape["opt_shape"],
                 max_shape=i.shape["max_shape"],
@@ -245,13 +245,13 @@ def _parse_compile_spec(compile_spec_: Dict[str, Any]) -> _ts_C.CompileSpec:
         for i in inputs:
             if i.shape_mode == Input._ShapeMode.STATIC:
                 ts_inputs.append(
-                    TSInput(
+                    TorchScriptInput(
                         shape=i.shape, dtype=i.dtype, format=i.format
                     )._to_internal()
                 )
             elif i.shape_mode == Input._ShapeMode.DYNAMIC:
                 ts_inputs.append(
-                    TSInput(
+                    TorchScriptInput(
                         min_shape=i.shape["min_shape"],
                         opt_shape=i.shape["opt_shape"],
                         max_shape=i.shape["max_shape"],

tests/py/ts/api/test_classes.py

@@ -104,7 +104,7 @@ def test_infer_from_example_tensor(self):
 
         example_tensor = torch.randn(shape).half()
         i = torchtrt.Input.from_tensor(example_tensor)
-        ts_i = torchtrt.ts.TSInput(shape=i.shape, dtype=i.dtype, format=i.format)
+        ts_i = torchtrt.ts.TorchScriptInput(shape=i.shape, dtype=i.dtype, format=i.format)
         self.assertTrue(self._verify_correctness(ts_i, target))
 
     def test_static_shape(self):
@@ -120,27 +120,27 @@ def test_static_shape(self):
         }
 
         i = torchtrt.Input(shape)
-        ts_i = torchtrt.ts.TSInput(shape=i.shape, dtype=i.dtype, format=i.format)
+        ts_i = torchtrt.ts.TorchScriptInput(shape=i.shape, dtype=i.dtype, format=i.format)
         self.assertTrue(self._verify_correctness(ts_i, target))
 
         i = torchtrt.Input(tuple(shape))
-        ts_i = torchtrt.ts.TSInput(shape=i.shape, dtype=i.dtype, format=i.format)
+        ts_i = torchtrt.ts.TorchScriptInput(shape=i.shape, dtype=i.dtype, format=i.format)
         self.assertTrue(self._verify_correctness(ts_i, target))
 
         i = torchtrt.Input(torch.randn(shape).shape)
-        ts_i = torchtrt.ts.TSInput(shape=i.shape, dtype=i.dtype, format=i.format)
+        ts_i = torchtrt.ts.TorchScriptInput(shape=i.shape, dtype=i.dtype, format=i.format)
         self.assertTrue(self._verify_correctness(ts_i, target))
 
         i = torchtrt.Input(shape=shape)
-        ts_i = torchtrt.ts.TSInput(shape=i.shape, dtype=i.dtype, format=i.format)
+        ts_i = torchtrt.ts.TorchScriptInput(shape=i.shape, dtype=i.dtype, format=i.format)
         self.assertTrue(self._verify_correctness(ts_i, target))
 
         i = torchtrt.Input(shape=tuple(shape))
-        ts_i = torchtrt.ts.TSInput(shape=i.shape, dtype=i.dtype, format=i.format)
+        ts_i = torchtrt.ts.TorchScriptInput(shape=i.shape, dtype=i.dtype, format=i.format)
         self.assertTrue(self._verify_correctness(ts_i, target))
 
         i = torchtrt.Input(shape=torch.randn(shape).shape)
-        ts_i = torchtrt.ts.TSInput(shape=i.shape, dtype=i.dtype, format=i.format)
+        ts_i = torchtrt.ts.TorchScriptInput(shape=i.shape, dtype=i.dtype, format=i.format)
         self.assertTrue(self._verify_correctness(ts_i, target))
 
     def test_data_type(self):
@@ -156,11 +156,11 @@ def test_data_type(self):
         }
 
         i = torchtrt.Input(shape, dtype=torchtrt.dtype.half)
-        ts_i = torchtrt.ts.TSInput(shape=i.shape, dtype=i.dtype, format=i.format)
+        ts_i = torchtrt.ts.TorchScriptInput(shape=i.shape, dtype=i.dtype, format=i.format)
         self.assertTrue(self._verify_correctness(ts_i, target))
 
         i = torchtrt.Input(shape, dtype=torch.half)
-        ts_i = torchtrt.ts.TSInput(shape=i.shape, dtype=i.dtype, format=i.format)
+        ts_i = torchtrt.ts.TorchScriptInput(shape=i.shape, dtype=i.dtype, format=i.format)
         self.assertTrue(self._verify_correctness(ts_i, target))
 
     def test_tensor_format(self):
@@ -176,11 +176,11 @@ def test_tensor_format(self):
         }
 
         i = torchtrt.Input(shape, format=torchtrt.TensorFormat.channels_last)
-        ts_i = torchtrt.ts.TSInput(shape=i.shape, dtype=i.dtype, format=i.format)
+        ts_i = torchtrt.ts.TorchScriptInput(shape=i.shape, dtype=i.dtype, format=i.format)
         self.assertTrue(self._verify_correctness(ts_i, target))
 
         i = torchtrt.Input(shape, format=torch.channels_last)
-        ts_i = torchtrt.ts.TSInput(shape=i.shape, dtype=i.dtype, format=i.format)
+        ts_i = torchtrt.ts.TorchScriptInput(shape=i.shape, dtype=i.dtype, format=i.format)
         self.assertTrue(self._verify_correctness(ts_i, target))
 
     def test_dynamic_shape(self):
@@ -200,7 +200,7 @@ def test_dynamic_shape(self):
         i = torchtrt.Input(
             min_shape=min_shape, opt_shape=opt_shape, max_shape=max_shape
         )
-        ts_i = torchtrt.ts.TSInput(
+        ts_i = torchtrt.ts.TorchScriptInput(
             min_shape=i.shape["min_shape"],
             opt_shape=i.shape["opt_shape"],
             max_shape=i.shape["max_shape"],
@@ -214,7 +214,7 @@ def test_dynamic_shape(self):
             opt_shape=tuple(opt_shape),
             max_shape=tuple(max_shape),
         )
-        ts_i = torchtrt.ts.TSInput(
+        ts_i = torchtrt.ts.TorchScriptInput(
             min_shape=i.shape["min_shape"],
             opt_shape=i.shape["opt_shape"],
             max_shape=i.shape["max_shape"],
@@ -229,7 +229,7 @@ def test_dynamic_shape(self):
             opt_shape=tensor_shape(opt_shape),
             max_shape=tensor_shape(max_shape),
         )
-        ts_i = torchtrt.ts.TSInput(
+        ts_i = torchtrt.ts.TorchScriptInput(
             min_shape=i.shape["min_shape"],
             opt_shape=i.shape["opt_shape"],
             max_shape=i.shape["max_shape"],