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