[ONNX] Use torch.export.Dim.AUTO in dynamo_export (pytorch#144356)

Align to the changes in pytorch#143158
Pull Request resolved: pytorch#144356
Approved by: https://github.com/justinchuby

Author: titaiwangms

torch/onnx/__init__.py

@@ -480,29 +480,20 @@ def forward(self, x, bias=None):
             )
 
         if export_options is not None and export_options.dynamic_shapes:
-            # Make all shapes dynamic
-            def _to_dynamic_shapes_mapper():
-                arg_order = 0
-
-                def _to_dynamic_shape(x):
-                    nonlocal arg_order
-                    if isinstance(x, torch.Tensor):
-                        rank = len(x.shape)
-                        dynamic_shape = {}
-                        for i in range(rank):
-                            dynamic_shape[i] = torch.export.Dim(
-                                f"arg_{arg_order}_dim_{i}"
-                            )
-                        arg_order += 1
-                        return dynamic_shape
-                    else:
-                        return None
-
-                return _to_dynamic_shape
+            # Make all shapes dynamic if it's possible
+            def _to_dynamic_shape(x):
+                if isinstance(x, torch.Tensor):
+                    rank = len(x.shape)
+                    dynamic_shape = {}
+                    for i in range(rank):
+                        dynamic_shape[i] = torch.export.Dim.AUTO  # type: ignore[attr-defined]
+                    return dynamic_shape
+                else:
+                    return None
 
             # model_args could be nested
             dynamic_shapes = _pytree.tree_map(
-                _to_dynamic_shapes_mapper(),
+                _to_dynamic_shape,
                 model_args,
             )
         else: