[Fix][Relax][Frontend][Torch] Keep zero-sized dims when reshaping

python/tvm/relax/frontend/torch/base_fx_graph_translator.py

@@ -129,6 +129,95 @@ 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_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.
+
+        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 or None 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."""
@@ -1946,7 +2035,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._emit_torch_reshape(x, new_shape)
 
     def _flatten(self, node: fx.Node) -> relax.Var:
         x = self.env[node.args[0]]
@@ -2286,14 +2375,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._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]]

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._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, size))
+        return self._emit_torch_reshape(x, size)
 
     ########## Symbolic Shape Constraints ##########
 

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)
@@ -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,56 @@ 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_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_roll():
     class Roll1(Module):
         def forward(self, x):
@@ -6942,6 +7001,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):