From b6311562867791ea592bbb6650a6ce6b2f723d7d Mon Sep 17 00:00:00 2001
From: YidaWeng <84164537+YidaWeng@users.noreply.github.com>
Date: Mon, 16 Nov 2026 00:12:00 +0800
Subject: [PATCH] [FIX][TOPI] Broadcast values in index_put

topi.index_put flattened the broadcast index shape and read values at each flat position, but never broadcast values to that shape. Values that needed broadcasting (for example x[:, [0, 2]] = v with v of shape (2, 1)) therefore indexed past the end of values and wrote garbage.

Broadcast values to the indexed-region shape following NumPy rules before building the store, matching torch.index_put semantics for both the overwrite and accumulate paths.

Add topi and Relax regression tests covering (2, 1), (2,) and scalar values, with accumulate disabled and enabled.
---
 python/tvm/topi/index_put.py              | 16 +++++++-
 tests/python/relax/test_op_manipulate.py  | 40 ++++++++++++++++++++
 tests/python/topi/test_topi_ir_builder.py | 46 +++++++++++++++++++++++
 3 files changed, 101 insertions(+), 1 deletion(-)

diff --git a/python/tvm/topi/index_put.py b/python/tvm/topi/index_put.py
index 1b45515f7f4d..406d9af8698b 100644
--- a/python/tvm/topi/index_put.py
+++ b/python/tvm/topi/index_put.py
@@ -17,12 +17,15 @@
 # pylint: disable=invalid-name
 """IndexPut operator"""
 
+import tvm_ffi
+
 from tvm import te
 from tvm_tirx import tirx
 from tvm_tirx.script.ir_builder import IRBuilder
 from tvm_tirx.tirx.script import ir_builder as T
 
 from . import utils
+from .broadcast import broadcast_to
 
 
 def index_put(data, indices, values, accumulate=False):
@@ -38,7 +41,8 @@ def index_put(data, indices, values, accumulate=False):
         Index tensors are broadcast together following NumPy broadcasting rules.
 
     values : tvm.te.Tensor
-        The values to place at the specified indices.
+        The values to place at the specified indices.  ``values`` is
+        broadcast to the shape of the indexed region following NumPy rules.
 
     accumulate : bool, optional
         Whether to accumulate (add) values rather than replace.
@@ -84,6 +88,16 @@ def index_put(data, indices, values, accumulate=False):
                         raise ValueError(f"Cannot broadcast index shapes: {index_shapes}")
         broadcast_shape.append(max_dim)
 
+    # The indexed region has the broadcast index shape.  Broadcast ``values``
+    # to that shape so that the flat index used below is also valid when
+    # ``values`` needs broadcasting (e.g. x[:, [0, 2]] = v with v of shape
+    # (2, 1)).  Without this, the flat index reads past the end of values.
+    if len(values.shape) != len(broadcast_shape) or any(
+        not tvm_ffi.structural_equal(values_dim, region_dim)
+        for values_dim, region_dim in zip(values.shape, broadcast_shape)
+    ):
+        values = broadcast_to(values, broadcast_shape)
+
     # Compute total number of elements after broadcasting
     index_len = 1
     for dim in broadcast_shape:
diff --git a/tests/python/relax/test_op_manipulate.py b/tests/python/relax/test_op_manipulate.py
index 137778c4ea6e..bf24e6604fa9 100644
--- a/tests/python/relax/test_op_manipulate.py
+++ b/tests/python/relax/test_op_manipulate.py
@@ -15,8 +15,10 @@
 # specific language governing permissions and limitations
 # under the License.
 # ruff: noqa: E731, F841
+import numpy as np
 import pytest
 
+import tvm
 import tvm_tirx
 import tvm_tirx.testing
 from tvm import relax
@@ -3532,5 +3534,43 @@ def test_one_hot_infer_ty():
         )
 
 
+@pytest.mark.parametrize("accumulate", [False, True])
+def test_index_put_broadcast_values(accumulate):
+    bb = relax.BlockBuilder()
+    data = relax.Var("data", R.Tensor((2, 3), "float32"))
+    indices_0 = relax.Var("indices_0", R.Tensor((2, 1), "int32"))
+    indices_1 = relax.Var("indices_1", R.Tensor((2,), "int32"))
+    values = relax.Var("values", R.Tensor((2, 1), "float32"))
+
+    with bb.function("main", [data, indices_0, indices_1, values]):
+        with bb.dataflow():
+            out = bb.emit(
+                relax.op.index_put(data, relax.Tuple([indices_0, indices_1]), values, accumulate)
+            )
+            gv = bb.emit_output(out)
+        bb.emit_func_output(gv)
+
+    exe = tvm.compile(bb.get(), target="llvm")
+    vm = relax.VirtualMachine(exe, tvm_tirx.cpu())
+
+    data_np = np.arange(6, dtype="float32").reshape(2, 3)
+    indices_0_np = np.array([[0], [1]], dtype="int32")
+    indices_1_np = np.array([0, 2], dtype="int32")
+    values_np = np.array([[10.0], [20.0]], dtype="float32")
+    expected = data_np.copy()
+    if accumulate:
+        np.add.at(expected, (indices_0_np, indices_1_np), np.broadcast_to(values_np, (2, 2)))
+    else:
+        expected[indices_0_np, indices_1_np] = values_np
+
+    out = vm["main"](
+        *(
+            tvm_tirx.runtime.tensor(value)
+            for value in (data_np, indices_0_np, indices_1_np, values_np)
+        )
+    )
+    tvm_tirx.testing.assert_allclose(out.numpy(), expected)
+
+
 if __name__ == "__main__":
     tvm_tirx.testing.main()
diff --git a/tests/python/topi/test_topi_ir_builder.py b/tests/python/topi/test_topi_ir_builder.py
index 8b25bf3cefd5..11e2ec2dcde1 100644
--- a/tests/python/topi/test_topi_ir_builder.py
+++ b/tests/python/topi/test_topi_ir_builder.py
@@ -172,6 +172,52 @@ def test_gpu_imperative_tensors(kind):
 
 
 @pytest.mark.skipif(not tvm_tirx.runtime.enabled("llvm"), reason="LLVM is not enabled")
+@pytest.mark.parametrize(
+    "values_shape, accumulate",
+    [
+        ((2, 1), False),
+        ((2,), False),
+        ((), False),
+        ((2, 1), True),
+        ((2,), True),
+    ],
+)
+def test_index_put_broadcast_values(values_shape, accumulate):
+    data = np.arange(6, dtype="float32").reshape(2, 3)
+    indices_0 = np.array([[0], [1]], dtype="int32")
+    indices_1 = np.array([0, 2], dtype="int32")
+    if values_shape == (2,):
+        values = np.array([10.0, 20.0], dtype="float32")
+    elif values_shape == (2, 1):
+        values = np.array([[10.0], [20.0]], dtype="float32")
+    else:
+        values = np.array(7.0, dtype="float32")
+
+    indexed = np.broadcast_to(values, (2, 2))
+    expected = data.copy()
+    if accumulate:
+        np.add.at(expected, (indices_0, indices_1), indexed)
+    else:
+        expected[indices_0, indices_1] = indexed
+
+    x = te.placeholder(data.shape, "float32")
+    i = te.placeholder(indices_0.shape, "int32")
+    j = te.placeholder(indices_1.shape, "int32")
+    v = te.placeholder(values.shape, "float32")
+    out = topi.index_put(x, [i, j], v, accumulate=accumulate)
+
+    executable = tvm_tirx.tirx.compile(te.create_function([x, i, j, v, out]), target="llvm")
+    actual = tvm_tirx.runtime.empty(data.shape, "float32")
+    executable(
+        tvm_tirx.runtime.tensor(data),
+        tvm_tirx.runtime.tensor(indices_0),
+        tvm_tirx.runtime.tensor(indices_1),
+        tvm_tirx.runtime.tensor(values),
+        actual,
+    )
+    tvm_tirx.testing.assert_allclose(actual.numpy(), expected, atol=1e-5, rtol=1e-5)
+
+
 @pytest.mark.parametrize(
     "kind",
     [