pytorch - 💡(How to fix) Fix `torch.compile` guard evaluation crash: `eval_is_non_overlapping_and_dense()` calling convention mismatch with `IsNonOverlappingAndDenseIndicator` [1 comments, 2 participants]

🐛 Describe the bug

torch.compile crashes during guard evaluation with:

AssertionError: Guard check failed: 0/3: eval_is_non_overlapping_and_dense() takes 2 positional arguments but 8 were given

Root Cause

There is a calling convention mismatch between how IsNonOverlappingAndDenseIndicator guard expressions are printed and how they are evaluated:

1. IsNonOverlappingAndDenseIndicator (SymPy Function in torch/fx/experimental/sym_node.py) takes flat positional args for an N-D tensor — e.g. (s0, s1, s2, s3, st0, st1, st2, st3) for a 4-D tensor.
2. eval_is_non_overlapping_and_dense (in torch/fx/experimental/symbolic_shapes.py) expects two sequence args: (sizes: Sequence, strides: Sequence).
3. SYMPY_INTERP maps the string "IsNonOverlappingAndDenseIndicator" directly to eval_is_non_overlapping_and_dense, so when the guard expression is eval()'d at runtime, 8 flat args are passed into a function expecting 2 sequences.
4. PythonPrinter has no custom _print_IsNonOverlappingAndDenseIndicator method, so SymPy's default printer emits the flat-arg form, triggering the mismatch.

When this triggers

This is a latent bug that rarely surfaces. It triggers when all three conditions are met:

• A tensor has symbolic (dynamic) sizes/strides
• The tensor layout cannot be resolved statically as contiguous
• The IsNonOverlappingAndDenseIndicator expression survives into the guard

In our case, this was triggered by a repeat_kv pattern using .expand() followed by .reshape() on a tensor with dynamic shapes under torch.compile.

Suggested Fix

Add a wrapper in SYMPY_INTERP that splits the flat args into two sequences before forwarding to eval_is_non_overlapping_and_dense:

# In torch/fx/experimental/symbolic_shapes.py

def _eval_is_non_overlapping_and_dense_interp(*args: int) -> int:
    dim = len(args) // 2
    return eval_is_non_overlapping_and_dense(list(args[:dim]), list(args[dim:]))

SYMPY_INTERP["IsNonOverlappingAndDenseIndicator"] = (
    _eval_is_non_overlapping_and_dense_interp
)

Alternatively, PythonPrinter could be taught to emit a call that packs the flat args into two lists before calling the function.

Minimal Reproducer

import torch

# repeat_kv pattern that triggers the bug under dynamic shapes
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(
        batch, num_key_value_heads, n_rep, slen, head_dim
    )
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)

compiled_fn = torch.compile(repeat_kv, dynamic=True)

x = torch.randn(2, 4, 128, 64)
# First call traces successfully
out1 = compiled_fn(x, 2)
# Subsequent call with different dynamic shape may trigger guard re-evaluation crash
x2 = torch.randn(3, 4, 256, 64)
out2 = compiled_fn(x2, 2)

(Note: exact repro may depend on PyTorch version and how dynamic shapes interact with the guard system. The bug was observed on PyTorch 2.8.0.)

Versions

• PyTorch version: 2.8.0
• Python version: 3.12
• OS: Linux

cc @chauhang @penguinwu @ezyang @bobrenjc93 @aditvenk @laithsakka