include user stacks with constraint violation error message (#152924)

Fixes #152918

Before:

```
File "/data/users/bobren/a/pytorch/torch/fx/experimental/symbolic_shapes.py", line 5588, in produce_guards_verbose
    raise ConstraintViolationError(
torch.fx.experimental.symbolic_shapes.ConstraintViolationError: Constraints violated (L['x'].size()[0])! For more information, run with TORCH_LOGS="+dynamic".
  - You marked L['x'].size()[0] as dynamic but your code specialized it to be a constant (5). Either remove the mark_dynamic or use a less strict API such as maybe_mark_dynamic or Dim.AUTO.
```

After:

```
File "/data/users/bobren/a/pytorch/torch/fx/experimental/symbolic_shapes.py", line 5588, in produce_guards_verbose
    raise ConstraintViolationError(
torch.fx.experimental.symbolic_shapes.ConstraintViolationError: Constraints violated (L['x'].size()[0])! For more information, run with TORCH_LOGS="+dynamic".
  - You marked L['x'].size()[0] as dynamic but your code specialized it to be a constant (5). Either remove the mark_dynamic or use a less strict API such as maybe_mark_dynamic or Dim.AUTO.

User stack:
  File "/home/bobren/local/a/pytorch/error.py", line 5, in foo
    return torch.randn(5) * x
```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/152924
Approved by: https://github.com/pianpwk
This commit is contained in:
bobrenjc93 2025-05-10 03:18:54 -07:00 committed by PyTorch MergeBot
parent 4c11b26158
commit 70c8047c2d

View File

@ -3542,6 +3542,8 @@ class ShapeEnv:
# with something like effect token tracking.
self.unbacked_alloc_order: dict[sympy.Symbol, int] = {}
self.specialization_stacks: dict[Source, traceback.StackSummary] = {}
self.trace_asserts = trace_asserts
from torch.fx.experimental.validator import translation_validation_enabled
@ -3624,6 +3626,7 @@ class ShapeEnv:
"_resimplify_floor_div_axioms",
"_expr_sym_node_id",
"_dde_suppressed",
"specialization_stacks",
)
# Mapping of the value of each to-be-compared field into the values that
@ -5159,10 +5162,16 @@ class ShapeEnv:
var_with_range = self._render_range_for_constraint_violation(
source, constraint
)
user_stack = self.specialization_stacks.get(source, None)
msg = (
f"You marked {self._debug_name(source)} as dynamic but your code "
f"specialized it to be a constant ({val}). Either remove the mark_dynamic "
f"or use a less strict API such as maybe_mark_dynamic or Dim.AUTO."
+ (
"\n\nUser stack:\n" + "".join(user_stack.format())
if user_stack
else ""
)
)
record_constraint_violation(
constraint.warn_only, self._debug_name(source), msg
@ -6371,6 +6380,10 @@ class ShapeEnv:
},
)
for source in self.var_to_sources.get(a, []):
if user_tb:
self.specialization_stacks[source] = user_tb
if config.print_specializations:
self.log.warning(
"Specializing %s to %s", self.var_to_sources[a][0].name(), tgt