8000 BUG: Deal with casting and non-contig/1-D indices · mattip/numpy@6cf9cc9 · GitHub
[go: up one dir, main page]

Skip to content

Commit 6cf9cc9

Browse files
committed
BUG: Deal with casting and non-contig/1-D indices
1 parent f064552 commit 6cf9cc9

File tree

2 files changed

+18
-4
lines changed

2 files changed

+18
-4
lines changed

numpy/core/src/umath/ufunc_object.c

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5907,15 +5907,21 @@ trivial_at_loop(PyArrayMethodObject *ufuncimpl, NPY_ARRAYMETHOD_FLAGS flags,
59075907
int buffersize=0, errormask = 0;
59085908
int res;
59095909
char *args[3];
5910-
npy_intp dimensions = iter->size;
59115910
npy_intp steps[3];
59125911
args[0] = (char *) iter->baseoffset;
5913-
args[1] = (char *) iter->outer_ptrs[0];
59145912
args[2] = (char *)PyArray_DATA(op2_array);
59155913
steps[0] = iter->fancy_strides[0];
5916-
steps[1] = iter->outer_strides[0];
59175914
steps[2] = PyArray_STRIDE(op2_array, 0);
5918-
res = ufuncimpl->contiguous_indexed_loop(context, args, &dimensions, steps, NULL);
5915+
5916+
npy_intp *inner_size = NpyIter_GetInnerLoopSizePtr(iter->outer);
5917+
5918+
do {
5919+
args[1] = (char *) iter->outer_ptrs[0];
5920+
steps[1] = iter->outer_strides[0];
5921+
5922+
res = ufuncimpl->contiguous_indexed_loop(
5923+
context, args, inner_size, steps, NULL);
5924+
} while (res == 0 && iter->outer_next(iter->outer));
59195925
/*
59205926
* An error should only be possible if `res != 0` is already set.
59215927
* But this is not strictly correct since old-style ufuncs (e.g. `power`

numpy/core/tests/test_ufunc.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2023,6 +2023,14 @@ def test_ufunc_at_inner_loops(self, typecode, ufunc):
20232023
assert w_at[0].category == w_loop[0].category
20242024
assert str(w_at[0].message)[:10] == str(w_loop[0].message)[:10]
20252025

2026+
def test_cast_index_fastpath(self):
2027+
arr = np.zeros(10)
2028+
values = np.ones(100000)
2029+
# index must be cast, which may be buffered in chunks:
2030+
index = np.zeros(len(values), dtype=np.uint8)
2031+
np.add.at(arr, index, values)
2032+
assert arr[0] == len(values)
2033+
20262034
def test_ufunc_at_multiD(self):
20272035
a = np.arange(9).reshape(3, 3)
20282036
b = np.array([[100, 100, 100], [200, 200, 200], [300, 300, 300]])

0 commit comments

Comments
 (0)
0