Refine the dynamic shapes based on the suggested fixes (#180)

Ref: pytorch/pytorch#134976
justinchuby · Sep 3, 2024 · 313d327 · 313d327
1 parent d8fea7c
commit 313d327
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 8 deletions.
diff --git a/src/torch_onnx/_capture_strategies.py b/src/torch_onnx/_capture_strategies.py
@@ -11,6 +11,7 @@
 from typing import TYPE_CHECKING, Any, Callable
 
 import torch
+import torch._dynamo.exc
 from torch.utils import _pytree
 
 from torch_onnx import _torchscript_converter
@@ -121,9 +122,20 @@ class TorchExportStrategy(CaptureStrategy):
     def _capture(
         self, model, args, kwargs, dynamic_shapes
     ) -> torch.export.ExportedProgram:
-        return torch.export.export(
-            model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes
-        )
+        try:
+            return torch.export.export(
+                model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes
+            )
+        except torch._dynamo.exc.UserError as exc:
+            # Refine the dynamic shapes based on the suggested fixes.
+            new_shapes = (
+                torch.export.dynamic_shapes.refine_dynamic_shapes_from_suggested_fixes(
+                    exc.msg, dynamic_shapes
+                )
+            )
+            return torch.export.export(
+                model, args, kwargs=kwargs, dynamic_shapes=new_shapes
+            )
 
     def _enter(self, model) -> None:
         model_repr = _take_first_line(repr(model))
@@ -149,9 +161,20 @@ class TorchExportNonStrictStrategy(CaptureStrategy):
     def _capture(
         self, model, args, kwargs, dynamic_shapes
     ) -> torch.export.ExportedProgram:
-        return torch.export.export(
-            model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes, strict=False
-        )
+        try:
+            return torch.export.export(
+                model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes, strict=False
+            )
+        except torch._dynamo.exc.UserError as exc:
+            # Refine the dynamic shapes based on the suggested fixes.
+            new_shapes = (
+                torch.export.dynamic_shapes.refine_dynamic_shapes_from_suggested_fixes(
+                    exc.msg, dynamic_shapes
+                )
+            )
+            return torch.export.export(
+                model, args, kwargs=kwargs, dynamic_shapes=new_shapes, strict=False
+            )
 
     def _enter(self, model) -> None:
         model_repr = _take_first_line(repr(model))

diff --git a/src/torch_onnx/_onnx_program.py b/src/torch_onnx/_onnx_program.py
@@ -15,7 +15,7 @@
 import onnxscript._framework_apis.torch_2_5 as onnxscript_apis
 import torch
 from onnxscript import ir
-from torch.utils import _pytree as pytree
+from torch.utils import _pytree
 
 if TYPE_CHECKING:
     import onnxruntime as ort  # type: ignore[import-untyped]
@@ -215,7 +215,7 @@ def _process_args(args, kwargs) -> tuple[torch.Tensor, ...]:
 
 
 def _flatten_inputs(model_args, model_kwargs):
-    flattened_args, _ = pytree.tree_flatten((model_args, model_kwargs))
+    flattened_args, _ = _pytree.tree_flatten((model_args, model_kwargs))
     return flattened_args