Merge pull request numpy#24281 from seberg/fixup-indexed

BUG: Further fixes to indexing loop and added tests

Author: charris

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 {

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@

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();

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;

numpy/core/tests/test_ufunc.py

@@ -2224,13 +2224,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):