From 8797d71139c842806f2ddd5411468976cba3d985 Mon Sep 17 00:00:00 2001 From: "kyler.xie" Date: Tue, 22 Sep 2026 09:46:51 +0800 Subject: [PATCH] [Fix][Relax][Torch] Respect zero-dimensional tensor dtype promotion Use tensor rank and torch.result_type when aligning binary operands. This preserves PyTorch semantics where a zero-dimensional float64 tensor does not promote a dimensioned float32 tensor. Keep the existing promotion behavior for unknown ranks and add regression coverage for both operand orders. --- .../frontend/torch/base_fx_graph_translator.py | 16 +++++++++++++--- .../test_frontend_from_exported_program.py | 18 ++++++++++++++++++ 2 files changed, 31 insertions(+), 3 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 3706411e4363..583a62f1df13 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -247,7 +247,9 @@ def _emit_torch_reshape(self, x, dims): return x @staticmethod - def _promote_common_dtype(lhs_dtype: str | None, rhs_dtype: str | None) -> str | None: + def _promote_common_dtype( + lhs_dtype: str | None, rhs_dtype: str | None, lhs_ndim: int, rhs_ndim: int + ) -> str | None: """Return the promoted dtype following PyTorch rules, or None if unsupported.""" import torch # type: ignore @@ -275,7 +277,13 @@ def _promote_common_dtype(lhs_dtype: str | None, rhs_dtype: str | None) -> str | if lhs_torch is None or rhs_torch is None: return None - promoted = torch.promote_types(lhs_torch, rhs_torch) + if lhs_ndim < 0 or rhs_ndim < 0: + promoted = torch.promote_types(lhs_torch, rhs_torch) + else: + promoted = torch.result_type( + torch.empty(() if lhs_ndim == 0 else (1,), dtype=lhs_torch, device="meta"), + torch.empty(() if rhs_ndim == 0 else (1,), dtype=rhs_torch, device="meta"), + ) return torch_to_tvm.get(promoted, None) @staticmethod @@ -698,7 +706,9 @@ def promote_binary_op_args(lhs, rhs): if isinstance(lhs_si, relax.TensorType) and isinstance( rhs_si, relax.TensorType ): - target_dtype = self._promote_common_dtype(lhs_si.dtype, rhs_si.dtype) + target_dtype = self._promote_common_dtype( + lhs_si.dtype, rhs_si.dtype, lhs_si.ndim, rhs_si.ndim + ) if target_dtype is not None: if lhs_si.dtype != target_dtype: lhs = self.block_builder.emit(relax.op.astype(lhs, target_dtype)) diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 22d5fee0d707..2c8b1550d8fe 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -1435,6 +1435,24 @@ def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tuple( verify_model(BinaryPromoteRHS(), example_args, {}, expected_promote_rhs) +@pytest.mark.parametrize("reverse", [False, True]) +def test_binary_dtype_promotion_zero_dim(reverse): + class Binary(Module): + def forward(self, lhs, rhs): + return lhs + rhs + + lhs = torch.tensor(1, dtype=torch.float64) + rhs = torch.tensor([1], dtype=torch.float32) + args = (rhs, lhs) if reverse else (lhs, rhs) + expected = Binary()(*args) + assert expected.dtype == torch.float32 + + mod = from_exported_program(export(Binary(), args)) + vm = relax.VirtualMachine(relax.build(mod, tvm.target.Target("llvm")), tvm.cpu()) + result = vm["main"](*[tvm.runtime.tensor(arg.numpy()) for arg in args])[0].numpy() + torch.testing.assert_close(torch.from_numpy(result), expected) + + operator_binary_2 = [ (operator.eq, R.equal), (operator.ne, R.not_equal),