Merge pull request #24150 from mattip/issue-24147

BUG: properly handle negative indexes in ufunc_at fast path
numpy · Jul 9, 2023 · c19ce9c · c19ce9c
2 parents 8b6024e + b774b3a
commit c19ce9c
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
@@ -5812,14 +5812,19 @@ trivial_at_loop(PyArrayMethodObject *ufuncimpl, NPY_ARRAYMETHOD_FLAGS flags,
         }
     }
 
-    npy_intp *inner_size = NpyIter_GetInnerLoopSizePtr(iter->outer);
-
     if (!(flags & NPY_METH_NO_FLOATINGPOINT_ERRORS)) {
         npy_clear_floatstatus_barrier((char *)context);
     }
 
     do {
-        args[1] = (char *) iter->outer_ptrs[0];
+        npy_intp *inner_size = NpyIter_GetInnerLoopSizePtr(iter->outer);
+        npy_intp * indxP = (npy_intp *)iter->outer_ptrs[0];
+        for (npy_intp i=0; i < *inner_size; i++) {
+            if (indxP[i] < 0) {
+                indxP[i] += iter->fancy_dims[0];
+            }
+        }
+        args[1] = (char *)indxP;
         steps[1] = iter->outer_strides[0];
 
         res = ufuncimpl->contiguous_indexed_loop(

diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
@@ -2224,6 +2224,14 @@ def test_ufunc_at_advanced(self):
         np.maximum.at(a, [0], 0)
         assert_equal(a, np.array([1, 2, 3]))
 
+    def test_at_negative_indexes(self):
+        a = np.arange(10)
+        indxs = np.array([-1, 1, -1, 2])
+        np.add.at(a, indxs, 1)
+        assert a[-1] == 11  # issue 24147
+        assert a[1] == 2
+        assert a[2] == 3
+
     def test_at_not_none_signature(self):
         # Test ufuncs with non-trivial signature raise a TypeError
         a = np.ones((2, 2, 2))