[export] don't duck size for DIM.AUTO (pytorch#134486)

Summary: apparently DIM.AUTO leads to duck sizing, I didn't catch this. Doing the least intrusive fix possible by using `torch._dynamo.maybe_mark_dynamic()` under the hood.

Test Plan: added test

Differential Revision: D61809344

Pull Request resolved: pytorch#134486
Approved by: https://github.com/avikchaudhuri

Author: pianpwk

test/export/test_export.py

@@ -2492,6 +2492,29 @@ def forward(self, x, y):
                 gm(torch.randn(33, 4), torch.randn(32, 4))
                 gm(torch.randn(128, 4), torch.randn(128, 4))
 
+    def test_dont_duck_size_for_auto_dynamic(self):
+        # for this use case, mark_dynamic() and AUTO should have same effect.
+        # check that same symbol gets allocated to both dims without raising constraint violation.
+        from torch.export.dynamic_shapes import DIM
+
+        AUTO, STATIC = DIM.AUTO, DIM.STATIC
+
+        class Foo(torch.nn.Module):
+            def forward(self, x, y):
+                # x: [s0, s1], y: [s0 + 1, 4]
+                assert y.shape[1] == 4
+                assert x.shape[0] == y.shape[0] - 1
+                return x * 2, y * 2
+
+        # duck sizing would make all static based on these sample inputs
+        inputs = (torch.randn(4, 4), torch.randn(5, 4))
+        shapes = {
+            "x": (AUTO, AUTO),
+            "y": (AUTO, AUTO),
+        }
+        ep = export(Foo(), inputs, dynamic_shapes=shapes)
+        ep.module()(torch.randn(6, 3), torch.randn(7, 4))
+
     @testing.expectedFailureRetraceability  # T183144629
     def test_map(self):
         class Module(torch.nn.Module):

torch/export/dynamic_shapes.py

@@ -835,6 +835,8 @@ def _marked_dynamic(tensor, i):
                 if _marked_dynamic(tensor, i) or dim == DIM.AUTO:
                     # don't have to specify anything if dynamic
                     # None also works, since assume_static_by_default=False
+                    if dim == DIM.AUTO:
+                        torch._dynamo.maybe_mark_dynamic(tensor, i)  # avoid duck sizing
                     continue
                 elif isinstance(dim, _Dim):
                     out[i] = dim
@@ -851,6 +853,8 @@ def _marked_dynamic(tensor, i):
             for i, val in enumerate(tensor.shape):
                 dim = shape[i]
                 if _marked_dynamic(tensor, i) or dim == DIM.AUTO:
+                    if dim == DIM.AUTO:
+                        torch._dynamo.maybe_mark_dynamic(tensor, i)  # avoid duck sizing
                     out.append(None)
                 elif isinstance(dim, _Dim):
                     out.append(dim)