auto

Author: xadupre

_unittests/ut_torch_export_patches/test_patch_torch.py

@@ -2,7 +2,12 @@
 from typing import Callable
 import torch
 from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex
-from onnx_diagnostic.ext_test_case import ExtTestCase, requires_torch, requires_transformers
+from onnx_diagnostic.ext_test_case import (
+    ExtTestCase,
+    requires_torch,
+    requires_transformers,
+    has_torch,
+)
 from onnx_diagnostic.torch_export_patches import torch_export_patches
 from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
 
@@ -252,28 +257,61 @@ def forward(self, x, ind1, ind2):
         expected = model(*inputs)
 
         dynamic_string = ({0: "A", 1: "B"}, {0: "C", 1: "D"}, {0: "E"})
+        # ({0: DYN, 1: DYN}, {0: DYN, 1: DYN}, {0: DYN})
+
         dynamic_shapes = use_dyn_not_str(dynamic_string)
-        with self.subTest(name="export 0/1 specialized due to hint of 1 for dimension"):
+        with self.subTest(
+            name="export 0/1 specialized due to hint of 1 for dimension",
+            dynamic_shapes=dynamic_shapes,
+        ):
             try:
                 torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes)
                 raise AssertionError("torch fixed that case")
             except ValueError as e:
                 self.assertIn("export 0/1 specialized due to hint of 1 for dimension", str(e))
 
-        with self.subTest(name="expected shape should be broadcastable to"):
-            try:
+        dynamic_shapes = use_dyn_not_str(dynamic_string, torch.export.Dim.AUTO)
+        if has_torch("2.9"):
+            with self.subTest(
+                name="expected shape should be broadcastable to (>= 2.9)",
+                dynamic_shapes=dynamic_shapes,
+            ):
+                try:
+                    with torch.fx.experimental._config.patch(backed_size_oblivious=True):
+                        torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes)
+                    raise AssertionError("torch fixed that case")
+                except RuntimeError as e:
+                    self.assertIn("expected shape should be broadcastable to", str(e))
+
+        if not has_torch("2.9"):
+            with self.subTest(
+                name="expected shape should be broadcastable to (< 2.9)",
+                dynamic_shapes=dynamic_shapes,
+            ):
                 with torch.fx.experimental._config.patch(backed_size_oblivious=True):
-                    torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes)
-                raise AssertionError("torch fixed that case")
-            except RuntimeError as e:
-                self.assertIn("expected shape should be broadcastable to", str(e))
+                    ep = torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes)
+                got = ep.module()(*inputs)
+                self.assertEqualArray(expected, got)
 
-        with self.subTest(name="patch for 0/1"):
+        with self.subTest(name="patch for 0/1", dynamic_shapes=dynamic_shapes):
             with torch_export_patches():
                 ep = torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes)
             got = ep.module()(*inputs)
             self.assertEqualArray(expected, got)
 
+        if has_torch("2.11"):
+            # Missing PR https://github.com/pytorch/pytorch/pull/164225
+            # Needs more thinking about the patch to apply for this particular example.
+            with self.subTest(
+                name="patch for 0/1 with oblivious", dynamic_shapes=dynamic_shapes
+            ):
+                with torch_export_patches(), torch.fx.experimental._config.patch(
+                    backed_size_oblivious=True
+                ):
+                    ep = torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes)
+                got = ep.module()(*inputs)
+                self.assertEqualArray(expected, got)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/torch_export_patches/patch_inputs.py

@@ -189,19 +189,23 @@ def convert_dynamic_axes_into_dynamic_shapes(
     return (), updated_kwargs, dynamic_shapes
 
 
-def use_dyn_not_str(dynamic_shapes: Any) -> Any:
+def use_dyn_not_str(dynamic_shapes: Any, default_value=None) -> Any:
     """
     Some functions returns dynamic shapes as string.
     This functions replaces them with ``torch.export.Dim.DYNAMIC``.
+    ``default_value=torch.export.Dim.AUTO`` changes the default value.
     """
     if isinstance(dynamic_shapes, list):
-        return [use_dyn_not_str(a) for a in dynamic_shapes]
+        return [use_dyn_not_str(a, default_value=default_value) for a in dynamic_shapes]
     if isinstance(dynamic_shapes, tuple):
-        return tuple(use_dyn_not_str(a) for a in dynamic_shapes)
+        return tuple(use_dyn_not_str(a, default_value=default_value) for a in dynamic_shapes)
     if isinstance(dynamic_shapes, dict):
-        return {k: use_dyn_not_str(v) for k, v in dynamic_shapes.items()}
+        return {
+            k: use_dyn_not_str(v, default_value=default_value)
+            for k, v in dynamic_shapes.items()
+        }
     if isinstance(dynamic_shapes, set):
-        return {use_dyn_not_str(a) for a in dynamic_shapes}
+        return {use_dyn_not_str(a, default_value=default_value) for a in dynamic_shapes}
     if isinstance(dynamic_shapes, str):
-        return torch.export.Dim.DYNAMIC
+        return torch.export.Dim.DYNAMIC if default_value is None else default_value
     return dynamic_shapes