fix reduction

Author: xadupre

onnxscript/function_libs/torch_lib/ops/core.py

@@ -4345,7 +4345,7 @@ def _1dint(i: int):
     assert len(n_none) == 1, f"Unable to handle that case: n_none={n_none}"
     unsq = op.Unsqueeze(indices[n_none[0]], _1dint(1))
     if n_none[0] == 0:
-        return op.ScatterND(x, unsq, values)
+        return op.ScatterND(x, unsq, values, reduction="add" if accumulate else "none")
 
     perm = list(range(len(x.shape)))
     perm[n_none[0]], perm[0] = perm[0], perm[n_none[0]]

tests/function_libs/torch_lib/e2e_ops_tests.py

@@ -287,6 +287,24 @@ def forward(self, x, index, update):
         )
         _testing.assert_onnx_program(onnx_program)
 
+    def test_index_put_55_2_25(self):
+        class Model(torch.nn.Module):
+            def forward(self, x, index, update):
+                return torch.ops.aten.index_put(x, [index], update, accumulate=True)
+
+        x = torch.ones((6, 5), dtype=torch.float32)
+        index = torch.tensor([4, 3], dtype=torch.int64)
+        update = (torch.arange(10) + 10).reshape((2, -1)).to(torch.float32)
+        onnx_program = torch.onnx.export(
+            Model(),
+            (x, index, update),
+            input_names=["x", "index", "update"],
+            output_names=["output"],
+            opset_version=18,
+            dynamo=True,
+        )
+        _testing.assert_onnx_program(onnx_program)
+
     def test_index_put_scatter_nd(self):
         class Model(torch.nn.Module):
             def forward(self, x, index, update):