Merge pull request #23467 from mattip/ufunc_at3

seberg · web-flow · commit 0a0346ade789 · 2023-03-27T11:55:44.000+02:00
BUG: in the fastest path form ufunc.at, properly increment args[2]
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
@@ -5935,6 +5935,9 @@ trivial_at_loop(PyArrayMethodObject *ufuncimpl, NPY_ARRAYMETHOD_FLAGS flags,
 
         res = ufuncimpl->contiguous_indexed_loop(
                 context, args, inner_size, steps, NULL);
+        if (args[2] != NULL) {
+            args[2] += (*inner_size) * steps[2];
+        }
     } while (res == 0 && iter->outer_next(iter->outer));
 
     if (res == 0 && !(flags & NPY_METH_NO_FLOATINGPOINT_ERRORS)) {
diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
@@ -2054,6 +2054,17 @@ def test_ufunc_at_negative(self):
         # If it is [-1, -1, -1, -100, 0] then the regular strided loop was used
         assert np.all(arr == [-1, -1, -1, -200, -1])
 
+    def test_ufunc_at_large(self):
+        # issue gh-23457
+        indices = np.zeros(8195, dtype=np.int16)
+        b = np.zeros(8195, dtype=float)
+        b[0] = 10
+        b[1] = 5
+        b[8192:] = 100
+        a = np.zeros(1, dtype=float)
+        np.add.at(a, indices, b)
+        assert a[0] == b.sum()
+
     def test_cast_index_fastpath(self):
         arr = np.zeros(10)
         values = np.ones(100000)
