From 18fc5c711fab0d9e5fc3e6aeded39b0dcb1ff18a Mon Sep 17 00:00:00 2001 From: Chen Yufan Date: Thu, 17 Sep 2026 17:07:16 +0800 Subject: [PATCH] [Fix][Relax][Frontend][Torch] Compare shapes without `==` on symbolic dims in the reshape shortcut `_reshape` skips an identity reshape by comparing the input shape with the target as `list(current_shape) == list(dims)`. On a symbolic dimension `==` builds a PrimExpr instead of answering, and Python then asks it for a truth value: ValueError: Cannot use and / or / not operator to Expr It only surfaces when the ranks match, because list equality compares lengths first. So `x.reshape(x.shape[0], -1)` on a rank-3 input imported fine while `x.reshape(x.shape[0], 0, x.shape[0])` raised; in the reshape sweep from #20255 that was 40 of the 938 cases, all with a symbol at a position of the target that lined up with the same symbol in the input. `_same_dims` compares dimension by dimension: static dims as integers, symbolic dims with tvm_ffi.structural_equal, and a static-vs-symbolic pair as different. A genuine identity with a symbolic batch, `x.reshape(x.shape[0], 2, 4)` on `(batch, 2, 4)`, is still recognised and emits no reshape; an expression written differently on the two sides (`s*2` vs `2*s`) is treated as different, which costs a no-op reshape and never a wrong shape. Reshape sweep (938 targets over inputs that mix symbolic and zero dims, output shape compared with torch.export's, symbols canonicalised): before 892 matched 46 mismatched (40 raising, 6 the sweep's renderer) after 932 matched 6 mismatched ( 0 raising, 6 the sweep's renderer) no case fails after this change that did not fail before it; 40 repaired Test: `x.reshape(x.shape[0], 0, x.shape[0])` on a dynamic `(batch, 0, 4)` input builds and matches torch, and `x.reshape(x.shape[0], 2, 4)` on `(batch, 2, 4)` emits no reshape. Fails against the previous head with the ValueError above. --- .../torch/base_fx_graph_translator.py | 25 ++++++++++++++++- .../test_frontend_from_exported_program.py | 28 +++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) 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 3ff3b596af3b..07dbab21a163 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -2556,11 +2556,34 @@ def _reshape(self, node: fx.Node) -> relax.Var: # Skip identity reshape current_shape = self.shape_of(x) - if current_shape is not None and list(current_shape) == list(dims): + if current_shape is not None and self._same_dims(current_shape, dims): return x return self._emit_torch_reshape(x, dims) + def _same_dims(self, lhs, rhs) -> bool: + """Whether two shapes are the same, dimension by dimension. + + Plain ``list(lhs) == list(rhs)`` is not usable here: for a symbolic dimension + ``==`` builds a PrimExpr rather than answering, and Python then asks it for a + truth value and raises. That only surfaces when the ranks match, since list + equality compares lengths first, which is why ``x.reshape(x.shape[0], -1)`` + worked and ``x.reshape(x.shape[0], 0, x.shape[0])`` on a rank-3 input did not. + """ + if len(lhs) != len(rhs): + return False + for a, b in zip(lhs, rhs): + a_static, b_static = self._static_dim(a), self._static_dim(b) + if a_static is not None or b_static is not None: + if a_static != b_static: + return False + continue + if not (isinstance(a, tvm.ir.Expr) and isinstance(b, tvm.ir.Expr)): + return False + if not tvm_ffi.structural_equal(a, b): + return False + return True + def _reshape_as(self, node: fx.Node) -> relax.Var: args = self.retrieve_args(node) x = args[0] diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index d78629f5737b..90fe2c4486f1 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -5287,6 +5287,34 @@ def forward(self, x): verify_model_numerically(model, example_args, dynamic_shapes={"x": {0: batch}}) +def test_reshape_symbolic_target_same_rank(): + # The identity-reshape shortcut compared the two shapes with list equality. On a + # symbolic dimension `==` builds a PrimExpr instead of answering, and Python then + # raises asking it for a truth value -- but only when the ranks match, since list + # equality compares lengths first. So `x.reshape(x.shape[0], -1)` imported and + # `x.reshape(x.shape[0], 0, x.shape[0])` on a rank-3 input raised ValueError. + class SymbolAtBothEnds(Module): + def forward(self, x): + return x.reshape(x.shape[0], 0, x.shape[0]) + + class Identity(Module): + def forward(self, x): + return x.reshape(x.shape[0], 2, 4) + + batch = torch.export.Dim("batch", min=1, max=64) + verify_model_numerically( + SymbolAtBothEnds(), + (torch.randn(3, 0, 4, dtype=torch.float32),), + dynamic_shapes={"x": {0: batch}}, + ) + # A genuine identity is still recognised with a symbolic dimension: no reshape emitted. + x = torch.randn(3, 2, 4, dtype=torch.float32) + mod = from_exported_program(export(Identity(), (x,), dynamic_shapes={"x": {0: batch}})) + bindings = mod["main"].body.blocks[0].bindings + assert not any("reshape" in str(b.value) for b in bindings), mod["main"] + verify_model_numerically(Identity(), (x,), dynamic_shapes={"x": {0: batch}}) + + def test_roll(): class Roll1(Module): def forward(self, x):