From 4d9bc2f296c1ec59d681fb94582b8ad55356fd92 Mon Sep 17 00:00:00 2001 From: Chen Yufan Date: Tue, 1 Sep 2026 18:36:46 +0800 Subject: [PATCH 1/3] [Fix][Relax][Frontend][Torch] Keep zero-sized dims when reshaping 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`. The torch frontend forwards torch's shape unchanged, so any target shape holding a literal `0` is silently reinterpreted: x = torch.randn(2, 0, 4) x.reshape(0, 4) # torch (0, 4) -> relax raises x.view(0, 4) # torch (0, 4) -> relax raises torch.flatten(x) # torch (0,) -> relax raises x.unflatten(0, ...) # rank grows -> IndexError from the zero-dim path `torch.flatten` on a `(0, 3)` input happens to work, because copying input dim 0 gives the same 0 the literal asked for. That coincidence is what hides the rest. When the input is statically empty, the dimension torch asks for can be written as `-1`, whose inference yields 0. Add `_torch_reshape_dims` and use it where a torch-supplied target shape reaches `relax.op.reshape`: `_reshape`, `_reshape_as`, `_flatten_impl`, `_unflatten`, `_as_strided`. Sites that derive the target from the input's own shape are unaffected, since "copy input dim" and the literal agree there. The rewrite is deliberately narrow. For a non-empty input torch rejects a zero in the target outright, and rewriting it to `-1` would silently produce a shape instead of surfacing that error, so those shapes are left alone. The three tests run the imported module rather than comparing against an expected TVMScript module: an `IRModule` holding `R.reshape(x, R.shape([0, 4]))` cannot be written as TVMScript, because re-parsing it applies the copy rule again and infers a different shape. That round-trip gap belongs to `relax.op.reshape` itself and is not addressed here. Co-authored-by: Claude --- .../torch/base_fx_graph_translator.py | 43 +++++++++++++++++-- .../torch/exported_program_translator.py | 4 +- .../test_frontend_from_exported_program.py | 30 +++++++++++++ 3 files changed, 72 insertions(+), 5 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 d600987cdd7b..e745ddb20e45 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -129,6 +129,43 @@ def shape_of(tensor): 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_dims(self, x, dims): + """Adapt a PyTorch target shape for ``relax.op.reshape``. + + 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``. When the input is statically empty, the + dimension PyTorch asks for can be written as ``-1`` instead, whose inference + yields ``0``. + + Every other case is left untouched. In particular, for a non-empty input PyTorch + rejects a zero in the target shape outright, and rewriting it to ``-1`` there would + silently produce a shape rather than surface the 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 + if 0 not in [self._static_dim(d) for d in shape]: + return dims + dims[target.index(0)] = -1 + return dims + @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.""" @@ -1946,7 +1983,7 @@ def _flatten_impl(self, x, start_dim, end_dim) -> relax.Var: + [flattened] + [shape[i] for i in range(end_dim + 1, len(shape))] ) - return self.block_builder.emit(relax.op.reshape(x, new_shape)) + return self.block_builder.emit(relax.op.reshape(x, self._torch_reshape_dims(x, new_shape))) def _flatten(self, node: fx.Node) -> relax.Var: x = self.env[node.args[0]] @@ -2286,14 +2323,14 @@ def _reshape(self, node: fx.Node) -> relax.Var: 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.block_builder.emit(relax.op.reshape(x, self._torch_reshape_dims(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.block_builder.emit(relax.op.reshape(x, self._torch_reshape_dims(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 ced0aa7b28bd..3ab7cb3864bd 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 @@ def _unflatten(self, node: fx.Node) -> relax.Var: 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.block_builder.emit(relax.op.reshape(x, self._torch_reshape_dims(x, new_shape))) ########## Creation ########## @@ -1477,7 +1477,7 @@ def _as_strided(self, node: fx.Node) -> relax.Var: f"size {size} is not supported" ) - return self.block_builder.emit(relax.op.reshape(x, size)) + return self.block_builder.emit(relax.op.reshape(x, self._torch_reshape_dims(x, size))) ########## Symbolic Shape Constraints ########## diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 7dc3c7356414..f2f2364055d8 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -5072,6 +5072,15 @@ def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple( 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): @@ -5233,6 +5242,19 @@ def main( 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_roll(): class Roll1(Module): def forward(self, x): @@ -6942,6 +6964,14 @@ def main( 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): From 0e11655c36511d5e477141b4f455920cbdb9df52 Mon Sep 17 00:00:00 2001 From: Chen Yufan Date: Thu, 3 Sep 2026 12:25:34 +0800 Subject: [PATCH 2/3] [Fix][Relax][Frontend][Torch] Handle targets with several zero-sized dims Review catch: rewriting only the first literal zero is not enough. A literal zero survives relax's copy rule exactly at a position whose input dimension is itself zero, so which zero needs rewriting depends on the input, and there can be more than one. (0, 3).reshape(0, 0) torch (0, 0) was (0, 3) (0, 3, 5).reshape(0, 0, 0) torch (0, 0, 0) was (0, 0... 3, 5) Both were silently wrong rather than an error. Rewrite the helper to pick the positions that cannot survive rather than the first zero, and to split the rewrite when several of them need `-1`: each step turns one such position into a real zero, which lets the next step spell it as a literal. Targets needing at most one rewrite -- every case seen in practice -- still emit a single reshape. Checked against numpy over 2132 valid reshapes (7 input shapes, ranks 1-3): no mismatches, longest chain 3 steps. Co-authored-by: Claude --- .../torch/base_fx_graph_translator.py | 78 ++++++++++++++----- .../torch/exported_program_translator.py | 4 +- .../test_frontend_from_exported_program.py | 21 +++++ 3 files changed, 82 insertions(+), 21 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 e745ddb20e45..3167ac85937b 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -141,30 +141,70 @@ def _static_dim(value): return const return None - def _torch_reshape_dims(self, x, dims): - """Adapt a PyTorch target shape for ``relax.op.reshape``. + 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``. When the input is statically empty, the - dimension PyTorch asks for can be written as ``-1`` instead, whose inference - yields ``0``. - - Every other case is left untouched. In particular, for a non-empty input PyTorch - rejects a zero in the target shape outright, and rewriting it to ``-1`` there would - silently produce a shape rather than surface the error. + 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. + + 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 + if 0 not in target or -1 in target or None in target: + return [dims] shape = self.shape_of(x) if shape is None: - return dims - if 0 not in [self._static_dim(d) for d in shape]: - return dims - dims[target.index(0)] = -1 - return dims + return [dims] + current = [self._static_dim(d) for d in shape] + if None in current or 0 not in current: + 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: + steps.append(dims) + 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: + kept = current[i] if i < len(current) else 1 + step.append(kept) + resulting.append(kept) + steps.append(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: @@ -1983,7 +2023,7 @@ def _flatten_impl(self, x, start_dim, end_dim) -> relax.Var: + [flattened] + [shape[i] for i in range(end_dim + 1, len(shape))] ) - return self.block_builder.emit(relax.op.reshape(x, self._torch_reshape_dims(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]] @@ -2323,14 +2363,14 @@ def _reshape(self, node: fx.Node) -> relax.Var: if current_shape is not None and list(current_shape) == list(dims): return x - return self.block_builder.emit(relax.op.reshape(x, self._torch_reshape_dims(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, self._torch_reshape_dims(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 3ab7cb3864bd..c28afed7aa1e 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 @@ def _unflatten(self, node: fx.Node) -> relax.Var: dim += len(x_shape) new_shape = x_shape[:dim] + sizes + x_shape[dim + 1 :] - return self.block_builder.emit(relax.op.reshape(x, self._torch_reshape_dims(x, new_shape))) + return self._emit_torch_reshape(x, new_shape) ########## Creation ########## @@ -1477,7 +1477,7 @@ def _as_strided(self, node: fx.Node) -> relax.Var: f"size {size} is not supported" ) - return self.block_builder.emit(relax.op.reshape(x, self._torch_reshape_dims(x, size))) + return self._emit_torch_reshape(x, size) ########## Symbolic Shape Constraints ########## diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index f2f2364055d8..db0f421257d8 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -5255,6 +5255,27 @@ def forward(self, x): 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_roll(): class Roll1(Module): def forward(self, x): From d9b542d3caf2d16b465911865520e776f86ff0b7 Mon Sep 17 00:00:00 2001 From: Chen Yufan Date: Fri, 4 Sep 2026 10:02:59 +0800 Subject: [PATCH 3/3] [Fix][Relax][Frontend][Torch] Rewrite zero dims when only some input dims are known Review catch: the guard declined the rewrite whenever any input dimension was symbolic, so a statically known zero next to a dynamic batch went unhandled. (batch, 0, 4).reshape(0, 4) torch (0, 4) was (s77, 4) Silently wrong again, and the wrong shape is not even empty. One known zero fixes the element count at zero whatever the symbols turn out to be, so require only that -- symbolic dimensions elsewhere are held as they stand, since a non-literal is not read as a copy, and rewritten in a later step once they have become real zeros. Also drop the trailing no-op reshape the loop emitted after a rewrite: once no position needs rewriting, the previous step has already produced the target. The dynamic case above now lowers to a single reshape rather than two. Checked against torch over 63 combinations of symbolic and static dims carrying a known zero (7 input layouts, 9 targets): 63 matched, against 27 before this change. The 2132-case static sweep is unchanged at no mismatches. Co-authored-by: Claude --- .../torch/base_fx_graph_translator.py | 22 ++++++++++++++----- .../test_frontend_from_exported_program.py | 20 +++++++++++++++-- 2 files changed, 35 insertions(+), 7 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 3167ac85937b..92caff08d4e0 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -167,8 +167,13 @@ def _torch_reshape_chain(self, x, 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 None in current or 0 not in current: + 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 = [] @@ -179,7 +184,11 @@ def _torch_reshape_chain(self, x, dims): if t == 0 and not (i < len(current) and current[i] == 0) ] if not unusable: - steps.append(dims) + 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 = [], [] @@ -194,10 +203,13 @@ def _torch_reshape_chain(self, x, dims): step.append(0) resulting.append(0) else: - kept = current[i] if i < len(current) else 1 - step.append(kept) - resulting.append(kept) + # 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): diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index db0f421257d8..7cfe1ac20439 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) @@ -5276,6 +5276,22 @@ def forward(self, x): 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_roll(): class Roll1(Module): def forward(self, x):