From 7b5e51a1acadb8730a102f1d867a268afe518297 Mon Sep 17 00:00:00 2001 From: Sebastian Berg Date: Fri, 28 Jul 2023 21:25:42 +0200 Subject: [PATCH] BUG: Further fixes to indexing loop and added tests This is a follow-up to gh-24272 which missed a few files. --- numpy/core/src/umath/_umath_tests.c.src | 11 +++++++--- .../src/umath/loops_arithm_fp.dispatch.c.src | 22 ++++++++++++++----- .../src/umath/loops_arithmetic.dispatch.c.src | 22 ++++++++++++++----- .../src/umath/loops_minmax.dispatch.c.src | 11 +++++++--- numpy/core/tests/test_ufunc.py | 22 +++++++++++++------ 5 files changed, 63 insertions(+), 25 deletions(-) diff --git a/numpy/core/src/umath/_umath_tests.c.src b/numpy/core/src/umath/_umath_tests.c.src index b427991e5463..b9e192706d00 100644 --- a/numpy/core/src/umath/_umath_tests.c.src +++ b/numpy/core/src/umath/_umath_tests.c.src @@ -375,13 +375,18 @@ INT32_negative_indexed(PyArrayMethod_Context *NPY_UNUSED(context), npy_intp const *steps, NpyAuxData *NPY_UNUSED(func)) { char *ip1 = args[0]; - char *indx = args[1]; + char *indxp = args[1]; npy_intp is1 = steps[0], isindex = steps[1]; npy_intp n = dimensions[0]; + npy_intp shape = steps[3]; npy_intp i; int32_t *indexed; - for(i = 0; i < n; i++, indx += isindex) { - indexed = (int32_t *)(ip1 + is1 * *(npy_intp *)indx); + for(i = 0; i < n; i++, indxp += isindex) { + npy_intp indx = *(npy_intp *)indxp; + if (indx < 0) { + indx += shape; + } + indexed = (int32_t *)(ip1 + is1 * indx); if (i == 3) { *indexed = -200; } else { diff --git a/numpy/core/src/umath/loops_arithm_fp.dispatch.c.src b/numpy/core/src/umath/loops_arithm_fp.dispatch.c.src index b72db5846f70..7ba3981e8119 100644 --- a/numpy/core/src/umath/loops_arithm_fp.dispatch.c.src +++ b/numpy/core/src/umath/loops_arithm_fp.dispatch.c.src @@ -258,14 +258,19 @@ NPY_NO_EXPORT int NPY_CPU_DISPATCH_CURFX(@TYPE@_@kind@_indexed) (PyArrayMethod_Context *NPY_UNUSED(context), char * const*args, npy_intp const *dimensions, npy_intp const *steps, NpyAuxData *NPY_UNUSED(func)) { char *ip1 = args[0]; - char *indx = args[1]; + char *indxp = args[1]; char *value = args[2]; npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2]; + npy_intp shape = steps[3]; npy_intp n = dimensions[0]; npy_intp i; @type@ *indexed; - for(i = 0; i < n; i++, indx += isindex, value += isb) { - indexed = (@type@ *)(ip1 + is1 * *(npy_intp *)indx); + for(i = 0; i < n; i++, indxp += isindex, value += isb) { + npy_intp indx = *(npy_intp *)indxp; + if (indx < 0) { + indx += shape; + } + indexed = (@type@ *)(ip1 + is1 * indx); *indexed = *indexed @OP@ *(@type@ *)value; } return 0; @@ -650,14 +655,19 @@ NPY_NO_EXPORT int NPY_CPU_DISPATCH_CURFX(@TYPE@_@kind@_indexed) (PyArrayMethod_Context *NPY_UNUSED(context), char * const*args, npy_intp const *dimensions, npy_intp const *steps, NpyAuxData *NPY_UNUSED(func)) { char *ip1 = args[0]; - char *indx = args[1]; + char *indxp = args[1]; char *value = args[2]; npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2]; + npy_intp shape = steps[3]; npy_intp n = dimensions[0]; npy_intp i; @ftype@ *indexed; - for(i = 0; i < n; i++, indx += isindex, value += isb) { - indexed = (@ftype@ *)(ip1 + is1 * *(npy_intp *)indx); + for(i = 0; i < n; i++, indxp += isindex, value += isb) { + npy_intp indx = *(npy_intp *)indxp; + if (indx < 0) { + indx += shape; + } + indexed = (@ftype@ *)(ip1 + is1 * indx); const @ftype@ b_r = ((@ftype@ *)value)[0]; const @ftype@ b_i = ((@ftype@ *)value)[1]; #if @is_mul@ diff --git a/numpy/core/src/umath/loops_arithmetic.dispatch.c.src b/numpy/core/src/umath/loops_arithmetic.dispatch.c.src index b6f12629807b..e07bb79808af 100644 --- a/numpy/core/src/umath/loops_arithmetic.dispatch.c.src +++ b/numpy/core/src/umath/loops_arithmetic.dispatch.c.src @@ -400,14 +400,19 @@ NPY_NO_EXPORT int NPY_CPU_DISPATCH_CURFX(@TYPE@_divide_indexed) (PyArrayMethod_Context *NPY_UNUSED(context), char * const*args, npy_intp const *dimensions, npy_intp const *steps, NpyAuxData *NPY_UNUSED(func)) { char *ip1 = args[0]; - char *indx = args[1]; + char *indxp = args[1]; char *value = args[2]; npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2]; + npy_intp shape = steps[3]; npy_intp n = dimensions[0]; npy_intp i; @type@ *indexed; - for(i = 0; i < n; i++, indx += isindex, value += isb) { - indexed = (@type@ *)(ip1 + is1 * *(npy_intp *)indx); + for(i = 0; i < n; i++, indxp += isindex, value += isb) { + npy_intp indx = *(npy_intp *)indxp; + if (indx < 0) { + indx += shape; + } + indexed = (@type@ *)(ip1 + is1 * indx); *indexed = floor_div_@TYPE@(*indexed, *(@type@ *)value); } return 0; @@ -486,14 +491,19 @@ NPY_NO_EXPORT int NPY_CPU_DISPATCH_CURFX(@TYPE@_divide_indexed) (PyArrayMethod_Context *NPY_UNUSED(context), char * const*args, npy_intp const *dimensions, npy_intp const *steps, NpyAuxData *NPY_UNUSED(func)) { char *ip1 = args[0]; - char *indx = args[1]; + char *indxp = args[1]; char *value = args[2]; npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2]; + npy_intp shape = steps[3]; npy_intp n = dimensions[0]; npy_intp i; @type@ *indexed; - for(i = 0; i < n; i++, indx += isindex, value += isb) { - indexed = (@type@ *)(ip1 + is1 * *(npy_intp *)indx); + for(i = 0; i < n; i++, indxp += isindex, value += isb) { + npy_intp indx = *(npy_intp *)indxp; + if (indx < 0) { + indx += shape; + } + indexed = (@type@ *)(ip1 + is1 * indx); @type@ in2 = *(@type@ *)value; if (NPY_UNLIKELY(in2 == 0)) { npy_set_floatstatus_divbyzero(); diff --git a/numpy/core/src/umath/loops_minmax.dispatch.c.src b/numpy/core/src/umath/loops_minmax.dispatch.c.src index 9d8667d3830a..236e2e2eb760 100644 --- a/numpy/core/src/umath/loops_minmax.dispatch.c.src +++ b/numpy/core/src/umath/loops_minmax.dispatch.c.src @@ -456,14 +456,19 @@ NPY_NO_EXPORT int NPY_CPU_DISPATCH_CURFX(@TYPE@_@kind@_indexed) (PyArrayMethod_Context *NPY_UNUSED(context), char *const *args, npy_intp const *dimensions, npy_intp const *steps, NpyAuxData *NPY_UNUSED(func)) { char *ip1 = args[0]; - char *indx = args[1]; + char *indxp = args[1]; char *value = args[2]; npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2]; npy_intp n = dimensions[0]; + npy_intp shape = steps[3]; npy_intp i; @type@ *indexed; - for(i = 0; i < n; i++, indx += isindex, value += isb) { - indexed = (@type@ *)(ip1 + is1 * *(npy_intp *)indx); + for(i = 0; i < n; i++, indxp += isindex, value += isb) { + npy_intp indx = *(npy_intp *)indxp; + if (indx < 0) { + indx += shape; + } + indexed = (@type@ *)(ip1 + is1 * indx); *indexed = SCALAR_OP(*indexed, *(@type@ *)value); } return 0; diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py index 03eeff285600..02c437021fe9 100644 --- a/numpy/core/tests/test_ufunc.py +++ b/numpy/core/tests/test_ufunc.py @@ -2215,13 +2215,21 @@ 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 + @pytest.mark.parametrize("dtype", + np.typecodes['AllInteger'] + np.typecodes['Float']) + @pytest.mark.parametrize("ufunc", + [np.add, np.subtract, np.divide, np.minimum, np.maximum]) + def test_at_negative_indexes(self, dtype, ufunc): + a = np.arange(0, 10).astype(dtype) + indxs = np.array([-1, 1, -1, 2]).astype(np.intp) + vals = np.array([1, 5, 2, 10], dtype=a.dtype) + + expected = a.copy() + for i, v in zip(indxs, vals): + expected[i] = ufunc(expected[i], v) + + ufunc.at(a, indxs, vals) + assert_array_equal(a, expected) assert np.all(indxs == [-1, 1, -1, 2]) def test_at_not_none_signature(self):