rebind_unbacked already records equivalences when a retraced binding site maps
an unbacked symbol to another symbol or to a constant. The same invariant applies
when the new value is a derived symbolic expression such as (u1 + 1) // 2: the
old symbol still has a concrete binding relationship and should be eliminated in
favor of that expression.
The previous assertion assumed any non-symbol replacement with free symbols was
invalid. That is too strong for legitimate derived unbacked shapes and makes the
binding logic reject a value it can represent with the existing ShapeEnv
replacement mechanism. Record the replacement with _eliminate_unbacked, which
is the existing path for replacing an unbacked symbol by a non-symbol expression.
This intentionally does not restore the reverted broad HOP fake-trace
suppression. The FlexAttention/HOP reproducer for #183677 now passes on current
main without that suppression, and the broad suppression was the source of the
internal cond, AOTInductor, Executorch, and FlexAttention regressions.
This change was authored with assistance from an AI assistant.
Test Plan:
python test/test_dynamic_shapes.py TestUnbacked.test_rebind_unbacked_to_symbolic_expression -vpython test/dynamo/test_misc.py MiscTests.test_cond_runtime_assert_generation -vpython test/functorch/test_control_flow.py TestControlFlowTraced.test_cond_functionalized_nested -vpython test/inductor/test_aot_inductor.py -k cond_non_tensor_predicates_dynamic_True_cpu -vpython - <<'PY'
import torch
from torch.nn.attention.flex_attention import create_block_mask, flex_attention
if not torch.cuda.is_available():
raise RuntimeError('CUDA unavailable')
device = 'cuda'
dtype = torch.float16
def score_mod(score, batch, head, q_idx, kv_idx):
return score
compiled_create_block_mask = torch.compile(create_block_mask, dynamic=True, fullgraph=True)
def create_dynamic_block_mask(q_batch, kv_batch):
q_len = q_batch.size(0)
kv_len = kv_batch.size(0)
def mask_mod(batch, head, q_idx, kv_idx):
q_group = q_batch[q_idx]
kv_group = kv_batch[kv_idx]
return (q_group == kv_group) & (q_group != -1) & (kv_group != -1)
return compiled_create_block_mask(mask_mod, B=None, H=None, Q_LEN=q_len, KV_LEN=kv_len, device=device, BLOCK_SIZE=128)
groups = torch.zeros(128, dtype=torch.int64, device=device)
block_mask = create_dynamic_block_mask(groups, groups)
q = torch.randn(1, 1, 128, 64, device=device, dtype=dtype, requires_grad=True)
k = torch.randn(1, 1, 128, 64, device=device, dtype=dtype, requires_grad=True)
v = torch.randn(1, 1, 128, 64, device=device, dtype=dtype, requires_grad=True)
compiled_flex_attention = torch.compile(flex_attention, fullgraph=True, dynamic=True, backend='aot_eager')
out = compiled_flex_attention(q, k, v, score_mod=score_mod, block_mask=block_mask)
out.sum().backward()
torch.cuda.synchronize()
print('ok', out.shape)
PYlintrunner -astack-info: PR: #183837, branch: sanketpurandare/stack/11
SOCIAL SHARE CARD GENERATOR