[Fix][Relax][Torch] Respect zero-dimensional tensor dtype promotion

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

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