Repository navigation
[Bug][Relax] torch.max(x, dim) on a bool tensor fails at runtime: "Unsupported input dtype: bool" #20942
Description
Expected behavior
torch.max(x, dim=0) on a bool tensor returns the values and the indices. It works in eager and torch.export, so the converted program should run.
Actual behavior
The conversion and tvm.compile pass, but the VM call fails. The same happens with keepdim=True and for x.max(dim=-1). An integer tensor works.
eager: [[True, False, True], [0, 0, 1]]
RuntimeError: Unsupported input dtype: bool
Possibly _max_dim lowers to relax.op.topk, whose CPU kernel does not accept bool.
Environment
- TVM: main @
2bf0532(0.28.dev0), built from source, targetllvm - 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, x):
return torch.max(x, dim=0)
x = torch.tensor([[True, False, False], [False, False, True]])
args = (x,)
print("eager:", [t.tolist() for t in M()(x)])
ep = torch.export.export(M(), args)
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(a.numpy()) for a in args])
print("tvm: ", [o.numpy().tolist() for o in out])
Activity
GoyalIshaan commentedon Nov 9, 2026
Could I work on the bool max-dim failure? I’d look at the topk lowering and CPU dtype handling, with regression tests for values and indices, negative dimensions, and keepdim.
AI disclosure: this allocation request was drafted and posted with AI assistance; I would do the implementation and submit the PR myself.