Repository navigation
[Bug][Relax] x[:, i] = v writes garbage values when v has to be broadcast #20925
Description
Expected behavior
v is broadcast to the indexed shape, here (2, 2), so both selected columns get [1, 2]. It works in eager and torch.export, so the converted program should give the same result.
Actual behavior
The output has wrong values, which change from run to run. The same happens for a v of shape (2,), for x[i, j] = v with index tensors that broadcast (i of shape (2, 1), j of shape (2,)), and for index_put_(..., accumulate=True). A v that already has the indexed shape, or a 0-d v, works.
eager: [[1.0, 0.0, 1.0], [2.0, 0.0, 2.0]]
tvm: [[1.0, 0.0, 2.0], [1.1350517561031018e-43, 0.0, 0.0]]
Possibly neither _index_put nor relax.op.index_put broadcasts values to the broadcast index shape, while topi.index_put reads values[k] for every index position.
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, v):
x = x.clone()
x[:, [0, 2]] = v
return x
x = torch.zeros(2, 3)
v = torch.tensor([[1.0], [2.0]])
args = (x, v)
print("eager:", M()(*args).tolist())
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: ", out[0].numpy().tolist())
Activity
GoyalIshaan commentedon Nov 6, 2026
Could I work on the
index_putvalues-broadcasting bug, including CPU regression tests for broadcasted values andaccumulate=True? Please let me know if this is available.AI-generated comment posted with Ishaan's authorization.
YidaWeng commentedon Nov 15, 2026
I had a working patch for this locally and just opened #20984. It broadcasts
valuesto the indexed-region shape intopi.index_put, where the flatvalues[k]access goes out of bounds, and adds topi plus Relax end-to-end regression tests. The new Relax test fails on current main and passes with the patch.I saw your earlier note asking to work on this. If you are actively working on it, let me know and we can consolidate; I am happy to review your patch or hand this over instead of keeping a competing PR.