[Fix][Relax][Frontend][Torch] Keep zero-sized dims when reshaping

[hiyufan]

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.

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.

[sylvesterkaczmarek]

_torch_reshape_dims() only rewrites the first literal zero. PyTorch treats every zero in the target as a real zero, so an empty (0, 3) tensor can validly reshape(0, 0) to (0, 0). This helper produces [-1, 0]; Relax then interprets the remaining 0 as “copy dim 1”, yielding (0, 3). Could we preserve all zero positions, or otherwise handle multiple-zero targets, and add a regression for this case?

[hiyufan]

You are right, and it is worse than a missed case — both of these were silently wrong rather than an error:

                            torch        before
(0, 3).reshape(0, 0)        (0, 0)       (0, 3)
(0, 3, 5).reshape(0, 0, 0)  (0, 0, 0)    (0, 3, 5)

Working through your example sharpened the rule for me. It is not that only the first zero is rewritten — it is that which zero needs rewriting depends on the input. A literal 0 survives the copy rule exactly at a position whose input dimension is itself 0, because copying reproduces the zero that was asked for. So in your (0, 3) case position 0 is already fine and position 1 is the one that has to change; rewriting position 0 both wasted the single -1 and left the real problem in place.

That also explains why (2, 0, 4).reshape(0, 0, 4) passed in my original testing and gave me false confidence: there the input dimension under the second zero happens to be 0, so the copy rule reproduced it by luck. Same coincidence as the flatten on (0, 3) case in the PR body.

Pushed 0e11655:

• pick the positions that cannot survive, rather than the first zero;
• when several of them need -1, split the rewrite — each step turns one such position into a real 0, which lets the next step spell it as a literal. Only one -1 per reshape, so this is the part that needs more than one step.

Targets needing at most one rewrite — every case I have seen in practice — still emit a single reshape. Longest chain over everything I tested is 3.

Verified against numpy over 2132 valid reshapes (7 input shapes including (0,3), (3,0), (0,3,5), (2,0,4), (0,0,4); ranks 1–3; dims drawn from {0,1,2,3,4,5,6,12,15,20,24}), keeping only the targets numpy itself accepts: no mismatches.

Regression added as test_reshape_multiple_zero_sized_dims, covering your case plus the others in that family:

(0, 3).reshape(0, 0)        (3, 0).reshape(0, 0)
(0, 3, 5).reshape(0, 0, 0)  (2, 0, 4).reshape(0, 0, 4)

It fails against the previous head of this PR and passes now, so the specific hole you found is pinned. Full suites: 24 failed / 416 passed against 24 failed / 412 passed on clean main — the 24 are pre-existing in my environment and identical either way, the four extra passes are these tests.

This is the third distinct symptom of the same root cause, after the raise and the IndexError, which I think strengthens the note at the bottom of the PR description: the durable fix is an allowzero-style option on relax.op.reshape so the frontend can state the intent directly instead of encoding it in -1. That would also close the round-trip gap, which this does not. Happy to implement that instead if you would prefer it over this workaround — just say which shape you want and I will rework it.

[tlopex]

The None in current guard bypasses this rewrite whenever any input dimension is symbolic, even if another statically known zero already proves that the input is always empty.

For example, torch.export accepts an input with shape (batch, 0, 4) and x.reshape(0, 4), producing an output with shape (0, 4). Here current is [None, 0, 4], so the helper returns [0, 4] unchanged. Relax then interprets the first zero as “copy batch”, yielding (batch, 4) or rejecting the reshape because the element counts do not match.

Please allow the rewrite when the input contains a known zero even if its other dimensions are symbolic, and add a dynamic-batch regression test for this case.

[hiyufan]

Confirmed and fixed in d9b542d. You are right, and the wrong answer is not even empty:

(batch, 0, 4).reshape(0, 4)     torch (0, 4)     before: (s77, 4)

The guard was if None in current or 0 not in current: return [dims], and the None in current half is simply wrong reasoning on my part. I wrote it to mean "I cannot see the whole shape, so stay out of the way", but a single statically known zero already fixes the element count at zero whatever the symbols turn out to be — the other dimensions do not need to be known for that. Now the guard asks only for a known zero.

The symbolic dimensions do still matter one level down, where a position that needs holding is filled in. A literal there would be re-read as a copy, so the position is held as the expression it already is, which is not a literal and so is not substituted; a later step rewrites it once it has become a real zero.

I also removed a no-op the loop was emitting: after a rewrite it appended the target once more, so your case lowered to two reshapes. Once nothing needs rewriting the previous step has already produced the target, so the extra one is only emitted when no rewrite happened at all. Your case is a single reshape now:

lv: R.Tensor((0, 4), dtype="float32") = R.reshape(x, R.shape([0, 4]))

Verification

I swept it rather than checking your example alone, since hand-picked cases are what let the previous hole through — a case that passed by coincidence read as confirmation. 7 input layouts mixing symbolic and static dims around a known zero, against 9 targets, compared to torch:

after   63 compared   63 matched    0 mismatched
before  63 compared   27 matched   36 mismatched
        (3,0,4) dyn=(0,)  target=(0,4)     torch=(0, 4)     tvm=('s13', 4)
        (3,0,4) dyn=(0,)  target=(0,)      torch=(0,)       tvm=('s13',)
        (3,0,4) dyn=(0,)  target=(0,0,4)   torch=(0, 0, 4)  tvm=('s13', 0, 4)

The 2132-case static sweep is unchanged at zero mismatches, so the symbolic path did not cost the static one anything.

Test

test_reshape_zero_sized_dim_dynamic_batch, using Dim("batch", min=1, max=64) over a (3, 0, 4) example. It fails on the previous head of this PR and passes now. verify_model_numerically gained a dynamic_shapes passthrough to carry it — the expected-IR form is not available here for the reason in the PR description: re-parsing R.reshape(x, R.shape([0, 4])) applies the copy rule again, which is the same round-trip gap an allowzero on relax.op.reshape would close.

Full suites: 24 failed / 417 passed against 24 failed / 412 passed on clean main; the 24 are pre-existing in my environment and identical either way. ruff format --check and ruff check clean.