From 11ea00aff73a378de37601baaccd5d8c24d2f6d9 Mon Sep 17 00:00:00 2001 From: mattip Date: Sun, 9 Jul 2023 10:10:09 +0300 Subject: [PATCH] BUG: properly handle negative indexes in ufunc_at fast path --- numpy/core/src/umath/ufunc_object.c | 11 ++++++++--- numpy/core/tests/test_ufunc.py | 8 ++++++++ 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c index 94cd73ef6c14..d992dd9602a6 100644 --- a/numpy/core/src/umath/ufunc_object.c +++ b/numpy/core/src/umath/ufunc_object.c @@ -5955,14 +5955,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 index 2ce6b4b69ea6..d6aa3e1056bf 100644 --- a/numpy/core/tests/test_ufunc.py +++ b/numpy/core/tests/test_ufunc.py @@ -2215,6 +2215,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))