Add test for `ShapeEnv` recording fallback. (#109944)
This PR adds a test for the previous PR in this stack: #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: https://github.com/pytorch/pytorch/pull/109944
Approved by: https://github.com/ezyang
diff --git a/test/dynamo/test_misc.py b/test/dynamo/test_misc.py
index e2894c2..fb09ce8 100644
--- a/test/dynamo/test_misc.py
+++ b/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 @@
)
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():
diff --git a/torch/fx/experimental/symbolic_shapes.py b/torch/fx/experimental/symbolic_shapes.py
index 2c5df3d..a6d2ddd 100644
--- a/torch/fx/experimental/symbolic_shapes.py
+++ b/torch/fx/experimental/symbolic_shapes.py
@@ -379,7 +379,7 @@
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:
