[export] refactor _Dim into Dim (pytorch#149891)

Summary: forward fix T218515233

Test Plan: test_export

Differential Revision: D71769231

Pull Request resolved: pytorch#149891
Approved by: https://github.com/jingsh, https://github.com/angelayi

Author: pianpwk

docs/source/export.rst

@@ -790,7 +790,7 @@ API Reference
 .. autofunction:: save
 .. autofunction:: load
 .. autofunction:: register_dataclass
-.. autofunction:: torch.export.dynamic_shapes.Dim
+.. autoclass:: torch.export.dynamic_shapes.Dim
 .. autofunction:: torch.export.exported_program.default_decompositions
 .. autofunction:: dims
 .. autoclass:: torch.export.dynamic_shapes.ShapesCollection

test/export/test_export.py

@@ -3852,7 +3852,7 @@ def forward(self, x):
 
         dynamic_shapes = (
             {"k": {"k": dim}},
-        )  # ValueError: Node type mismatch; expected <class 'list'>, but got .*_Dim.*.
+        )  # ValueError: Node type mismatch; expected <class 'list'>, but got .*Dim.*.
         with self.assertRaisesRegex(
             torch._dynamo.exc.UserError,
             re.escape(
@@ -12362,7 +12362,7 @@ def test_dynamic_shapes_serdes_user_errors(self):
 
         self.assertExpectedInline(
             _load_dynamic_shapes(spec, from_dict=False),
-            """[[<class 'torch._export.serde.dynamic_shapes.dx'>]]""",
+            """[[Dim('dx', min=4, max=16)]]""",
         )
 
         # check incorrect info in dims

test/onnx/exporter/test_small_models_e2e.py

@@ -450,7 +450,7 @@ def forward(
         )
 
         dynamic_shapes = (
-            {0: torch.export.Dim("dim_x", min=3)},  # _Dim
+            {0: torch.export.Dim("dim_x", min=3)},  # Dim
             [("custom_name_axis_ys_0",), (torch.export.Dim.AUTO,)],  # custom name
             {
                 "a": {0: torch.export.Dim.AUTO},

torch/_export/serde/dynamic_shapes.py

@@ -6,7 +6,6 @@
 from torch.export.dynamic_shapes import (
     _check_dynamic_shapes,
     _DerivedDim,
-    _Dim,
     _DimHint,
     _tree_map_with_path,
     Dim,
@@ -19,7 +18,7 @@
 @dataclasses.dataclass
 class RootDim:
     """
-    This represents a _Dim object.
+    This represents a Dim object.
     """
 
     min: int
@@ -150,7 +149,7 @@ def _standardize_shapes(path, tensor, shape):  # type: ignore[no-untyped-def]
         return out
 
     def _track_dim_from_dims(
-        val: Union[None, int, _DimHint, _Dim]
+        val: Union[None, int, _DimHint, Dim]
     ) -> Union[None, int, str]:
         """
         Tracks dims, ranges, derived dims from the standardized dynamic_shapes spec.
@@ -160,7 +159,7 @@ def _track_dim_from_dims(
         if isinstance(val, _DimHint):  # store enum as string
             return val.__class__.__name__ + "." + val.type.name
 
-        assert isinstance(val, _Dim)
+        assert isinstance(val, Dim)
 
         # track root dim
         root = val.root if isinstance(val, _DerivedDim) else val  # type: ignore[attr-defined]
@@ -297,7 +296,7 @@ def _load_dynamic_shapes(
 
     def deserialize_shape(
         val: Union[None, int, str]
-    ) -> Union[None, int, _Dim, _DimHint]:
+    ) -> Union[None, int, Dim, _DimHint]:
         if val is None or isinstance(val, int):
             return val
         elif val == "_DimHint.AUTO":