chore: revert FX changes

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

Author: peri044

py/torch_tensorrt/fx/__init__.py

@@ -8,7 +8,7 @@
     tensorrt_converter,
 )
 from .fx2trt import TRTInterpreter, TRTInterpreterResult  # noqa
-from .input_tensor_spec import generate_input_specs  # noqa
+from .input_tensor_spec import generate_input_specs, InputTensorSpec  # noqa
 from .lower_setting import LowerSetting  # noqa
 from .trt_module import TRTModule  # noqa
 from .lower import compile  # usort: skip  #noqa

py/torch_tensorrt/fx/fx2trt.py

@@ -14,7 +14,7 @@
 from torch.fx.passes.shape_prop import TensorMetadata
 
 from .converter_registry import CONVERTERS
-from .._Input import Input
+from .input_tensor_spec import InputTensorSpec
 from .observer import Observer
 from .utils import get_dynamic_dims, LowerPrecision, torch_dtype_to_trt
 
@@ -36,7 +36,7 @@ class TRTInterpreter(torch.fx.Interpreter):
     def __init__(
         self,
         module: torch.fx.GraphModule,
-        input_specs: List[Input],
+        input_specs: List[InputTensorSpec],
         explicit_batch_dimension: bool = False,
         explicit_precision: bool = False,
         logger_level=None,
@@ -79,7 +79,6 @@ def __init__(
         ] = dict()
 
     def validate_input_specs(self):
-        # import pdb; pdb.set_trace()
         for shape, _, _, shape_ranges, has_batch_dim in self.input_specs:
             if not self.network.has_implicit_batch_dimension:
                 assert (

py/torch_tensorrt/fx/tools/common_fx2trt.py

@@ -12,8 +12,7 @@
 from torch.fx.passes import shape_prop
 from torch.fx.passes.infra.pass_base import PassResult
 from torch.testing._internal.common_utils import TestCase
-from torch_tensorrt.fx import TRTInterpreter, TRTModule
-from torch_tensorrt._Input import Input
+from torch_tensorrt.fx import InputTensorSpec, TRTInterpreter, TRTModule
 from torch_tensorrt.fx.passes.lower_basic_pass_aten import (
     compose_bmm,
     compose_chunk,
@@ -212,7 +211,7 @@ def run_test(self, mod, inputs, expected_ops, rtol=1e-03, atol=1e-03):
         mod = torch.fx.symbolic_trace(mod)
         shape_prop.ShapeProp(mod).propagate(*inputs)
         mod = NormalizeArgs(mod).transform()
-        interp = TRTInterpreter(mod, Input.from_tensors(inputs))
+        interp = TRTInterpreter(mod, InputTensorSpec.from_tensors(inputs))
         super().run_test(mod, inputs, expected_ops, None, interp, rtol, atol)
 
     def run_test_custom_compare_results(
@@ -230,7 +229,7 @@ def run_test_custom_compare_results(
         mod = torch.fx.symbolic_trace(mod)
         shape_prop.ShapeProp(mod).propagate(*inputs)
         mod = NormalizeArgs(mod).transform()
-        interp = TRTInterpreter(mod, Input.from_tensors(inputs))
+        interp = TRTInterpreter(mod, InputTensorSpec.from_tensors(inputs))
         super().run_test_custom_compare_results(
             mod, inputs, expected_ops, interp, comparators, fp16_mode=fp16_mode
         )
@@ -259,14 +258,14 @@ def run_test(
             mod = pass_tracer(mod, inputs)
 
         if test_implicit_batch_dim:
-            interp = TRTInterpreter(mod, Input.from_tensors(inputs))
+            interp = TRTInterpreter(mod, InputTensorSpec.from_tensors(inputs))
             super().run_test(
                 mod, inputs, expected_ops, unexpected_ops, interp, rtol, atol, precision
             )
 
         if test_explicit_batch_dim:
             interp = TRTInterpreter(
-                mod, Input.from_tensors(inputs), explicit_batch_dimension=True
+                mod, InputTensorSpec.from_tensors(inputs), explicit_batch_dimension=True
             )
             super().run_test(
                 mod, inputs, expected_ops, unexpected_ops, interp, rtol, atol, precision
@@ -275,7 +274,7 @@ def run_test(
         if test_explicit_precision:
             interp = TRTInterpreter(
                 mod,
-                Input.from_tensors(inputs),
+                InputTensorSpec.from_tensors(inputs),
                 explicit_precision=test_explicit_precision,
             )
             super().run_test(
@@ -284,7 +283,7 @@ def run_test(
 
             interp = TRTInterpreter(
                 mod,
-                Input.from_tensors(inputs),
+                InputTensorSpec.from_tensors(inputs),
                 explicit_batch_dimension=True,
                 explicit_precision=test_explicit_precision,
             )
@@ -304,12 +303,12 @@ def run_test_with_assert_error(
         mod = acc_tracer.trace(mod, inputs)
 
         if test_implicit_batch_dim:
-            interp = TRTInterpreter(mod, Input.from_tensors(inputs))
+            interp = TRTInterpreter(mod, InputTensorSpec.from_tensors(inputs))
             super().run_test_with_error(mod, inputs, interp, expect_error)
 
         if test_explicit_batch_dim:
             interp = TRTInterpreter(
-                mod, Input.from_tensors(inputs), explicit_batch_dimension=True
+                mod, InputTensorSpec.from_tensors(inputs), explicit_batch_dimension=True
             )
             super().run_test_with_error(mod, inputs, interp, expect_error)
 
@@ -323,7 +322,7 @@ def run_test_with_dynamic_shape(
         atol=1e-03,
     ):
         mod.eval()
-        inputs = Input.create_inputs_from_specs(input_specs)
+        inputs = InputTensorSpec.create_inputs_from_specs(input_specs)
         mod = acc_tracer.trace(mod, inputs)
         interp = TRTInterpreter(mod, input_specs, explicit_batch_dimension=True)
         super().run_test(mod, inputs, expected_ops, unexpected_ops, interp, rtol, atol)
@@ -393,7 +392,7 @@ def run_test(
         if test_explicit_batch_dim:
             interp = TRTInterpreter(
                 mod,
-                Input.from_tensors(inputs),
+                InputTensorSpec.from_tensors(inputs),
                 explicit_batch_dimension=True,
             )
             super().run_test(
@@ -403,7 +402,7 @@ def run_test(
         if test_explicit_precision:
             interp = TRTInterpreter(
                 mod,
-                Input.from_tensors(inputs),
+                InputTensorSpec.from_tensors(inputs),
                 explicit_precision=test_explicit_precision,
             )
             super().run_test(
@@ -412,7 +411,7 @@ def run_test(
 
             interp = TRTInterpreter(
                 mod,
-                Input.from_tensors(inputs),
+                InputTensorSpec.from_tensors(inputs),
                 explicit_batch_dimension=True,
                 explicit_precision=test_explicit_precision,
             )
@@ -430,7 +429,7 @@ def run_test_with_dynamic_shape(
         atol=1e-03,
     ):
         mod.eval()
-        inputs = Input.create_inputs_from_specs(input_specs)
+        inputs = InputTensorSpec.create_inputs_from_specs(input_specs)
         mod = self.generate_graph(mod, inputs, expected_ops, unexpected_ops, None)
 
         interp = TRTInterpreter(
@@ -440,7 +439,7 @@ def run_test_with_dynamic_shape(
         )
         # Since the lowering is based on optimal shape. We need to test with
         # different shape(for ex. max shape) for testing dynamic shape
-        inputs_max = Input.create_inputs_from_max_specs(input_specs)
+        inputs_max = InputTensorSpec.create_inputs_from_max_specs(input_specs)
         super().run_test(
             mod, inputs_max, expected_ops, unexpected_ops, interp, rtol, atol
         )