chore: refactor code

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

py/torch_tensorrt/fx/input_tensor_spec.py

@@ -117,6 +117,43 @@ def from_tensors(cls, tensors: Sequence[torch.Tensor]) -> List["InputTensorSpec"
         assert isinstance(tensors, (list, tuple))
         return [cls.from_tensor(t) for t in tensors]
 
+    @classmethod
+    def from_input(cls, input_obj: Input) -> "InputTensorSpec":
+        """
+        Produce a list of InputTenosrSpec named tuples which contain
+        the information of all the given PyTorch tensors.
+
+        Args:
+            tensors (Iterable[torch.Tensor]): A list of PyTorch tensors.
+
+        Returns:
+            A list of InputTensorSpec named tuples.
+        """
+        assert isinstance(input_obj, Input)
+        input_spec = None
+        if isinstance(input_obj.shape, dict):
+            min_shape = input_obj.shape["min_shape"]
+            opt_shape = input_obj.shape["opt_shape"]
+            max_shape = input_obj.shape["max_shape"]
+            dyn_shape = []
+            for min, opt, max in zip(min_shape, opt_shape, max_shape):
+                if min == opt == max:
+                    dyn_shape.append(min)
+                else:
+                    dyn_shape.append(-1)
+            dtype = input_obj.torch_dtype
+            input_spec = cls(
+                shape=dyn_shape,
+                dtype=dtype,
+                shape_ranges=[(min_shape, opt_shape, max_shape)],
+            )
+        else:
+            shape = input_obj.shape
+            dtype = input_obj.torch_dtype
+            input_spec = cls(shape=shape, dtype=dtype)
+
+        return input_spec
+
     @classmethod
     def from_tensors_with_dynamic_batch_size(
         cls,

py/torch_tensorrt/fx/passes/lower_pass_manager_builder.py

@@ -61,37 +61,6 @@
 # ----------------------------------------------------------------------
 
 
-def convert_input_to_spec(input_obj: Any):
-    input_spec = None
-    if isinstance(input_obj, _Input.Input):
-        if isinstance(input_obj.shape, dict):
-            min_shape = input_obj.shape["min_shape"]
-            opt_shape = input_obj.shape["opt_shape"]
-            max_shape = input_obj.shape["max_shape"]
-            dyn_shape = []
-            for min, opt, max in zip(min_shape, opt_shape, max_shape):
-                if min == opt == max:
-                    dyn_shape.append(min)
-                else:
-                    dyn_shape.append(-1)
-            dtype = input_obj.torch_dtype
-            input_spec = InputTensorSpec(
-                shape=dyn_shape,
-                dtype=dtype,
-                shape_ranges=[(min_shape, opt_shape, max_shape)],
-            )
-        else:
-            shape = input_obj.shape
-            dtype = input_obj.torch_dtype
-            input_spec = InputTensorSpec(shape=shape, dtype=dtype)
-    else:
-        raise ValueError(
-            "Invalid input type provided in the FX lowering. Expected type: torch_tensorrt.Input or torch.Tensor"
-        )
-
-    return input_spec
-
-
 def wrapper(fn: Callable, input) -> Callable:
     @wraps(fn)
     def wrapped_fn(gm):
@@ -303,7 +272,7 @@ def build_trt_lower_pipeline(
         self._trt_input = []
         for input_obj in input:
             if isinstance(input_obj, _Input.Input):
-                self._trt_input.append(convert_input_to_spec(input_obj))
+                self._trt_input.append(InputTensorSpec.from_input(input_obj))
 
         self._additional_input = additional_input
         passes = []
@@ -325,7 +294,7 @@ def build_aten2trt_lower_pipeline(
         self._trt_input = []
         for input_obj in input:
             if isinstance(input_obj, _Input.Input):
-                self._trt_input.append(convert_input_to_spec(input_obj))
+                self._trt_input.append(InputTensorSpec.from_input(input_obj))
 
         self._additional_input = additional_input
         passes = []