split the PR, add dummy inference for converter test

py/torch_tensorrt/dynamo/conversion/_conversion.py

@@ -6,6 +6,8 @@
 import tensorrt as trt
 import torch
 from torch._subclasses.fake_tensor import FakeTensor
+from torch.fx.experimental.proxy_tensor import unset_fake_temporarily
+from torch_tensorrt._Device import Device
 from torch_tensorrt._enums import dtype
 from torch_tensorrt._features import ENABLED_FEATURES
 from torch_tensorrt._Input import Input
@@ -16,6 +18,7 @@
     TRTInterpreterResult,
 )
 from torch_tensorrt.dynamo.runtime import PythonTorchTensorRTModule, TorchTensorRTModule
+from torch_tensorrt.dynamo.utils import get_model_device, get_torch_inputs
 
 logger = logging.getLogger(__name__)
 
@@ -57,6 +60,55 @@ def infer_module_output_dtypes(
     return get_output_dtypes(outputs, truncate_double)
 
 
+def infer_module_output_dtypes_for_test(
+    module: torch.fx.GraphModule,
+    inputs: Sequence[Input],
+    device: Device,
+    kwarg_inputs: Optional[dict[str, Any]] = None,
+    truncate_double: bool = False,
+) -> List[dtype]:
+    """
+    This function performs model inference to determine the output dtypes
+    and truncates them accordingly. inputs can be either arg_inputs or flattened input list.
+    If it is flattened list, kwarg_inputs should be None, as it is already included in the flattened input.
+    """
+    # TODO: We can also determine output dtypes from the module.graph based on node metadata.
+    # However, our converter tests use fx.symbolic_trace which sometimes does not provide metadata,
+    # so we stick to the model inference approach currently.
+    with unset_fake_temporarily():
+        # Get the device on which the model exists
+        # For large models, this can be done on CPU to save GPU memory allocation for TRT.
+        device = get_model_device(module)
+        torch_inputs = get_torch_inputs(inputs, device)
+        if kwarg_inputs is None:
+            kwarg_inputs = {}
+        torch_kwarg_inputs = get_torch_inputs(kwarg_inputs, device)
+        module_outputs = module(*torch_inputs, **torch_kwarg_inputs)
+        if not isinstance(module_outputs, (list, tuple)):
+            module_outputs = [module_outputs]
+
+    # Int64 outputs can sometimes be generated from within other operators
+    # such as aten.sum - such outputs can be truncated
+    output_dtypes = []
+    for output in module_outputs:
+        output_ = output
+        # We don't need to check if output is nested here because the input module will be flattened
+        if not isinstance(output, torch.Tensor):
+            if isinstance(output, str):
+                raise ValueError(
+                    f"Received an output type {type(output)} that's not in the acceptable datatypes (https://pytorch.org/docs/stable/tensor_attributes.html#torch.dtype)"
+                )
+            else:
+                output_ = torch.tensor(output)
+
+        if truncate_double and output_.dtype == dtype.float64:
+            output_dtypes.append(dtype.float32)
+        else:
+            output_dtypes.append(dtype._from(output_.dtype))
+
+    return output_dtypes
+
+
 def interpret_module_to_result(
     module: torch.fx.GraphModule,
     inputs: Sequence[Input],

tests/py/dynamo/conversion/harness.py

@@ -17,7 +17,9 @@
 
 # Use interpreter, input spec, and test case from fx_ts_compat to test Dynamo Converter Registry
 from torch_tensorrt.dynamo.conversion import TRTInterpreter
-from torch_tensorrt.dynamo.conversion._conversion import infer_module_output_dtypes
+from torch_tensorrt.dynamo.conversion._conversion import (
+    infer_module_output_dtypes_for_test,
+)
 from torch_tensorrt.dynamo.lowering import (
     get_decompositions,
     post_lowering,
@@ -273,7 +275,7 @@ def run_test(
         atol=ATOL,
         precision=dtype.f32,
         check_dtype=True,
-        use_dynamo_tracer=True,
+        use_dynamo_tracer=False,
         enable_passes=False,
         propagate_shapes=False,
         int32_reqd=False,
@@ -326,8 +328,10 @@ def run_test(
 
         output_dtypes = None
         if check_dtype:
-            output_dtypes = infer_module_output_dtypes(
+            output_dtypes = infer_module_output_dtypes_for_test(
                 mod,
+                input_specs,
+                compilation_settings.device,
                 truncate_double=compilation_settings.truncate_double,
             )
 
@@ -399,7 +403,7 @@ def run_test_with_dynamic_shape(
         rtol=RTOL,
         atol=ATOL,
         output_dtypes=None,
-        use_dynamo_tracer=True,
+        use_dynamo_tracer=False,
         enable_passes=False,
         use_example_tensors=True,
         pyt_inputs=None,
@@ -426,8 +430,10 @@ def run_test_with_dynamic_shape(
         )
 
         if check_dtype:
-            output_dtypes = infer_module_output_dtypes(
+            output_dtypes = infer_module_output_dtypes_for_test(
                 mod,
+                input_specs,
+                compilation_settings.device,
                 truncate_double=compilation_settings.truncate_double,
             )
 

tests/py/dynamo/conversion/test_acos_aten.py

@@ -66,7 +66,7 @@ def forward(self, input):
             (
                 "3d_dim_dtype_float",
                 (1, 1, 1),
-                (2, 2, 3),
+                (1, 2, 3),
                 (3, 3, 3),
                 torch.float,
                 torch.float,

tests/py/dynamo/conversion/test_acosh_aten.py

@@ -58,15 +58,15 @@ def forward(self, input):
             (
                 "3d_dim_dtype_float",
                 (1, 1, 1),
-                (2, 2, 3),
+                (1, 2, 3),
                 (3, 3, 3),
                 torch.float,
                 torch.float,
             ),
             (
                 "3d_dim_dtype_int32",
                 (1, 1, 1),
-                (2, 2, 4),
+                (1, 2, 4),
                 (2, 3, 5),
                 torch.int32,
                 torch.float,

tests/py/dynamo/conversion/test_any.py

@@ -191,7 +191,7 @@ class TestAnyConverterDynamic(DispatchTestCase):
             (
                 "3d_dynamic_float",
                 (2, 1, 1),
-                (2, 2, 2),
+                (2, 2, 1),
                 (3, 2, 4),
                 torch.float,
             ),
@@ -234,7 +234,7 @@ def forward(self, x):
             (
                 "3d_dynamic_dim_float",
                 (2, 1, 1),
-                (2, 2, 2),
+                (2, 2, 1),
                 (3, 2, 4),
                 torch.float,
                 2,
@@ -252,7 +252,7 @@ def forward(self, x):
             (
                 "3d_dynamic_dim_bool",
                 (2, 1, 1),
-                (2, 2, 2),
+                (2, 2, 1),
                 (3, 2, 4),
                 torch.bool,
                 0,
@@ -285,7 +285,7 @@ def forward(self, x):
             (
                 "3d_dynamic_dims_float",
                 (2, 1, 1),
-                (2, 2, 2),
+                (2, 2, 1),
                 (3, 2, 4),
                 torch.float,
                 [1, 2],

tests/py/dynamo/conversion/test_arange_aten.py

@@ -56,7 +56,6 @@ def forward(self, end_tensor):
             use_example_tensors=False,
             check_dtype=False,
             pyt_inputs=[pyt_input],
-            use_dynamo_tracer=False,
         )
 
 

tests/py/dynamo/conversion/test_cat_aten.py

@@ -42,14 +42,12 @@ def forward(self, x, y):
                 min_shape=(16, 2, 3),
                 opt_shape=(16, 3, 3),
                 max_shape=(16, 32, 3),
-                name="x",
             ),
             Input(
                 dtype=torch.float32,
                 min_shape=(16, 2, 3),
                 opt_shape=(16, 16, 3),
                 max_shape=(16, 32, 3),
-                name="y",
             ),
         ]
         self.run_test_with_dynamic_shape(

tests/py/dynamo/conversion/test_full_aten.py

@@ -50,11 +50,7 @@ def forward(self, shape):
             )
         ]
         self.run_test_with_dynamic_shape(
-            full(),
-            inputs,
-            use_example_tensors=False,
-            check_dtype=False,
-            use_dynamo_tracer=False,
+            full(), inputs, use_example_tensors=False, check_dtype=False
         )
 
     @parameterized.expand(

tests/py/dynamo/conversion/test_ge_aten.py

@@ -1,7 +1,6 @@
 import torch
 import torch.nn as nn
 from parameterized import parameterized
-from torch.export import Dim
 from torch.testing._internal.common_utils import run_tests
 from torch_tensorrt import Input
 
@@ -84,40 +83,10 @@ def forward(self, lhs_val, rhs_val):
             inputs,
         )
 
-    @parameterized.expand(
-        [
-            ("3d_2d", (2, 2, 2), (2, 3, 2), (2, 4, 2), (2, 1), (3, 1), (4, 1)),
-        ]
-    )
-    def test_ge_dynamic_tensor_torch_export(self, *args):
-        class ge(nn.Module):
-            def forward(self, lhs_val, rhs_val):
-                return torch.ops.aten.ge.Tensor(lhs_val, rhs_val)
-
-        input_specs = [
-            Input(
-                min_shape=args[1],
-                opt_shape=args[2],
-                max_shape=args[3],
-            ),
-            Input(
-                min_shape=args[4],
-                opt_shape=args[5],
-                max_shape=args[6],
-            ),
-        ]
-        dyn_dim = Dim("dyn_dim", min=2, max=4)
-        torch_export_dynamic_shapes = {"lhs_val": {1: dyn_dim}, "rhs_val": {0: dyn_dim}}
-
-        self.run_test_with_dynamic_shape(
-            ge(),
-            input_specs,
-            torch_export_dynamic_shapes=torch_export_dynamic_shapes,
-        )
-
     @parameterized.expand(
         [
             ("2d_2d", (2, 3), (4, 3), (5, 3), (2, 3), (4, 3), (5, 3)),
+            ("3d_2d", (2, 2, 2), (2, 3, 2), (2, 4, 2), (2, 1), (3, 1), (4, 1)),
         ]
     )
     def test_ge_dynamic_tensor(self, *args):

tests/py/dynamo/conversion/test_isinf_aten.py

@@ -50,12 +50,7 @@ def forward(self, input):
                 max_shape=(5, 3, 3),
                 dtype=torch.float32,
                 torch_tensor=torch.tensor(
-                    (
-                        [
-                            [[2.7, float("-inf"), 1.1], [4.7, -2.3, float("inf")]],
-                            [[2.7, float("-inf"), 1.1], [4.7, -2.3, float("inf")]],
-                        ]
-                    ),
+                    ([[[2.7, float("-inf"), 1.1], [4.7, -2.3, float("inf")]]]),
                     dtype=torch.float32,
                 ).cuda(),
             )
@@ -77,9 +72,7 @@ def forward(self, input):
                 opt_shape=(3, 2),
                 max_shape=(5, 3),
                 dtype=torch.int,
-                torch_tensor=torch.tensor(
-                    ([[-3, 2], [-2, 1], [1, 2]]), dtype=torch.int
-                ).cuda(),
+                torch_tensor=torch.tensor(([[-3, 2]]), dtype=torch.int).cuda(),
             )
         ]
         self.run_test_with_dynamic_shape(

tests/py/dynamo/conversion/test_isnan_aten.py

@@ -52,18 +52,7 @@ def forward(self, input):
                 max_shape=(5, 3, 3),
                 dtype=torch.float32,
                 torch_tensor=torch.tensor(
-                    (
-                        [
-                            [
-                                [3.2, float("nan"), 3.1],
-                                [float("inf"), 1.1, float("nan")],
-                            ],
-                            [
-                                [3.2, float("nan"), 3.1],
-                                [float("inf"), 1.1, float("nan")],
-                            ],
-                        ]
-                    ),
+                    ([[[3.2, float("nan"), 3.1], [float("inf"), 1.1, float("nan")]]]),
                     dtype=torch.float32,
                 ).cuda(),
             )

tests/py/dynamo/conversion/test_logical_and_aten.py

@@ -37,7 +37,7 @@ def forward(self, lhs_val, rhs_val):
             (
                 "3d_dim_dtype_bool",
                 (1, 1, 1),
-                (2, 2, 3),
+                (1, 2, 3),
                 (3, 3, 3),
                 torch.bool,
             ),