add run_test_with_dynamic_shape change

tests/py/dynamo/conversion/harness.py

@@ -3,7 +3,7 @@
 import logging
 import time
 import unittest
-from typing import Callable, List, Optional, Tuple
+from typing import Any, Callable, List, Optional, Tuple
 
 import torch
 import torch_tensorrt
@@ -12,6 +12,7 @@
 from torch_tensorrt import Input
 from torch_tensorrt._enums import dtype
 from torch_tensorrt.dynamo import _defaults
+from torch_tensorrt.dynamo._defaults import default_device
 from torch_tensorrt.dynamo._settings import CompilationSettings
 
 # Use interpreter, input spec, and test case from fx_ts_compat to test Dynamo Converter Registry
@@ -223,10 +224,22 @@ def generate_graph(
         use_dynamo_tracer: bool,
         enable_passes: bool,
         propagate_shapes: bool = False,
+        torch_export_dynamic_shapes: Optional[Any] = None,
     ):
         mod = mod.eval()
         if use_dynamo_tracer:
-            exported_program = torch_tensorrt.dynamo.trace(mod, tuple(original_inputs))
+            if torch_export_dynamic_shapes is not None:
+                device = default_device()
+                torch_export_inputs = get_torch_inputs(original_inputs, device)
+                exported_program = torch.export.export(
+                    mod,
+                    tuple(torch_export_inputs),
+                    dynamic_shapes=torch_export_dynamic_shapes,
+                )
+            else:
+                exported_program = torch_tensorrt.dynamo.trace(
+                    mod, tuple(original_inputs)
+                )
             exported_program = pre_export_lowering(exported_program)
             exported_program = exported_program.run_decompositions(
                 get_decompositions(False)
@@ -387,13 +400,16 @@ def run_test_with_dynamic_shape(
         propagate_shapes=False,
         check_dtype=True,
         make_refittable=False,
+        # this field is optional, in case user wants to specify custom dynamic_shapes rules for the testcase
+        torch_export_dynamic_shapes: Optional[Any] = None,
     ):
         mod = self.generate_graph(
             mod,
             input_specs,
             use_dynamo_tracer=use_dynamo_tracer,
             enable_passes=enable_passes,
             propagate_shapes=propagate_shapes,
+            torch_export_dynamic_shapes=torch_export_dynamic_shapes,
         )
 
         # Previous instance of the interpreter auto-casted 64-bit inputs

tests/py/dynamo/conversion/test_ge_aten.py

@@ -1,6 +1,7 @@
 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
 
@@ -85,10 +86,40 @@ def forward(self, lhs_val, rhs_val):
 
     @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_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)),
+        ]
+    )
     def test_ge_dynamic_tensor(self, *args):
         class ge(nn.Module):
             def forward(self, lhs_val, rhs_val):