fix: Add support for Dynamic Shapes

pytorch · Dec 8, 2023 · 9b78a0e · 9b78a0e
1 parent 998e560
commit 9b78a0e
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 4 deletions.
diff --git a/py/torch_tensorrt/dynamo/_DryRunTracker.py b/py/torch_tensorrt/dynamo/_DryRunTracker.py
@@ -229,6 +229,21 @@ def input_formatter_helper(shapes: Any, dtypes: Any) -> str:
         if isinstance(shapes, tuple) and all(isinstance(elt, int) for elt in shapes):
             return f"Tensor: {shapes}@{str(dtypes)[6:]}, "
 
+        # Base case - dynamic shape, single dtype
+        elif (
+            isinstance(shapes, dict)
+            and len(shapes) == 3
+            and all(
+                (
+                    isinstance(shape, tuple)
+                    and all(isinstance(elt, int) for elt in shape)
+                    and k in ("min_shape", "opt_shape", "max_shape")
+                )
+                for k, shape in shapes.items()
+            )
+        ):
+            return f"Tensor: {shapes}@{str(dtypes)[6:]}, "
+
         # Shapes is a sequence
         elif isinstance(shapes, (list, tuple)):
             formatted_str = "List[" if isinstance(shapes, list) else "Tuple("

diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py
@@ -244,7 +244,7 @@ def compile_module(
     dryrun_tracker.total_ops_in_graph = total_ops
     dryrun_tracker.supported_ops_in_graph = num_supported_ops
     dryrun_tracker.graph_input_shapes = parse_complex_tensor_structs(
-        sample_inputs, "shape", tuple
+        sample_inputs, "shape", lambda x: dict(x) if isinstance(x, dict) else tuple(x)
     )
     dryrun_tracker.graph_input_dtypes = parse_complex_tensor_structs(
         sample_inputs, "torch_dtype"
@@ -356,7 +356,9 @@ def compile_module(
             )
 
         subgraph_data.subgraph_input_shapes = parse_complex_tensor_structs(
-            submodule_inputs, "shape", tuple
+            submodule_inputs,
+            "shape",
+            lambda x: dict(x) if isinstance(x, dict) else tuple(x),
         )
         subgraph_data.subgraph_input_dtypes = parse_complex_tensor_structs(
             submodule_inputs, "torch_dtype"
@@ -367,7 +369,9 @@ def compile_module(
         )
 
         subgraph_data.subgraph_output_shapes = parse_complex_tensor_structs(
-            submodule_outputs, "shape", tuple
+            submodule_outputs,
+            "shape",
+            lambda x: dict(x) if isinstance(x, dict) else tuple(x),
         )
         subgraph_data.subgraph_output_dtypes = parse_complex_tensor_structs(
             submodule_outputs, "dtype"
@@ -395,7 +399,7 @@ def compile_module(
         sample_outputs = [sample_outputs]
 
     dryrun_tracker.graph_output_shapes = parse_complex_tensor_structs(
-        sample_outputs, "shape", tuple
+        sample_outputs, "shape", lambda x: dict(x) if isinstance(x, dict) else tuple(x)
     )
     dryrun_tracker.graph_output_dtypes = parse_complex_tensor_structs(
         sample_outputs, "dtype"