[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.

[sylvesterkaczmarek]

The reshape handling now preserves multiple literal-zero target dimensions by chaining rewrites, and the symbolic-input case no longer bails merely because another dimension is unknown when a statically known zero already fixes the element count. The multiple-zero and dynamic-batch regressions cover both issues raised in review.

[tlopex]

None in target in _torch_reshape_chain still skips the fix when the target contains a symbolic dimension. For input (batch, 2, 0, 4), x.flatten(1, 2) imports as (batch, 2, 4) instead of (batch, 0, 4). Could you drop that check, preserve the symbolic dimensions, and add a regression where the target shape also contains batch?

[hiyufan]

Confirmed and fixed in 2ac8744. You are right, and it is the same mistake as the last round on the other half of the argument — I guarded against a symbol in the input last time, and left the identical bail-out for a symbol in the target.

input (batch, 2, 0, 4)              torch           before
x.flatten(1, 2)                     (batch, 0, 4)   (batch, 2, 4)
x.reshape(x.shape[0], 0, 4)         (batch, 0, 4)   (batch, 2, 4)
x.reshape(0, x.shape[0])            (0, batch)      (batch, batch)

The third one I found while checking your case and it is worse: both dimensions come back as batch.

The change

Dropping the check is the whole fix, exactly as you said:

-        if 0 not in target or -1 in target or None in target:
+        if 0 not in target or -1 in target:

The rewrite already leaves a non-literal alone at any position it does not touch, so the symbol is preserved without further work. The reasoning I got wrong is worth naming: I wrote that guard to mean "the target is not fully known, so stay out of the way", but relax only reads a literal as "copy the input dimension". A symbol is not a literal, so it was never at risk — and standing aside left the literal zero next to it being read as a copy, which is the one thing the rewrite exists to prevent.

Regression

test_reshape_zero_sized_dim_symbolic_target covers a symbol ahead of the zero, a target built from x.shape[0], and a zero ahead of the symbol. All three fail at the previous head:

FAILED test_reshape_zero_sized_dim_symbolic_target  (against d9b542d)

What I checked this time

Three rounds of this helper have each been "the case I happened to try passes", so I built the sweep around targets that can reach a symbol rather than around examples. 938 comparable cases — reshape and view over every target drawn from {0, 2, 4, x.shape[k]} at lengths 2 and 3, every valid flatten(start, end), and unflatten mixing zeros, -1 and symbols — each compared against the shape torch.export itself recorded, with symbol names canonicalised so (s13, 0, 4) and (batch, 0, 4) count as equal:

                     before        after
matched                 578          892
mismatched              360           46

cases failing after this change that did not fail before it:   0
cases this change repairs:                                   314

The 46 left are not this change:

• 40 are relax raising ValueError: Cannot use and / or / not operator to Expr, hint: use tvm.tirx.all / tvm.tirx.any for a target that mixes a symbol with a literal, e.g. x.reshape(x.shape[0], 0, x.shape[0]) on (batch, 0, 4). That reproduces byte-for-byte with this commit reverted, so it is a separate defect somewhere below the frontend. Happy to open an issue or a PR for it if you would like — it is a Python and/or applied to a PrimExpr, not a semantics question.
• 6 are my sweep printing s13 * T.int64(2) where torch prints 2*s13. Same expression: arith.Analyzer().simplify turns the one compound case, s13 * T.int64(2) * T.int64(4), into s13 * T.int64(8) against torch's 8*s13.

I should also say that my first version of that sweep reported 938/938 matched, which was wrong twice over: int() on a backed SymInt silently returns its hint, so every symbol compared as a constant, and the models were defined in a closure, which made torch.export refuse all 1000 dynamic cases and left only static ones being compared. It now prints the skip count by reason so an empty sweep cannot read as a clean one. Mentioning it because a green differential test that is not testing anything is exactly how the previous two rounds got past me.

Also

ruff check and ruff format --check clean on the touched files. test_frontend_from_exported_program.py has the same failure set with and without this change (9 either way, all pre-existing in my local torch 2.13 environment), verified by diffing the two lists rather than the counts.

Rebased onto 0eaf1cb while pushing, no conflicts. The local verification above was run before that rebase; my build tree needs the 3rdparty/tvm-ffi bump to compile against current main and is rebuilding now, so CI is the authority on the rebased tree until it finishes. If anything turns up there I will say so here rather than wait to be asked.

[hiyufan]

Retracting the caveat in my last comment: my build tree is current now (it needed the 3rdparty/tvm-ffi bump the rebase brought in), so everything above has been re-run on the rebased tree at 2ac8744, not only on the pre-rebase branch. Identical results:

sweep, rebased tree     938 compared, 892 matched, 46 mismatched
                        (40 relax raising, 6 my renderer) - same as before the rebase

test_frontend_from_exported_program.py
  with this commit      9 failures
  with it reverted      9 failures
  head-only failures    none

Nothing new turned up, so there is nothing to correct. Jenkins so far: docker, arm and wasm green, cpu and gpu still running.

[sylvesterkaczmarek]

Rechecked the current revision. The reshape helper now handles multiple zero-sized dimensions and symbolic targets, with numerical regressions covering the added cases. Current checks are green. Looks good from my side.

[tlopex]

RemoveRedundantReshape breaks the new reshape chain: for input (0, 3, 5), x.reshape(0, 0, 0) returns (0, 3, 5) after the pass. Combining the calls reinterprets the zeros as input-dimension copies and removes the reshape entirely. Could you have a change for that?

[hiyufan]

Confirmed and fixed in bebfc56. Reproduced your case exactly before touching anything:

x: (0, 3, 5), x.reshape(0, 0, 0)
  torch                (0, 0, 0)
  after import         (0, 0, 0)
  after the pass       (0, 3, 5)     both reshapes gone, the function returns x

Why it happens

relax.op.reshape resolves its target when the call is built — ConvertNewShapeToExpr replaces a literal 0 at position i with the input's dimension i. The pass rewrites reshape(reshape(x, s1), s2) to reshape(x, s2) without accounting for that, so an already-resolved target holding a literal zero asks for something different once it sits on a different input:

reshape(reshape(x, [0, 0, 5]), [0, 0, 0])    (0, 0, 0)
reshape(x, [0, 0, 0])                        (0, 3, 5)   all three copied back from x

The no-op rule then sees a reshape to the input's own shape and drops it, which is how both calls disappear.

This is not specific to the chain in this PR — any second target holding a literal zero over a differing intermediate dimension hits it — but the chain is what makes it reachable from a plain torch import, so it belongs here.

The change

Rather than reason about which zeros are safe to move, validate the rewrite: re-resolve the target against the new input and keep the fusion only when it comes back unchanged.

def _can_reparent(arg: Expr, output_shape: Expr) -> bool:
    try:
        reparented = relax.op.reshape(arg, output_shape)
    except Exception:
        return False          # reshape cannot resolve a 0 or -1 without a known input shape
    return tvm_ffi.structural_equal(reparented.args[1], output_shape)

It was 24 cases, not one

Your report was one case, so I re-ran the 938-case sweep from the last round with RemoveRedundantReshape applied after import:

                                    matched   mismatched
pass as it stands                       868           70
pass fixed                              892           46
no pass at all                          892           46

cases failing with the fixed pass that did not fail without it:   0
cases the fix repairs:                                           24

892/46 is exactly the no-pass result, so the pass no longer changes any outcome — and the 46 left are the same pre-existing ones from the last round (40 ValueError: Cannot use and / or / not operator to Expr, 6 my sweep printing s0*2 where torch prints 2*s0).

Tests

test_remove_redundant_reshape_pass_keeps_zero_sized_chain asserts the pair survives — it fails without the change, with At index 1 diff: 3 != 0. test_remove_redundant_reshape_pass_still_combines_without_zero_dims asserts an ordinary chain still collapses to a single reshape, so the guard cannot quietly turn the pass off. The three existing tests are untouched and pass.

Removed the file-level # ruff: noqa: F401 from the test file while there: the new tests use relax, so ruff reports it as an unused directive.

On the rebased tree

Rebased onto 7fca2e1. Everything above was re-run there after rebuilding, not only on the old base:

test_remove_redundant_reshape.py                                  5 passed
test_frontend_from_exported_program.py -k zero_sized              6 passed
test_frontend_from_exported_program.py, head vs branch reverted   9 failures either way, none head-only
sweep with the pass / without the pass, on 7fca2e1                892/46 both
ruff check + ruff format --check on all five touched files        clean

One thing worth passing on, since it cost me a confusing hour: after that rebase the numeric tests all failed with AttributeError: module 'tvm_ffi' has no attribute 'structural_map', because python/tvm/relax/utils.py on current main needs a newer apache-tvm-ffi than was installed. Reinstalling from 3rdparty/tvm-ffi fixed it. Nothing to do with this PR, but a stale environment reads exactly like a broken branch.

[sylvesterkaczmarek]

Re-reviewed current bebfc56 after the RemoveRedundantReshape regression fix. The pass now re-resolves the target against the proposed new input and only reparents when the resolved shape is structurally unchanged, which preserves literal-zero semantics while still collapsing ordinary reshape chains. The added zero-sized-chain and ordinary-chain regressions cover both sides. No blocker from me on this head.