[Bug][Relax] torch.where with branches of different dtypes fails in from_exported_program #20940

Expected behavior

torch.where promotes its two branches to a common dtype, e.g. int64 and float32 give float32. It works in eager and torch.export, so from_exported_program should convert it. torch.select_scatter with a src of another dtype also hits this, since it is decomposed into where.

Actual behavior

The conversion fails, with either branch order. With bool and float32 branches the conversion passes, but tvm.compile fails with x and y must have the same dtype. Branches of the same dtype work.

eager: [1.0, 1.5] torch.float32
TypeError: Binary operators must have the same datatype for both operands.  However, ... "relax.where" ... uses datatype T.int64 on the LHS (Type of R.Tensor((2,), dtype="int64")), and datatype T.float32 on the RHS (Type of R.Tensor((2,), dtype="float32")).

Possibly _where passes both branches to relax.op.where as is, without the dtype promotion that _binary_op does.

Environment

• TVM: main @ 2bf0532 (0.28.dev0), built from source, target llvm
• PyTorch: 2.12.0+cu130
• Python: 3.12, Ubuntu x86_64

Steps to reproduce

import torch
import tvm
import tvm_tirx
from tvm import relax
from tvm.relax.frontend.torch import from_exported_program


class M(torch.nn.Module):
    def forward(self, c, a, b):
        return torch.where(c, a, b)


c = torch.tensor([True, False])
a = torch.tensor([1, 2])
b = torch.tensor([0.5, 1.5])
print("eager:", M()(c, a, b).tolist(), M()(c, a, b).dtype)

ep = torch.export.export(M(), (c, a, b))
mod = from_exported_program(ep)
ex = tvm.compile(mod, target="llvm")
vm = relax.VirtualMachine(ex, tvm_tirx.cpu())
out = vm["main"](*[tvm_tirx.runtime.tensor(t.numpy()) for t in (c, a, b)])
print("tvm:  ", out[0].numpy().tolist(), out[0].numpy().dtype)