[export] refactor _Dim into Dim (#149891)

Summary:
X-link: meta-pytorch/torchrec#2847


X-link: pytorch/executorch#9559

forward fix T218515233

Test Plan: test_export

Differential Revision: D71769231

Author: pianpwk

test/export/test_export.py

@@ -3835,7 +3835,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(
@@ -12311,7 +12311,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":

torch/export/dynamic_shapes.py

@@ -70,83 +70,118 @@ def STATIC():
         return _DimHint(_DimHintType.STATIC)
 
 
-class _Dim(type):
+class Dim:
     """
-    Metaclass for :func:`Dim` types.
+    :func:`Dim` constructs a type analogous to a named symbolic integer with a range.
+    It can be used to describe multiple possible values of a dynamic tensor dimension.
+    Note that different dynamic dimensions of the same tensor, or of different tensors,
+    can be described by the same type.
+
+    Args:
+        name (str): Human-readable name for debugging.
+        min (Optional[int]): Minimum possible value of given symbol (inclusive)
+        max (Optional[int]): Maximum possible value of given symbol (inclusive)
+
+    Returns:
+        A type that can be used in dynamic shape specifications for tensors.
     """
 
-    @staticmethod
-    def readable(name, min_, max_):
+    AUTO = _DimHint.AUTO()
+    DYNAMIC = _DimHint.DYNAMIC()
+    STATIC = _DimHint.STATIC()
+
+    def __init__(
+        self, name: str, *, min: Optional[int] = None, max: Optional[int] = None
+    ):
         from torch.utils._sympy.numbers import int_oo
 
-        if min_ == 2:
-            min_ = None
-        if max_ == int_oo:
-            max_ = None
-        if min_ is None and max_ is None:
-            return f"Dim('{name}')"
-        if min_ is None:
-            return f"Dim('{name}', max={max_})"
-        if max_ is None:
-            return f"Dim('{name}', min={min_})"
-        return f"Dim('{name}', min={min_}, max={max_})"
+        _min = 0 if min is None else min
+        _max = int_oo if max is None else max
+        assert _max > _min, f"Cannot create Dim with inconsistent min={min}, max={max}"
+        assert name.isidentifier(), f"Dim name must be a valid identifier, got {name}"
+        self.__name__ = name
+        self.min = _min
+        self.max = _max
 
-    def __add__(cls, other):
+    def __add__(self, other) -> "Dim":
         # e.g., dim + 1
         if type(other) is not int:
             raise NotImplementedError(
-                f"Attempted to add {other} to {cls.__name__}, where an integer was expected. "
+                f"Attempted to add {other} to {self.__name__}, where an integer was expected. "
                 "(Only increasing linear operations with integer coefficients are supported.)"
             )
-        return cls._derive(lambda x: x + other)
+        return self._derive(lambda x: x + other)
 
-    def __radd__(cls, other):
-        return cls + other
+    def __radd__(self, other) -> "Dim":
+        return self + other
 
-    def __sub__(cls, other):
+    def __sub__(self, other) -> "Dim":
         # e.g., dim - 1
         if type(other) is not int:
             raise NotImplementedError(
-                f"Attempted to subtract {other} from {cls.__name__}, where an integer was expected. "
+                f"Attempted to subtract {other} from {self.__name__}, where an integer was expected. "
                 "(Only increasing linear operations with integer coefficients are supported.)"
             )
-        return cls._derive(lambda x: x - other)
+        return self._derive(lambda x: x - other)
 
-    def __rsub__(cls, other):
+    def __rsub__(self, other) -> "Dim":
         raise NotImplementedError(
-            f"Attempted to negate {cls.__name__}. "
+            f"Attempted to negate {self.__name__}. "
             "(Only increasing linear operations with integer coefficients are supported.)"
         )
 
-    def __mul__(cls, other):
+    def __mul__(self, other) -> "Dim":
         # e.g., dim * 2
         if type(other) is not int or other <= 0:
             raise NotImplementedError(
-                f"Attempted to multiply {other} with {cls.__name__}, where a positive integer was expected. "
+                f"Attempted to multiply {other} with {self.__name__}, where a positive integer was expected. "
                 "(Only increasing linear operations with integer coefficients are supported.)"
             )
-        return cls._derive(lambda x: x * other)
+        return self._derive(lambda x: x * other)
 
-    def __rmul__(cls, other):
-        return cls * other
+    def __rmul__(self, other) -> "Dim":
+        return self * other
 
-    def _derived_name(cls, fn):
+    def _derived_name(self, fn) -> str:
         from sympy import sympify
 
-        return str(fn(sympify(cls.__name__)))
+        return str(fn(sympify(self.__name__)))
 
-    def _derive(cls, fn):
-        return _DerivedDim(cls._derived_name(fn), (int,), {"root": cls, "fn": fn})
+    def _derive(self, fn) -> "Dim":
+        return _DerivedDim(self._derived_name(fn), self, fn)
 
+    @staticmethod
+    def readable(name: str, min_: int, max_: int) -> str:
+        from torch.utils._sympy.numbers import int_oo
 
-class _StaticDim(_Dim):
+        if min_ == 2:
+            min_ = None  # type: ignore[assignment]
+        if max_ == int_oo:
+            max_ = None  # type: ignore[assignment]
+        if min_ is None and max_ is None:
+            return f"Dim('{name}')"
+        if min_ is None:
+            return f"Dim('{name}', max={max_})"
+        if max_ is None:
+            return f"Dim('{name}', min={min_})"
+        return f"Dim('{name}', min={min_}, max={max_})"
+
+    def __repr__(self):
+        return Dim.readable(self.__name__, self.min, self.max)
+
+
+class _StaticDim(Dim):
     """
-    Meta class for static :func:`Dim` types.
+    Class for static :func:`Dim` types.
 
     This class is only for setting and checking static dim constraints,
     and the user should never interact with it.
     """
 
+    def __init__(self, value: int):
+        self.__name__ = str(value)
+        self.value = value
+
     @property
     def min(self):
         return self.value  # type: ignore[attr-defined]
@@ -156,9 +191,9 @@ def max(self):
         return self.value  # type: ignore[attr-defined]
 
 
-class _DerivedDim(_Dim):
+class _DerivedDim(Dim):
     """
-    Metaclass for derived :func:`Dim` types.
+    Class for derived :func:`Dim` types.
 
     Currently we only support increasing linear expressions with integer coefficients.
     In other words, a derived Dim can always be written in the form Ax + B, where
@@ -172,6 +207,11 @@ class _DerivedDim(_Dim):
     The range of a derived Dim is computed by mapping `fn` over the range of its `root`.
     """
 
+    def __init__(self, name: str, root: Dim, fn: Callable):
+        self.__name__ = name
+        self.root = root
+        self.fn = fn
+
     @property
     def min(self):
         # assume that self.fn is an increasing function
@@ -218,50 +258,17 @@ def _derive(self, fn):
         # As a consequence, roots are always regular Dims (i.e., not derived Dims).
         return _DerivedDim(
             self._derived_name(fn),
-            (int,),
-            {"root": self.root, "fn": lambda x: fn(self.fn(x))},  # type: ignore[attr-defined]
+            self.root,
+            lambda x: fn(self.fn(x)),
         )
 
-
-class Dim(type):
-    """
-    :func:`Dim` constructs a type analogous to a named symbolic integer with a range.
-    It can be used to describe multiple possible values of a dynamic tensor dimension.
-    Note that different dynamic dimensions of the same tensor, or of different tensors,
-    can be described by the same type.
-
-    Args:
-        name (str): Human-readable name for debugging.
-        min (Optional[int]): Minimum possible value of given symbol (inclusive)
-        max (Optional[int]): Maximum possible value of given symbol (inclusive)
-
-    Returns:
-        A type that can be used in dynamic shape specifications for tensors.
-    """
-
-    AUTO = _DimHint.AUTO()
-    DYNAMIC = _DimHint.DYNAMIC()
-    STATIC = _DimHint.STATIC()
-
-    def __new__(
-        metacls, name: str, *, min: Optional[int] = None, max: Optional[int] = None
-    ):
-        from torch.utils._sympy.numbers import int_oo
-
-        _min = 0 if min is None else min
-        _max = int_oo if max is None else max
-        assert _max > _min, f"Cannot create Dim with inconsistent min={min}, max={max}"
-        assert name.isidentifier(), f"Dim name must be a valid identifier, got {name}"
-        dim = _Dim(name, (int,), {"min": _min, "max": _max})
-        dim.__module__ = getattr(
-            inspect.getmodule(inspect.stack()[1][0]), "__name__", "__main__"
-        )
-        return dim
+    def __repr__(self):
+        return self.__name__
 
 
 def dims(
     *names: str, min: Optional[int] = None, max: Optional[int] = None
-) -> tuple[_Dim, ...]:
+) -> tuple[Dim, ...]:
     """
     Util to create multiple :func:`Dim` types.
 
@@ -722,8 +729,8 @@ def check_same_bounds(dim):
         if dim.__name__ in bounds:
             min_, max_ = bounds[dim.__name__]
             if dim.min != min_ or dim.max != max_:
-                this_ = _Dim.readable(dim.__name__, min_, max_)
-                that_ = _Dim.readable(dim.__name__, dim.min, dim.max)
+                this_ = Dim.readable(dim.__name__, min_, max_)
+                that_ = Dim.readable(dim.__name__, dim.min, dim.max)
                 raise UserError(
                     UserErrorType.INVALID_INPUT,
                     f"Found different definitions {this_} and {that_} "
@@ -735,7 +742,7 @@ def check_same_bounds(dim):
     def check_symbols(path, tensor, shape):
         if isinstance(shape, dict):
             for i, dim in shape.items():
-                if isinstance(dim, _Dim):
+                if isinstance(dim, Dim):
                     check_same_bounds(dim)
                 elif dim is None:
                     _warn_on_None_dynamic_shape_dimension()
@@ -750,7 +757,7 @@ def check_symbols(path, tensor, shape):
                     )
         elif isinstance(shape, (tuple, list)):
             for i, dim in enumerate(shape):
-                if isinstance(dim, _Dim):
+                if isinstance(dim, Dim):
                     check_same_bounds(dim)
                 elif dim is None:
                     _warn_on_None_dynamic_shape_dimension()
@@ -911,7 +918,7 @@ def root_value():
                 ),
             )
         else:
-            assert isinstance(dim, _Dim)
+            assert isinstance(dim, Dim)
             constraint = _Constraint(  # type: ignore[assignment]
                 id(tensor),
                 i,
@@ -924,7 +931,7 @@ def root_value():
 
     def update_symbols(path, tensor, shape):
         def _create_static_dim(tensor, i, value):
-            return _StaticDim(str(value), (int,), {"value": value})
+            return _StaticDim(value)
 
         # clean out decorators from user side, or previous export call
         # we also delete these attributes in non_strict_utils.py/make_constraints()
@@ -936,7 +943,7 @@ def _create_static_dim(tensor, i, value):
 
         if isinstance(shape, dict):
             for i, dim in shape.items():
-                if isinstance(dim, (int, _Dim)):
+                if isinstance(dim, (int, Dim)):
                     if isinstance(dim, int):
                         dim = _create_static_dim(tensor, i, dim)
                     constraint = to_constraint(dim, tensor, i)
@@ -953,7 +960,7 @@ def _create_static_dim(tensor, i, value):
                     torch._dynamo.mark_static(tensor, i)
         elif isinstance(shape, (tuple, list)):
             for i, dim in enumerate(shape):
-                if isinstance(dim, (int, _Dim)):
+                if isinstance(dim, (int, Dim)):
                     if isinstance(dim, int):
                         dim = _create_static_dim(tensor, i, dim)
                     constraint = to_constraint(dim, tensor, i)
@@ -1002,14 +1009,14 @@ def _get_dim_name_mapping(
     name_to_dim = {}
     for dim in tree_flatten(
         dynamic_shapes,
-        is_leaf=lambda x: isinstance(x, _Dim),
+        is_leaf=lambda x: isinstance(x, Dim),
     )[0]:
         if dim is None:
             # NOTE: this must denote a non-Tensor or automatic at this point.
             continue
         if isinstance(dim, int):
             continue
-        elif isinstance(dim, _Dim):
+        elif isinstance(dim, Dim):
             name_to_dim[dim.__name__] = dim
             if isinstance(dim, _DerivedDim):
                 name_to_dim[dim.root.__name__] = dim.root  # type: ignore[attr-defined]
@@ -1092,7 +1099,7 @@ def refine_dynamic_shapes_from_suggested_fixes(
     # track derived dim roots
     roots: set[str] = set()
     for k, c in shape_fixes.items():
-        assert isinstance(c, (int, _Dim, _DerivedDim, sympy.Expr))
+        assert isinstance(c, (int, Dim, _DerivedDim, sympy.Expr))
         if isinstance(c, sympy.Expr):  # check dim/derived dim expression
             assert _is_supported_equivalence(c)
             shape_fixes[k] = c