BUG: Deal with casting and non-contig/1-D indices

2 files changed, 18 insertions, 4 deletions

diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 551315fa4..efcfa73fe 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -5907,15 +5907,21 @@ trivial_at_loop(PyArrayMethodObject *ufuncimpl, NPY_ARRAYMETHOD_FLAGS flags,
     int buffersize=0, errormask = 0;
     int res;
     char *args[3];
-    npy_intp dimensions = iter->size;
     npy_intp steps[3];
     args[0] = (char *) iter->baseoffset;
-    args[1] = (char *) iter->outer_ptrs[0];
     args[2] = (char *)PyArray_DATA(op2_array);
     steps[0] = iter->fancy_strides[0];
-    steps[1] = iter->outer_strides[0];
     steps[2] = PyArray_STRIDE(op2_array, 0);
-    res = ufuncimpl->contiguous_indexed_loop(context, args, &dimensions, steps, NULL);
+
+    npy_intp *inner_size = NpyIter_GetInnerLoopSizePtr(iter->outer);
+
+    do {
+        args[1] = (char *) iter->outer_ptrs[0];
+        steps[1] = iter->outer_strides[0];
+
+        res = ufuncimpl->contiguous_indexed_loop(
+                context, args, inner_size, steps, NULL);
+    } while (res == 0 && iter->outer_next(iter->outer));
     /*
      * An error should only be possible if `res != 0` is already set.
      * But this is not strictly correct since old-style ufuncs (e.g. `power`
diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
index c8d13918f..f31482666 100644
--- a/numpy/core/tests/test_ufunc.py
+++ b/numpy/core/tests/test_ufunc.py
@@ -2023,6 +2023,14 @@ class TestUfunc:
             assert w_at[0].category == w_loop[0].category
             assert str(w_at[0].message)[:10] == str(w_loop[0].message)[:10]
 
+    def test_cast_index_fastpath(self):
+        arr = np.zeros(10)
+        values = np.ones(100000)
+        # index must be cast, which may be buffered in chunks:
+        index = np.zeros(len(values), dtype=np.uint8)
+        np.add.at(arr, index, values)
+        assert arr[0] == len(values)
+
     def test_ufunc_at_multiD(self):
         a = np.arange(9).reshape(3, 3)
         b = np.array([[100, 100, 100], [200, 200, 200], [300, 300, 300]])