[export] refactor DimHints for type errors (pytorch#149424)

Differential Revision: D71414367

Pull Request resolved: pytorch#149424
Approved by: https://github.com/justinchuby, https://github.com/avikchaudhuri

Author: pianpwk

test/inductor/test_cutlass_backend.py

@@ -736,7 +736,7 @@ def forward(self, x, w):
             model = MyModel()
             M, N, K = 128, 64, 64
             dynamic_shapes = {
-                "x": {0: Dim.DYNAMIC},  # type: ignore[attr-defined]
+                "x": {0: Dim.DYNAMIC},
                 "w": None,
             }
 

torch/_export/converter.py

@@ -1514,7 +1514,7 @@ def retrace_as_exported_program(
     ):
         dynamic_shapes = _tree_map_with_path(
             lambda path, x: (
-                [Dim.AUTO] * x.dim() if isinstance(x, torch.Tensor) else None  # type: ignore[attr-defined]
+                [Dim.AUTO] * x.dim() if isinstance(x, torch.Tensor) else None
             ),
             self.sample_args,
         )

torch/_export/serde/dynamic_shapes.py

@@ -137,7 +137,7 @@ def _standardize_shapes(path, tensor, shape):  # type: ignore[no-untyped-def]
         if not isinstance(tensor, torch.Tensor):
             return None
         if shape is None:
-            return [Dim.STATIC] * len(tensor.shape)  # type: ignore[attr-defined]
+            return [Dim.STATIC] * len(tensor.shape)
 
         out = []
         if isinstance(shape, dict):
@@ -158,7 +158,7 @@ def _track_dim_from_dims(
         if val is None or isinstance(val, int):  # non-tensor input or static
             return val
         if isinstance(val, _DimHint):  # store enum as string
-            return val.__class__.__name__ + "." + val.name
+            return val.__class__.__name__ + "." + val.type.name
 
         assert isinstance(val, _Dim)
 
@@ -290,9 +290,9 @@ def _load_dynamic_shapes(
             modulus, remainder = sympy.polys.polytools.div(expr, symbol)
             ddim = dim_cache[name]
             if modulus != 1:
-                ddim = int(modulus) * ddim
+                ddim = int(modulus) * ddim  # type: ignore[assignment, operator]
             if remainder != 0:
-                ddim = ddim + int(remainder)
+                ddim = ddim + int(remainder)  # type: ignore[assignment, operator]
             dim_cache[_expr] = ddim  # cache derived dims
 
     def deserialize_shape(
@@ -301,9 +301,11 @@ def deserialize_shape(
         if val is None or isinstance(val, int):
             return val
         elif val == "_DimHint.AUTO":
-            return _DimHint.AUTO
+            return _DimHint.AUTO()
+        elif val == "_DimHint.DYNAMIC":
+            return _DimHint.DYNAMIC()
         elif val == "_DimHint.STATIC":
-            return _DimHint.STATIC
+            return _DimHint.STATIC()
         if not isinstance(val, str):
             raise UserError(
                 UserErrorType.INVALID_INPUT,
@@ -316,6 +318,6 @@ def deserialize_shape(
                 "Expected dims in `spec['dynamic_shapes']` to be tracked in `spec['dims']`, "
                 f"got {val} which is not in {dims.keys()}",
             )
-        return dim_cache[val]
+        return dim_cache[val]  # type: ignore[return-value]
 
     return tree_map(deserialize_shape, dynamic_shapes)

torch/export/dynamic_shapes.py

@@ -40,7 +40,7 @@
 log = logging.getLogger(__name__)
 
 
-class _DimHint(Enum):
+class _DimHintType(Enum):
     """
     Enum for dynamic shape hints.
     - AUTO means automatic inference of shape (static or dynamic).
@@ -53,6 +53,23 @@ class _DimHint(Enum):
     DYNAMIC = auto()
 
 
+@dataclasses.dataclass
+class _DimHint:
+    type: _DimHintType
+
+    @staticmethod
+    def AUTO():
+        return _DimHint(_DimHintType.AUTO)
+
+    @staticmethod
+    def DYNAMIC():
+        return _DimHint(_DimHintType.DYNAMIC)
+
+    @staticmethod
+    def STATIC():
+        return _DimHint(_DimHintType.STATIC)
+
+
 class _Dim(type):
     """
     Metaclass for :func:`Dim` types.
@@ -206,7 +223,7 @@ def _derive(self, fn):
         )
 
 
-def Dim(name: str, *, min: Optional[int] = None, max: Optional[int] = None):
+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.
@@ -222,22 +239,24 @@ def Dim(name: str, *, min: Optional[int] = None, max: Optional[int] = None):
         A type that can be used in dynamic shape specifications for tensors.
     """
 
-    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
+    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
 
-Dim.AUTO = _DimHint.AUTO  # type: ignore[attr-defined]
-Dim.STATIC = _DimHint.STATIC  # type: ignore[attr-defined]
-Dim.DYNAMIC = _DimHint.DYNAMIC  # type: ignore[attr-defined]
+        _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 dims(
@@ -249,7 +268,7 @@ def dims(
     Returns:
         A tuple of :func:`Dim` types.
     """
-    return tuple(Dim(name, min=min, max=max) for name in names)
+    return tuple(Dim(name, min=min, max=max) for name in names)  # type: ignore[misc]
 
 
 @dataclasses.dataclass
@@ -923,11 +942,11 @@ def _create_static_dim(tensor, i, value):
                     constraint = to_constraint(dim, tensor, i)
                     symbols[dim.__name__].append(constraint)
                 elif isinstance(dim, _DimHint):
-                    if dim == _DimHint.AUTO:
+                    if dim.type == _DimHintType.AUTO:
                         torch._dynamo.maybe_mark_dynamic(tensor, i)
-                    elif dim == _DimHint.STATIC:
+                    elif dim.type == _DimHintType.STATIC:
                         torch._dynamo.mark_static(tensor, i)
-                    elif dim == _DimHint.DYNAMIC:
+                    elif dim.type == _DimHintType.DYNAMIC:
                         torch._dynamo.mark_dynamic(tensor, i)
                     constraints.append(_RelaxedConstraint(id(tensor), i))
                 elif dim is None:
@@ -940,11 +959,11 @@ def _create_static_dim(tensor, i, value):
                     constraint = to_constraint(dim, tensor, i)
                     symbols[dim.__name__].append(constraint)
                 elif isinstance(dim, _DimHint):
-                    if dim == _DimHint.AUTO:
+                    if dim.type == _DimHintType.AUTO:
                         torch._dynamo.maybe_mark_dynamic(tensor, i)
-                    elif dim == _DimHint.STATIC:
+                    elif dim.type == _DimHintType.STATIC:
                         torch._dynamo.mark_static(tensor, i)
-                    elif dim == _DimHint.DYNAMIC:
+                    elif dim.type == _DimHintType.DYNAMIC:
                         torch._dynamo.mark_dynamic(tensor, i)
                     constraints.append(_RelaxedConstraint(id(tensor), i))
                 elif dim is None:
@@ -1063,7 +1082,7 @@ def refine_dynamic_shapes_from_suggested_fixes(
             expr = sympy.sympify(expr)
             if isinstance(expr, sympy.Number):
                 # static, integer
-                shape_fixes[name] = int(expr)
+                shape_fixes[name] = int(expr)  # type: ignore[assignment]
             else:
                 # relation or derived dim
                 shape_fixes[name] = expr

torch/onnx/__init__.py

@@ -474,7 +474,7 @@ def _to_dynamic_shape(x):
                     rank = len(x.shape)
                     dynamic_shape = {}
                     for i in range(rank):
-                        dynamic_shape[i] = torch.export.Dim.AUTO  # type: ignore[attr-defined]
+                        dynamic_shape[i] = torch.export.Dim.AUTO
                     return dynamic_shape
                 else:
                     return None

torch/onnx/_internal/exporter/_dynamic_shapes.py

@@ -63,15 +63,14 @@ def from_dynamic_axes_to_dynamic_shapes(
                     "The axis in dynamic_axes must be in the form of: dict[int, str] or list[int]."
                 )
             dynamic_shapes[input_name] = {
-                k: torch.export.Dim.AUTO  # type: ignore[attr-defined]
-                for k, _ in axes.items()
+                k: torch.export.Dim.AUTO for k, _ in axes.items()
             }
         elif isinstance(axes, list):
             if any(not isinstance(k, int) for k in axes):
                 raise ValueError(
                     "The axis in dynamic_axes must be in the form of: dict[int, str] or list[int]."
                 )
-            dynamic_shapes[input_name] = {k: torch.export.Dim.AUTO for k in axes}  # type: ignore[attr-defined]
+            dynamic_shapes[input_name] = {k: torch.export.Dim.AUTO for k in axes}
         elif axes is None:
             dynamic_shapes[input_name] = None
         else:
@@ -203,15 +202,15 @@ def convert_str_to_export_dim(
             converted_axes_dict: dict[int, _Dim | _DimHint | None] = {}
             for axis, dim in axes.items():
                 if isinstance(dim, str):
-                    converted_axes_dict[axis] = torch.export.Dim.AUTO  # type: ignore[attr-defined]
+                    converted_axes_dict[axis] = torch.export.Dim.AUTO
                 else:
                     converted_axes_dict[axis] = dim
             dynamic_shapes_with_export_dim.append(converted_axes_dict)
         elif isinstance(axes, (list, tuple)):
             converted_axes_list: list[_Dim | _DimHint | None] = []
             for dim in axes:
                 if isinstance(dim, str):
-                    converted_axes_list.append(torch.export.Dim.AUTO)  # type: ignore[attr-defined]
+                    converted_axes_list.append(torch.export.Dim.AUTO)
                 else:
                     converted_axes_list.append(dim)
             dynamic_shapes_with_export_dim.append(converted_axes_list)