[Fix][Relax][Frontend][Torch] Compare shapes without == on symbolic dims in the reshape shortcut

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]

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):