apache / apache/tvm

[Relax][Frontend][Torch] `_squeeze` does not validate `dim`: out-of-bounds tuple dims are silently converted to `squeeze(None)` via `from_fx` (native PyTorch raises `IndexError`)

Open Beginner friendly
#20,321 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
13.7k
Forks
4k
Avg merge
2d 1h
Merged PRs (30d)
112

Description

### Expected behavior

`tvm.relax.frontend.torch._squeeze`
(`python/tvm/relax/frontend/torch/base_fx_graph_translator.py:2523-2543`) reads the
`dim`/`dims` argument of a PyTorch `squeeze` call, drops out-of-range axes from a
list/tuple, and falls back to `dim=None` when the filtered list is empty:

```python
if isinstance(dim, list | tuple) and len(dim) > 0:
shape = self.shape_of(x)
valid_dims = []
for d in dim:
axis = d if d >= 0 else len(shape) + d
if axis < len(shape):
valid_dims.append(d)
# If no valid dims, use None to squeeze all size-1 dimensions
dim = valid_dims if valid_dims else None
return self.block_builder.emit(relax.op.squeeze(x, dim))
```

An out-of-range `dim` should be rejected with a clear, frontend-level error, matching
native PyTorch, which raises `IndexError: Dimension out of range` for the same call. It
should not be silently dropped and reinterpreted as a **different** operation.

### Actual behavior

When `dim` is a list/tuple whose elements are all positive out-of-range, `valid_dims`
becomes empty, `dim` is reset to `None`, and the call is **silently converted into
`squeeze(None)`** (remove all size-1 dimensions). `torch.fx.symbolic_trace` records an
out-of-range tuple `dim` without validating it, so the legacy `from_fx` path imports the
model and emits a **differently-shaped tensor** where native PyTorch raises `IndexError`:

```
shape=(2, 3) squeeze((5,)) torch=IndexError tvm=OK (2, 3)
shape=(2, 1, 3) squeeze((5,)) torch=IndexError tvm=OK (2, 3) <- size-1 dim silently removed
shape=(1, 2, 1) squeeze((5,)) torch=IndexError tvm=OK (2,) <- both size-1 dims removed
```

`(2, 1, 3)` is the clearest case: the module is declared to keep its shape, and the
frontend instead returns a tensor with the size-1 axis removed.

The other argument forms are rejected, but only by the C++ side of `relax.op.squeeze`,
with an opaque op-level error rather than a frontend one — so the forms behave
inconsistently:

```
shape=(2, 3) squeeze(5) torch=IndexError tvm=InternalError
shape=(2, 3) squeeze(-5) torch=IndexError tvm=InternalError
shape=(2, 3) squeeze((-5,)) torch=IndexError tvm=InternalError
```

```
tvm.error.InternalError: In Op(relax.squeeze), the input axis 5 is out of range.
The input tensor has 2 dimensions, so axis should be in range [-2, 2).
```

The negative out-of-range case is kept by the filter (`len(shape) + d` is always
`< len(shape)` for `d < 0`), so it reaches the op; the positive out-of-range case is
dropped by the filter and never reaches it. Both should be a frontend error.

Additional context:

- **Scope** — Only the legacy `from_fx` path is affected. `torch.export` rejects an
out-of-range `dim` (scalar or tuple) at trace time with `IndexError`, so
`from_exported_program` never reaches `_squeeze` with a bad dim. Because the model is
invalid in native PyTorch anyway (it would crash on any input), this is a
robustness/validation gap rather than a wrong result on a valid model — but a model
with a latent, never-exercised bad `dim` silently produces a differently-shaped module
instead of surfacing the error.
- **In-bounds dims are unaffected** — control cases (`squeeze((0,2))` on `(1,2,1)`,
`squeeze((1,))` on `(2,1,3)`) match PyTorch exactly.
- **Misleading comment** — the filter comment says "filter out axes where dimension is
not 1", but the code only bounds-checks axes; it never inspects the size. The comment
does not describe what the code does, and the fallback it feeds does not match PyTorch
semantics.
- **Where to validate** — this converter is shared by `from_fx` and
`from_exported_program` (both `squeeze`, `squeeze.dim` and `squeeze.dims` dispatch to
it), so a single fix covers all of them and makes the error message consistent with the
range check that `relax.op.squeeze` already performs.

### Environment

- OS: Linux
- TVM: main branch (`60a9871a9`, re-verified 2026-09-12; also observed on `390af87345`)
- Python: 3.11
- torch: 2.10.0+cu128
- target: `llvm`

### Steps to reproduce

```python
"""Repro: TVM torch frontend silently converts out-of-bounds tuple dims to squeeze(None)."""
import numpy as np
import torch
import torch.nn as nn
from torch.fx import symbolic_trace

import tvm
from tvm import relax
from tvm.relax.frontend.torch import from_fx

class M(nn.Module):
def __init__(self, dim):
super().__init__()
self.dim = dim

def forward(self, t):
return t.squeeze(self.dim)

def run(shape, dim):
m = M(dim).eval()
x = np.arange(int(np.prod(shape))).reshape(shape).astype(np.float32) + 1
xt = torch.tensor(x)
try: # native PyTorch ground truth
ref = m(xt).numpy().shape
ref_txt = f"OK {ref}"
except Exception as e:
ref_txt = f"{type(e).__name__}"
try: # TVM legacy from_fx path
gm = symbolic_trace(m)
mod = from_fx(gm, [((*shape,), "float32")])
ex = relax.build(mod, target="llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())
out = vm["main"](x)
arr = out[0] if hasattr(out, "__len__") and len(out) else out
tv_txt = f"OK {tuple(np.asarray(arr.numpy()).shape)}"
except Exception as e:
tv_txt = f"{type(e).__name__}"
print(f" shape={shape} dim={dim} torch={ref_txt} tvm={tv_txt}")

if __name__ == "__main__":
print("tvm:", tvm.__version__, "| torch:", torch.__version__)
print("\n# Out-of-bounds dims in a tuple silently become 'squeeze(None)':")
run((2, 3), (5,))
run((2, 1, 3), (5,))
run((1, 2, 1), (5,))
print("\n# Same out-of-bounds dim as a scalar int raises an opaque op-level error:")
run((2, 3), 5)
run((2, 3), -5)
run((2, 3), (-5,))
print("\n# In-bounds dims are unaffected (control):")
run((1, 2, 1), (0, 2))
run((2, 1, 3), (1,))
```

Actual output:

```
tvm: 0.24.dev0 | torch: 2.10.0+cu128

# Out-of-bounds dims in a tuple silently become 'squeeze(None)':
shape=(2, 3) dim=(5,) torch=IndexError tvm=OK (2, 3)
shape=(2, 1, 3) dim=(5,) torch=IndexError tvm=OK (2, 3)
shape=(1, 2, 1) dim=(5,) torch=IndexError tvm=OK (2,)

# Same out-of-bounds dim as a scalar int raises an opaque op-level error:
shape=(2, 3) dim=5 torch=IndexError tvm=InternalError
shape=(2, 3) dim=-5 torch=IndexError tvm=InternalError
shape=(2, 3) dim=(-5,) torch=IndexError tvm=InternalError

# In-bounds dims are unaffected (control):
shape=(1, 2, 1) dim=(0, 2) torch=OK (2,) tvm=OK (2,)
shape=(2, 1, 3) dim=(1,) torch=OK (2, 3) tvm=OK (2, 3)
```

### Triage

* needs-triage
* bug
* relax
* frontend/torch

Contributor guide

No contributing guide indexed for this repository

Research direction

Start in python/tvm/relax/frontend/torch/base_fx_graph_translator.py:2523-2543, focusing on _squeeze and its squeeze, squeeze.dim, and squeeze.dims dispatch. Run the provided reproduction to compare from_fx with native PyTorch for scalar and tuple dimensions. Done means every out-of-range dimension produces a clear frontend error, while the listed in-bounds control cases remain unchanged.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
compilers, frontend, machine-learning
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
76/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.