This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new f986a66336 [Fix][Relax][Frontend][Torch] Keep zero-sized dims when
reshaping (#20255)
f986a66336 is described below
commit f986a66336bd3373c1d972e240feddc79b0552ba
Author: Chen Yufan <[email protected]>
AuthorDate: Wed Sep 16 02:40:40 2026 +0800
[Fix][Relax][Frontend][Torch] Keep zero-sized dims when reshaping (#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.
---------
Co-authored-by: Claude <[email protected]>
---
.../frontend/torch/base_fx_graph_translator.py | 100 ++++++++++++++++++++-
.../frontend/torch/exported_program_translator.py | 4 +-
.../relax/transform/remove_redundant_reshape.py | 25 +++++-
.../relax/test_frontend_from_exported_program.py | 94 ++++++++++++++++++-
.../python/relax/test_remove_redundant_reshape.py | 52 ++++++++++-
5 files changed, 266 insertions(+), 9 deletions(-)
diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
index 5b723350cd..3ff3b596af 100644
--- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
+++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
@@ -152,6 +152,100 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
return tensor.shape
raise ValueError(f"Unsupported type: {type(tensor)}")
+ @staticmethod
+ def _static_dim(value):
+ """Return ``value`` as a Python int when it is a compile-time
constant, else ``None``."""
+ if isinstance(value, bool):
+ return None
+ if isinstance(value, int):
+ return value
+ const = getattr(value, "value", None)
+ if isinstance(const, int) and not isinstance(const, bool):
+ return const
+ return None
+
+ def _torch_reshape_chain(self, x, dims):
+ """Return the relax reshape targets that reproduce PyTorch's ``dims``.
+
+ 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", which is
+ ONNX ``Reshape`` with ``allowzero=0``. A literal ``0`` therefore only
survives at a
+ position whose input dimension is itself ``0``; anywhere else it
silently becomes
+ that input dimension.
+
+ A position that does not survive can be written as ``-1``, whose
inference yields
+ ``0`` for an empty input. Only one ``-1`` is allowed per reshape, so
when several
+ positions need it the rewrite is split: each step turns one of them
into a real
+ ``0``, which lets the next step spell that position as a literal.
Targets with at
+ most one such position - every case seen in practice - stay a single
reshape.
+
+ Only a literal is read as a copy, so a symbolic dimension in the
target is carried
+ through untouched and does not stop the literal zeros beside it from
being
+ rewritten. ``x.flatten(1, 2)`` on ``(batch, 2, 0, 4)`` asks for
``(batch, 0, 4)``,
+ where the zero still has to survive.
+
+ Shapes that do not need the rewrite are returned unchanged. In
particular, for a
+ non-empty input PyTorch rejects a zero in the target outright, and
rewriting it
+ would produce a shape rather than surface that error.
+ """
+ dims = list(dims)
+ target = [self._static_dim(d) for d in dims]
+ if 0 not in target or -1 in target:
+ return [dims]
+ shape = self.shape_of(x)
+ if shape is None:
+ return [dims]
+ shape = list(shape)
+ current = [self._static_dim(d) for d in shape]
+ if 0 not in current:
+ # Without a statically known zero the input is not known to be
empty, and
+ # PyTorch rejects a zero in the target for a non-empty input. A
symbolic
+ # dimension elsewhere does not change that: one known zero already
fixes the
+ # element count at zero whatever the symbols turn out to be.
+ return [dims]
+
+ steps = []
+ while True:
+ unusable = [
+ i
+ for i, t in enumerate(target)
+ if t == 0 and not (i < len(current) and current[i] == 0)
+ ]
+ if not unusable:
+ if not steps:
+ # Nothing needed rewriting; emit the target as given.
+ steps.append(dims)
+ # Otherwise the last step already produced the target shape,
since every
+ # remaining zero now sits over an input dimension that is zero
as well.
+ return steps
+ rewritten = unusable[0]
+ step, resulting = [], []
+ for i, t in enumerate(target):
+ if i == rewritten:
+ step.append(-1)
+ resulting.append(0)
+ elif t != 0:
+ step.append(dims[i])
+ resulting.append(t)
+ elif i < len(current) and current[i] == 0:
+ step.append(0)
+ resulting.append(0)
+ else:
+ # Hold this position as it stands -- a symbolic dimension
included, since
+ # it is not a literal and so is not read as a copy -- and
rewrite it in a
+ # later step, once it has become a real zero.
+ step.append(shape[i] if i < len(shape) else 1)
+ resulting.append(current[i] if i < len(current) else 1)
+ steps.append(step)
+ shape = [0 if i == rewritten else step[i] for i in
range(len(step))]
+ current = resulting
+
+ def _emit_torch_reshape(self, x, dims):
+ """Emit the reshape(s) giving ``dims`` PyTorch's meaning. See
_torch_reshape_chain."""
+ for step in self._torch_reshape_chain(x, dims):
+ x = self.block_builder.emit(relax.op.reshape(x, step))
+ return x
+
@staticmethod
def _promote_common_dtype(lhs_dtype: str | None, rhs_dtype: str | None) ->
str | None:
"""Return the promoted dtype following PyTorch rules, or None if
unsupported."""
@@ -2125,7 +2219,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
+ [flattened]
+ [shape[i] for i in range(end_dim + 1, rank)]
)
- return self.block_builder.emit(relax.op.reshape(x, new_shape))
+ return self._emit_torch_reshape(x, new_shape)
def _flatten(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
@@ -2465,14 +2559,14 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
if current_shape is not None and list(current_shape) == list(dims):
return x
- return self.block_builder.emit(relax.op.reshape(x, dims))
+ return self._emit_torch_reshape(x, dims)
def _reshape_as(self, node: fx.Node) -> relax.Var:
args = self.retrieve_args(node)
x = args[0]
other = args[1]
dims = self.shape_of(other)
- return self.block_builder.emit(relax.op.reshape(x, dims))
+ return self._emit_torch_reshape(x, dims)
def _scatter(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py
b/python/tvm/relax/frontend/torch/exported_program_translator.py
index 2029dcc121..18c0ef96d8 100644
--- a/python/tvm/relax/frontend/torch/exported_program_translator.py
+++ b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -1191,7 +1191,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
dim += len(x_shape)
new_shape = x_shape[:dim] + sizes + x_shape[dim + 1 :]
- return self.block_builder.emit(relax.op.reshape(x, new_shape))
+ return self._emit_torch_reshape(x, new_shape)
########## Creation ##########
@@ -1477,7 +1477,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
f"size {size} is not supported"
)
- return self.block_builder.emit(relax.op.reshape(x, size))
+ return self._emit_torch_reshape(x, size)
########## Symbolic Shape Constraints ##########
diff --git a/python/tvm/relax/transform/remove_redundant_reshape.py
b/python/tvm/relax/transform/remove_redundant_reshape.py
index 8d3709e167..0cc1e11900 100644
--- a/python/tvm/relax/transform/remove_redundant_reshape.py
+++ b/python/tvm/relax/transform/remove_redundant_reshape.py
@@ -27,6 +27,27 @@ from tvm.relax.dpl import is_op, rewrite_call, wildcard
from . import function_pass
+def _can_reparent(arg: Expr, output_shape: Expr) -> bool:
+ """Whether ``reshape(arg, output_shape)`` still asks for ``output_shape``.
+
+ A literal ``0`` in a reshape target means "copy the corresponding input
dimension",
+ and ``relax.op.reshape`` resolves it against the input it is handed.
Moving such a
+ target onto a different input therefore changes what it asks for. For
+ ``x: (0, 3, 5)``, ``reshape(reshape(x, [0, 0, 5]), [0, 0, 0])`` is ``(0,
0, 0)``,
+ while the combined ``reshape(x, [0, 0, 0])`` copies all three dimensions
back and is
+ ``(0, 3, 5)``.
+
+ Rather than reason about which zeros are safe, re-resolve the target
against the new
+ input and keep the rewrite only when it comes back unchanged.
+ """
+ try:
+ reparented = relax.op.reshape(arg, output_shape)
+ except Exception: # pylint: disable=broad-except
+ # reshape cannot resolve a 0 or -1 without a known input shape.
+ return False
+ return tvm_ffi.structural_equal(reparented.args[1], output_shape)
+
+
@function_pass(opt_level=0)
class RemoveRedundantReshape:
"""
@@ -70,7 +91,9 @@ class RemoveRedundantReshape:
if self.repeated_reshape in matches:
output_shape = matches[self.repeated_reshape].args[1]
- return relax.op.reshape(arg, output_shape)
+ if _can_reparent(arg, output_shape):
+ return relax.op.reshape(arg, output_shape)
+ return expr
elif self.no_op_reshape in matches:
output_shape = matches[self.no_op_reshape].args[1]
diff --git a/tests/python/relax/test_frontend_from_exported_program.py
b/tests/python/relax/test_frontend_from_exported_program.py
index b70774be65..43065d9c13 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -62,12 +62,12 @@ def verify_model(
tvm.ir.assert_structural_equal(mod, expected, map_free_vars=map_free_vars)
-def verify_model_numerically(torch_model, example_args, rtol=1e-7, atol=1e-7):
+def verify_model_numerically(torch_model, example_args, rtol=1e-7, atol=1e-7,
dynamic_shapes=None):
"""Verify model by comparing numerical outputs between PyTorch and TVM."""
with torch.no_grad():
pytorch_output = torch_model(*example_args)
- exported_program = export(torch_model, args=example_args)
+ exported_program = export(torch_model, args=example_args,
dynamic_shapes=dynamic_shapes)
mod = from_exported_program(exported_program)
target = tvm.target.Target("llvm")
ex = relax.build(mod, target)
@@ -5315,6 +5315,15 @@ def test_flatten():
verify_model(Flatten(), example_args, {}, expected1)
+def test_flatten_zero_sized_dim():
+ class Flatten(Module):
+ def forward(self, x):
+ return torch.flatten(x)
+
+ verify_model_numerically(Flatten(), (torch.randn(2, 0, 4,
dtype=torch.float32),))
+ verify_model_numerically(Flatten(), (torch.randn(2, 3, 0,
dtype=torch.float32),))
+
+
def test_meshgrid():
class Meshgrid1(Module):
def forward(self, input1, input2):
@@ -5476,6 +5485,79 @@ def test_reshape_as():
verify_model(ReshapeAs(), example_args, {}, expected1)
+def test_reshape_zero_sized_dim():
+ class Reshape(Module):
+ def forward(self, x):
+ return x.reshape(0, 4)
+
+ class ReshapeTrailing(Module):
+ def forward(self, x):
+ return x.reshape(3, 0)
+
+ verify_model_numerically(Reshape(), (torch.randn(2, 0, 4,
dtype=torch.float32),))
+ verify_model_numerically(ReshapeTrailing(), (torch.randn(0, 3,
dtype=torch.float32),))
+
+
+def test_reshape_multiple_zero_sized_dims():
+ # A literal zero only survives relax's copy rule at a position whose input
dimension is
+ # itself zero, so targets holding several zeros need more than one
position rewritten.
+ class TwoZeros(Module):
+ def forward(self, x):
+ return x.reshape(0, 0)
+
+ class ThreeZeros(Module):
+ def forward(self, x):
+ return x.reshape(0, 0, 0)
+
+ class ZeroPastInputRank(Module):
+ def forward(self, x):
+ return x.reshape(0, 0, 4)
+
+ verify_model_numerically(TwoZeros(), (torch.randn(0, 3,
dtype=torch.float32),))
+ verify_model_numerically(TwoZeros(), (torch.randn(3, 0,
dtype=torch.float32),))
+ verify_model_numerically(ThreeZeros(), (torch.randn(0, 3, 5,
dtype=torch.float32),))
+ verify_model_numerically(ZeroPastInputRank(), (torch.randn(2, 0, 4,
dtype=torch.float32),))
+
+
+def test_reshape_zero_sized_dim_dynamic_batch():
+ # One statically known zero fixes the element count at zero whatever the
symbolic
+ # dimension turns out to be, so the literal zero in the target still has
to survive.
+ # Reading it as "copy the batch" gives a non-empty shape that torch never
produces.
+ class Reshape(Module):
+ def forward(self, x):
+ return x.reshape(0, 4)
+
+ batch = torch.export.Dim("batch", min=1, max=64)
+ verify_model_numerically(
+ Reshape(),
+ (torch.randn(3, 0, 4, dtype=torch.float32),),
+ dynamic_shapes={"x": {0: batch}},
+ )
+
+
+def test_reshape_zero_sized_dim_symbolic_target():
+ # Only a literal is read as "copy the input dimension", so a symbolic
dimension in the
+ # target is not one: it is carried through, and the literal zero beside it
still has to
+ # be rewritten. Skipping the rewrite because the target is not fully
static reads that
+ # zero as a copy and gives a non-empty shape torch never produces.
+ class Flatten(Module):
+ def forward(self, x):
+ return x.flatten(1, 2)
+
+ class KeepBatch(Module):
+ def forward(self, x):
+ return x.reshape(x.shape[0], 0, 4)
+
+ class ZeroBeforeBatch(Module):
+ def forward(self, x):
+ return x.reshape(0, x.shape[0])
+
+ batch = torch.export.Dim("batch", min=1, max=64)
+ example_args = (torch.randn(3, 2, 0, 4, dtype=torch.float32),)
+ for model in (Flatten(), KeepBatch(), ZeroBeforeBatch()):
+ verify_model_numerically(model, example_args, dynamic_shapes={"x": {0:
batch}})
+
+
def test_roll():
class Roll1(Module):
def forward(self, x):
@@ -7265,6 +7347,14 @@ def test_unflatten():
verify_model(Unflatten1(), example_args, {}, Expected)
+def test_unflatten_zero_sized_dim():
+ class Unflatten(Module):
+ def forward(self, x):
+ return x.unflatten(0, (2, -1))
+
+ verify_model_numerically(Unflatten(), (torch.randn(2, 0,
dtype=torch.float32),))
+
+
def test_gather():
class Gather0(Module):
def forward(self, data, indices):
diff --git a/tests/python/relax/test_remove_redundant_reshape.py
b/tests/python/relax/test_remove_redundant_reshape.py
index f5a865940b..bf8c790e01 100644
--- a/tests/python/relax/test_remove_redundant_reshape.py
+++ b/tests/python/relax/test_remove_redundant_reshape.py
@@ -14,7 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: F401
"""
Test relax transform - Eliminate redundant reshape operations
@@ -113,5 +112,56 @@ def test_remove_redundant_reshape_pass_three_arg():
_run_pass_compare_output(Before, Expected)
+def _return_shape(mod):
+ ret = mod["main"].ret_ty
+ field = ret.fields[0] if hasattr(ret, "fields") else ret
+ return [int(dim) for dim in field.shape]
+
+
+def test_remove_redundant_reshape_pass_keeps_zero_sized_chain():
+ # A literal 0 in a reshape target means "copy the corresponding input
dimension",
+ # and relax.op.reshape resolves it against the input it is handed.
Combining these
+ # two calls re-reads the zeros against x and asks for (0, 3, 5) instead of
the
+ # (0, 0, 0) the pair produces, so the pair has to survive the pass.
+ #
+ # Built with the block builder rather than TVMScript because reshape
resolves its
+ # target at construction: the printed R.shape([0, 0, 5]) is the resolved
shape, and
+ # parsing it back would resolve those zeros a second time.
+ bb = relax.BlockBuilder()
+ x = relax.Var("x", relax.TensorType([0, 3, 5], "float32"))
+ with bb.function("main", [x]):
+ with bb.dataflow():
+ lv = bb.emit(relax.op.reshape(x, [0, -1, 5]))
+ gv = bb.emit_output(relax.op.reshape(lv, [0, 0, -1]))
+ bb.emit_func_output(gv)
+ before = bb.get()
+ assert _return_shape(before) == [0, 0, 0]
+
+ after = DeadCodeElimination()(RemoveRedundantReshape()(before))
+ assert _return_shape(after) == [0, 0, 0], (
+ "combining the reshapes re-read the literal zeros against x and
changed the shape"
+ )
+
+
+def test_remove_redundant_reshape_pass_still_combines_without_zero_dims():
+ # The guard above must not stop the pass doing its job on an ordinary
chain.
+ bb = relax.BlockBuilder()
+ x = relax.Var("x", relax.TensorType([1, 1001, 1, 1], "float32"))
+ with bb.function("main", [x]):
+ with bb.dataflow():
+ lv = bb.emit(relax.op.reshape(x, [1, 1001, 1]))
+ gv = bb.emit_output(relax.op.reshape(lv, [1, 1001]))
+ bb.emit_func_output(gv)
+ after = DeadCodeElimination()(RemoveRedundantReshape()(bb.get()))
+ assert _return_shape(after) == [1, 1001]
+ reshapes = [
+ binding
+ for block in after["main"].body.blocks
+ for binding in block.bindings
+ if isinstance(binding.value, tvm.relax.Call)
+ ]
+ assert len(reshapes) == 1, "the chain without zeros should still collapse
to one reshape"
+
+
if __name__ == "__main__":
tvm.testing.main()
