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..92caff08d4e0 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,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]] diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index ced0aa7b28bd..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, 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 ########## diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 7dc3c7356414..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) @@ -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):