Revert "hang dim hint constants off Dim (pytorch#134484)"

This reverts commit c142af7.

Reverted pytorch#134484 on behalf of https://github.com/facebook-github-bot due to Diff reverted internally ([comment](pytorch#134484 (comment)))

Author: pytorchmergebot

docs/source/export.rst

@@ -676,6 +676,7 @@ API Reference
 .. autofunction:: save
 .. autofunction:: load
 .. autofunction:: register_dataclass
+.. autoclass:: torch.export.dynamic_shapes.DIM
 .. autofunction:: torch.export.dynamic_shapes.Dim
 .. autofunction:: dims
 .. autoclass:: torch.export.dynamic_shapes.ShapesCollection

test/export/test_export.py

@@ -1773,6 +1773,8 @@ def forward(self, x, y, y1, z):
         )
 
     def test_static_dim_constraints(self):
+        from torch.export.dynamic_shapes import DIM
+
         class Foo(torch.nn.Module):
             def __init__(self) -> None:
                 super().__init__()
@@ -1801,7 +1803,7 @@ def forward(self, x, y, z):
             ((dx, None), (dy, 4), (dz, 3)),
             ((None, 6), (5, None), (None, None)),
             ((4, 6), {0: None, 1: 4}, {0: None, 1: 3}),
-            (None, None, (Dim.STATIC, Dim.STATIC)),
+            (None, None, (DIM.STATIC, DIM.STATIC)),
         ]:
             ep = export(foo, inputs, dynamic_shapes=dynamic_shapes)
             self.assertEqual(foo(*inputs), ep.module()(*inputs))
@@ -1948,7 +1950,9 @@ def forward(self, inp: Inp):
                 self.assertEqual(str(tuple(node.meta["val"].shape)), f"({sym},)")
 
     def test_mismatched_dynamic_shapes(self):
-        AUTO, STATIC = Dim.AUTO, Dim.STATIC
+        from torch.export.dynamic_shapes import DIM
+
+        AUTO, STATIC = DIM.AUTO, DIM.STATIC
 
         class M(torch.nn.Module):
             def forward(self, x):
@@ -1980,7 +1984,7 @@ def forward(self, x):
             + re.escape(
                 "specified at `dynamic_shapes[0]['k']['k'][0]` "
                 "(expected either a list/tuple of dimensions, or a dict mapping indices to dimensions,"
-                " where each dimension is None, an int, a Dim, Dim.AUTO, or Dim.STATIC)"
+                " where each dimension is None, an int, a Dim, DIM.AUTO, or DIM.STATIC)"
             ),
         ):
             export(M(), inputs, dynamic_shapes=dynamic_shapes)
@@ -2057,7 +2061,7 @@ def forward(self, x):
         with self.assertRaisesRegex(
             torch._dynamo.exc.UserError,
             re.escape(
-                "Specifying both `Dim.AUTO` and `Dim` or `DerivedDim` in `dynamic_shapes` is not well supported at the moment, "
+                "Specifying both `DIM.AUTO` and `Dim` or `DerivedDim` in `dynamic_shapes` is not well supported at the moment, "
                 "and can easily lead to constraint violation errors or obscure errors in torch.export."
             ),
         ):
@@ -2461,7 +2465,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
     def test_mark_and_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.
-        AUTO, STATIC = Dim.AUTO, Dim.STATIC
+        from torch.export.dynamic_shapes import DIM
+
+        AUTO, STATIC = DIM.AUTO, DIM.STATIC
 
         class Foo(torch.nn.Module):
             def forward(self, x, y):
@@ -2489,7 +2495,9 @@ def forward(self, x, y):
     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.
-        AUTO, STATIC = Dim.AUTO, Dim.STATIC
+        from torch.export.dynamic_shapes import DIM
+
+        AUTO, STATIC = DIM.AUTO, DIM.STATIC
 
         class Foo(torch.nn.Module):
             def forward(self, x, y):
@@ -6868,7 +6876,9 @@ def test_automatic_dynamic_shapes_simple_equality(self):
         # The next 3 test cases tests for automatic dynamic shapes specs, verifying that automatic dynamism
         # leads to replacement symbols being set for equalities, and inferred relationships being checked
         # with runtime asserts. Check that we specialize to static values when the program says so.
-        AUTO, STATIC = Dim.AUTO, Dim.STATIC
+        from torch.export.dynamic_shapes import DIM
+
+        AUTO, STATIC = DIM.AUTO, DIM.STATIC
 
         # case 1: direct equality between symbols
         class SimpleEquality(torch.nn.Module):
@@ -6933,7 +6943,9 @@ def forward(self, x, y, z):
         )
 
     def test_automatic_dynamic_shapes_constant_relation(self):
-        AUTO, STATIC = Dim.AUTO, Dim.STATIC
+        from torch.export.dynamic_shapes import DIM
+
+        AUTO, STATIC = DIM.AUTO, DIM.STATIC
 
         # case 2: related by constant: s0 + 4 = s1
         class OffBy4(torch.nn.Module):
@@ -6976,7 +6988,9 @@ def forward(self, x, y):
         )
 
     def test_automatic_dynamic_shapes_linear_relation(self):
-        AUTO, STATIC = Dim.AUTO, Dim.STATIC
+        from torch.export.dynamic_shapes import DIM
+
+        AUTO, STATIC = DIM.AUTO, DIM.STATIC
 
         # case 3: linear relation
         class LinearRel(torch.nn.Module):

torch/_export/non_strict_utils.py

@@ -24,10 +24,10 @@
 from torch.export.dynamic_shapes import (
     _check_dynamic_shapes,
     _combine_args,
-    _DimHint,
     _process_dynamic_shapes,
     _transform_shapes_for_default_dynamic,
     _tree_map_with_path,
+    DIM,
 )
 from torch.export.graph_signature import CustomObjArgument
 from torch.fx.experimental import _config as config
@@ -351,7 +351,7 @@ def make_constraints(
                 # we want the symbol, not its replacement, which could be an expression. Maybe
                 # there's a better way to do this, e.g., by (re)computing value ranges for expressions?
                 dim = shape_spec[i] if shape_spec else None
-                if dim is None or isinstance(dim, _DimHint):
+                if dim is None or isinstance(dim, DIM):
                     range_constraints[d.node.expr] = shape_env.var_to_range[
                         d.node._expr
                     ]

torch/export/dynamic_shapes.py

@@ -30,21 +30,20 @@
 
 __all__ = [
     "Constraint",
+    "DIM",
     "Dim",
     "dims",
     "refine_dynamic_shapes_from_suggested_fixes",
 ]
 
 
-class _DimHint(Enum):
+class DIM(Enum):
     """
-    Enum for dynamic shape hints.
-    - AUTO means automatic inference of shape (static or dynamic).
-    - STATIC means static shape (always specialized).
+    Enum for automatic/static dynamic shapes.
     """
 
-    AUTO = auto()
     STATIC = auto()
+    AUTO = auto()
 
 
 class _Dim(type):
@@ -215,7 +214,6 @@ def Dim(name: str, *, min: Optional[int] = None, max: Optional[int] = None):
     Returns:
         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
@@ -229,10 +227,6 @@ def Dim(name: str, *, min: Optional[int] = None, max: Optional[int] = None):
     return dim
 
 
-Dim.AUTO = _DimHint.AUTO  # hint for automatic inference of shape (static or dynamic)
-Dim.STATIC = _DimHint.STATIC  # hint for static shape (always specialized)
-
-
 def dims(*names: str, min: Optional[int] = None, max: Optional[int] = None):
     """
     Util to create multiple :func:`Dim` types.
@@ -674,32 +668,32 @@ def check_symbols(path, tensor, shape):
             for i, dim in shape.items():
                 if isinstance(dim, _Dim):
                     check_same_bounds(dim)
-                elif not (isinstance(dim, (int, _DimHint)) or dim is None):
+                elif not (isinstance(dim, (int, DIM)) or dim is None):
                     raise UserError(
                         UserErrorType.INVALID_INPUT,
                         f"Unexpected dimension mapped to index {i} in input tensor shape {shape} "
                         f"specified at `dynamic_shapes{keystr(path)}` "
-                        f"(expected None, an int, a Dim, Dim.AUTO, or Dim.STATIC, but got {dim} instead)",
+                        f"(expected None, an int, a Dim, DIM.AUTO, or DIM.STATIC, but got {dim} instead)",
                         case_name="dynamic_shapes_validation",
                     )
         elif isinstance(shape, (tuple, list)):
             for i, dim in enumerate(shape):
                 if isinstance(dim, _Dim):
                     check_same_bounds(dim)
-                elif not (isinstance(dim, (int, _DimHint)) or dim is None):
+                elif not (isinstance(dim, (int, DIM)) or dim is None):
                     raise UserError(
                         UserErrorType.INVALID_INPUT,
                         f"Unexpected dimension #{i} in input tensor shape {shape} "
                         f"specified at `dynamic_shapes{keystr(path)}` "
-                        f"(expected None, an int, a Dim, Dim.AUTO, or Dim.STATIC, but got {dim} instead)",
+                        f"(expected None, an int, a Dim, DIM.AUTO, or DIM.STATIC, but got {dim} instead)",
                         case_name="dynamic_shapes_validation",
                     )
         elif shape is not None:
             raise UserError(
                 UserErrorType.INVALID_INPUT,
                 f"Unexpected input tensor shape {shape} specified at `dynamic_shapes{keystr(path)}` "
                 f"(expected either a list/tuple of dimensions, or a dict mapping indices to dimensions,"
-                f" where each dimension is None, an int, a Dim, Dim.AUTO, or Dim.STATIC)",
+                f" where each dimension is None, an int, a Dim, DIM.AUTO, or DIM.STATIC)",
                 case_name="dynamic_shapes_validation",
             )
 
@@ -746,18 +740,18 @@ def check_shape(path, t, dynamic_shape):
 
     _tree_map_with_path(check_shape, combined_args, dynamic_shapes, tree_name="inputs")
 
-    # raise user warning if both Dim.AUTO & Dims are specified in dynamic_shapes
+    # raise user warning if both DIM.AUTO & Dims are specified in dynamic_shapes
     flat_dynamic_shapes = _flatten_dynamic_shapes(combined_args, dynamic_shapes)
     flatter_dynamic_shapes, _ = tree_flatten(flat_dynamic_shapes)
     if any(isinstance(s, _Dim) for s in flatter_dynamic_shapes) and any(
-        s == _DimHint.AUTO for s in flatter_dynamic_shapes
+        s == DIM.AUTO for s in flatter_dynamic_shapes
     ):
         raise UserError(
             UserErrorType.INVALID_INPUT,
-            "Specifying both `Dim.AUTO` and `Dim` or `DerivedDim` in `dynamic_shapes` is not well supported at the moment, "
+            "Specifying both `DIM.AUTO` and `Dim` or `DerivedDim` in `dynamic_shapes` is not well supported at the moment, "
             "and can easily lead to constraint violation errors or obscure errors in torch.export. Dim/DerivedDims "
-            "expect all equal or related dimensions to be specified, and does not yet compose well with `Dim.AUTO`. "
-            "We suggest using `Dim.AUTO` mixed with `None` for auto-dynamic + static shapes, plus torch._check(dim >= min), "
+            "expect all equal or related dimensions to be specified, and does not yet compose well with `DIM.AUTO`. "
+            "We suggest using `DIM.AUTO` mixed with `None` for auto-dynamic + static shapes, plus torch._check(dim >= min), "
             "torch._check(dim <= max) calls in your program to specify min/max ranges, or `Dim`/`DerivedDim` mixed with `None` "
             "if you want to assert on the exact specification of your program's dynamic shapes behavior.",
             case_name="dynamic_shapes_validation",
@@ -779,8 +773,8 @@ def _transform_shapes_for_default_dynamic(
       for all dims governed by this symbol (i.e. relations, equality, linear relations, etc.)
 
     For export.export(), historically dynamism for unspecified dims has been undesirable, so the semantics are:
-    - Dim.AUTO: dynamic, allocated a symbol
-    - None/unspecified/Dim.STATIC: static
+    - DIM.AUTO: dynamic, allocated a symbol
+    - None/unspecified/DIM.STATIC: static
     - Dim/DerivedDims: also a strict assertion
 
     To allow both APIs to follow the same process for producing constraints, this function converts dynamic_shapes
@@ -790,8 +784,8 @@ def _transform_shapes_for_default_dynamic(
     An example conversion might look like, for a 3-d input tensor:
 
         input spec: {
-            0: Dim.AUTO,
-            1: None,  # or Dim.STATIC
+            0: DIM.AUTO,
+            1: None,  # or DIM.STATIC
             2: Dim("dx"),
         }
         output spec: {
@@ -838,10 +832,10 @@ def _marked_dynamic(tensor, i):
             out = {}
             for i, val in enumerate(tensor.shape):
                 dim = shape.get(i, None)
-                if _marked_dynamic(tensor, i) or dim == _DimHint.AUTO:
+                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 == _DimHint.AUTO:
+                    if dim == DIM.AUTO:
                         torch._dynamo.maybe_mark_dynamic(tensor, i)  # avoid duck sizing
                     continue
                 elif isinstance(dim, _Dim):
@@ -852,22 +846,22 @@ def _marked_dynamic(tensor, i):
                     out[i] = dim
                 else:
                     # make explicitly static
-                    assert dim is None or dim == _DimHint.STATIC
+                    assert dim is None or dim == DIM.STATIC
                     out[i] = val
         elif isinstance(shape, (tuple, list)):
             out = []
             for i, val in enumerate(tensor.shape):
                 dim = shape[i]
-                if _marked_dynamic(tensor, i) or dim == _DimHint.AUTO:
-                    if dim == _DimHint.AUTO:
+                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)
                 elif isinstance(dim, int):
                     out.append(dim)
                 else:
-                    assert dim is None or dim == _DimHint.STATIC
+                    assert dim is None or dim == DIM.STATIC
                     out.append(val)
             out = type(shape)(out)  # type: ignore[assignment]
         else:
@@ -1046,7 +1040,7 @@ def _get_dim_name_mapping(
         dynamic_shapes,
         is_leaf=lambda x: isinstance(x, _Dim),
     )[0]:
-        if isinstance(dim, (int, _DimHint)) or dim is None:
+        if isinstance(dim, (int, DIM)) or dim is None:
             continue
         name_to_dim[dim.__name__] = dim
         if isinstance(dim, _DerivedDim):
@@ -1128,11 +1122,9 @@ def refine_dynamic_shapes_from_suggested_fixes(
             name, expr = fix.split(" = ")
             expr = sympy.sympify(expr)
             if isinstance(expr, sympy.Number):
-                # static, integer
-                shape_fixes[name] = int(expr)  # type: ignore[assignment]
+                shape_fixes[name] = int(expr)  # static, integer
             else:
-                # relation or derived dim
-                shape_fixes[name] = expr
+                shape_fixes[name] = expr  # relation or derived dim
 
     name_to_dim = _get_dim_name_mapping(dynamic_shapes)