hiyufan opened a new pull request, #20255:
URL: https://github.com/apache/tvm/pull/20255
### Problem
PyTorch reads a literal `0` in a target shape as a real zero-sized
dimension. `relax.op.reshape` reads it as *"copy the corresponding input
dimension"* — ONNX `Reshape` with `allowzero=0`. The torch frontend forwards
torch's shape unchanged, so any target shape holding a literal `0` is silently
reinterpreted.
```python
import torch
x = torch.randn(2, 0, 4)
x.reshape(0, 4) # torch (0, 4)
x.view(0, 4) # torch (0, 4)
torch.flatten(x) # torch (0,)
torch.randn(2, 0).unflatten(0, (2, -1)) # torch (2, 1, 0)
```
On `main` these import as:
| expression | input | torch | frontend on `main` |
| --- | --- | --- | --- |
| `x.reshape(0, 4)` | `(2, 0, 4)` | `(0, 4)` | `ValueError: Reshape expects
the new shape to be convertible…` |
| `x.view(0, 4)` | `(2, 0, 4)` | `(0, 4)` | same `ValueError` |
| `x.reshape(0)` | `(2, 0, 4)` | `(0,)` | same `ValueError` |
| `x.reshape(3, 0)` | `(0, 3)` | `(3, 0)` | same `ValueError` |
| `torch.flatten(x)` | `(2, 0, 4)` | `(0,)` | same `ValueError` |
| `torch.flatten(x)` | `(2, 3, 0)` | `(0,)` | same `ValueError` |
| `x.unflatten(0, (2, -1))` | `(2, 0)` | `(2, 1, 0)` | `IndexError: Index 2
out of bounds 2` |
| `torch.flatten(x)` | `(0, 3)` | `(0,)` | `(0,)` — happens to work |
The last row is why this is easy to miss: copying input dim 0 there gives
back the same `0` the literal asked for, so the one case people usually try
looks fine.
The `IndexError` comes from the same rule. `ConvertNewShapeToExpr` resolves
a zero with `array_ref.Set(i, shape_ty->values.value()[i])`, indexing the
*input* shape at the new shape's position, so a target of higher rank than the
input reads past the end.
Zero-sized tensors are not exotic in exported models — a detector with no
proposals, an empty batch, an empty mask — and they reach
`reshape`/`view`/`flatten` on ordinary code paths.
### Fix
When the input is statically empty, the dimension torch asks for can be
written as `-1` instead, whose inference yields `0`. `_torch_reshape_dims` does
that rewrite, applied where a torch-supplied target shape reaches
`relax.op.reshape`: `_reshape`, `_reshape_as`, `_flatten_impl`, `_unflatten`,
`_as_strided`.
The other `relax.op.reshape` call sites in the frontend derive their target
from the input's own shape, where "copy input dim" and the literal agree, so
they are left alone.
**The rewrite is deliberately narrow.** It only fires when the input is
statically empty. For a non-empty input torch rejects a zero in the target
outright, and rewriting it to `-1` there would turn an error into a silently
wrong shape:
```
input (2, 3), target [0, 2] torch: rejects
today raises ValueError <- correct
unconditional 0 -> -1 R.Tensor((3, 2)) <- wrong, and
silent
this PR (guard declines) raises ValueError <- unchanged
```
### Verification
All 17 shape cases I exercised now agree with PyTorch (6 previously raised).
Zero-sized behaviour of `squeeze`, `permute`, `expand`, `cat` and `sum` was
already correct and is unchanged.
Built with LLVM and ran the imported module: `x.reshape(0, 4)` on `(2, 0,
4)` builds, runs, and returns shape `(0, 4)`.
`tests/python/relax/test_frontend_from_fx.py` +
`tests/python/relax/test_frontend_from_exported_program.py`:
- clean `main`: 24 failed, 412 passed, 3 skipped
- with this change: 24 failed, **415** passed, 3 skipped
The 24 failures are pre-existing on `main` in my environment (`test_dtypes`
and friends), identical before and after. The three additional passes are the
new tests, which fail on `main` and pass with the fix.
`ruff format --check` and `ruff check` are clean.
### One thing I want to flag
The three tests run the imported module instead of comparing against an
expected TVMScript module, because **the resulting `IRModule` cannot be written
as TVMScript**. The frontend emits
```
lv: R.Tensor((0, 4), dtype="float32") = R.reshape(x, R.shape([0, 4]))
```
which executes correctly, but re-parsing it applies the copy rule again and
infers `(2, 4)`, so the annotation no longer matches and the module is rejected
as not well-formed. That round-trip gap lives in `relax.op.reshape`, not in the
frontend, and this PR does not try to close it.
So this fixes the observable behaviour but leaves the underlying ambiguity
in place. **The more complete fix is probably an `allowzero`-style option on
`relax.op.reshape`** (the ONNX frontend already carries `allowzero` and works
around the same rule by routing through a dynamic shape expression), with the
torch frontend opting in — that would also make the emitted IR round-trip. That
is a change to a core op's interface, so I did not want to make that call
unilaterally. **If you would prefer that shape, I am happy to implement it
instead and close this.**
Also worth noting: `_flatten_impl` is touched here and also by #20245. The
hunks are independent and should merge cleanly; happy to rebase either way.
---
This change was prepared with AI assistance (Claude). I have reviewed and
verified it, and can speak to it in review.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
