Add test for ShapeEnv recording fallback. (pytorch#109944)

This PR adds a test for the previous PR in this stack: pytorch#109904. In summary, it calls
functions decorated with `@record_shapeenv_event`, that don't have an explicit `ShapeEnv`
parameter, with arguments that don't hold a `ShapeEnv` instance.

Pull Request resolved: pytorch#109944
Approved by: https://github.com/ezyang

test/dynamo/test_misc.py

@@ -51,6 +51,9 @@
 from torch.ao.quantization.quantize_fx import prepare_qat_fx
 from torch.fx.experimental.recording import NotEqualError, replay_shape_env_events
 from torch.fx.experimental.symbolic_shapes import (
+    _constrain_range_for_size,
+    constrain_range,
+    constrain_unify,
     ConstraintViolationError,
     expect_true,
     ShapeEnv,
@@ -7582,6 +7585,18 @@ def test_shape_env_equal_runtime_assert(self):
         )
         self._replay_and_check(main)
 
+    def test_shape_env_recorded_function_fallback(self):
+        # Make sure the record/replay mechanism for ShapeEnv will fallback
+        # if no ShapeEnv instance is found.
+        constrain_range(5, min=2, max=10)
+        constrain_unify(5, 5)
+
+        self.assertExpectedRaisesInline(
+            AssertionError,
+            lambda: _constrain_range_for_size(5, min=2, max=10),
+            """can only constrain range for SymInt""",
+        )
+
     def test_default_dtype_change(self):
         @torch.compile
         def foo():

torch/fx/experimental/symbolic_shapes.py

@@ -379,7 +379,7 @@ def _constrain_range_for_size(a, min: Optional[int] = None, max: Optional[int] =
     if isinstance(a, (SymFloat, SymBool)):
         raise ValueError("Constraining SymFloat/SymBool is nyi")
 
-    assert isinstance(a, SymInt)
+    assert isinstance(a, SymInt), "can only constrain range for SymInt"
     assert isinstance(a.node.expr, sympy.Symbol), "constraining non-Symbols NYI"
 
     if min is None: