diff options
| author | Sebastian Berg <sebastianb@nvidia.com> | 2023-02-09 08:36:11 +0100 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2023-02-09 08:36:11 +0100 |
| commit | 3aa60665601efdcbc49c54841a599e69a2cbd1bc (patch) | |
| tree | 1d023e578c32b4e7cd40d2f45bb383330b56d497 /numpy | |
| parent | c4c0bbd36e597f24c481ffd736527cc48309e481 (diff) | |
| parent | 3b3d035eb6ec52353020ad2a76de1326795276e7 (diff) | |
| download | numpy-3aa60665601efdcbc49c54841a599e69a2cbd1bc.tar.gz | |
Merge pull request #23179 from mhvk/ufunc-at-fix-scalar-value
BUG/ENH: Fix fast index loops for 1-el array / allow scalar value
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/core/src/umath/ufunc_object.c | 9 | ||||
| -rw-r--r-- | numpy/core/tests/test_ufunc.py | 9 |
2 files changed, 16 insertions, 2 deletions
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c index bfb4a8143..a159003de 100644 --- a/numpy/core/src/umath/ufunc_object.c +++ b/numpy/core/src/umath/ufunc_object.c @@ -5915,7 +5915,12 @@ trivial_at_loop(PyArrayMethodObject *ufuncimpl, NPY_ARRAYMETHOD_FLAGS flags, steps[2] = 0; } else { args[2] = (char *)PyArray_DATA(op2_array); - steps[2] = PyArray_STRIDE(op2_array, 0); + if (PyArray_NDIM(op2_array) == 0 + || PyArray_DIM(op2_array, 0) <= 1) { + steps[2] = 0; + } else { + steps[2] = PyArray_STRIDE(op2_array, 0); + } } npy_intp *inner_size = NpyIter_GetInnerLoopSizePtr(iter->outer); @@ -6435,7 +6440,7 @@ ufunc_at(PyUFuncObject *ufunc, PyObject *args) */ if ((ufuncimpl->contiguous_indexed_loop != NULL) && (PyArray_NDIM(op1_array) == 1) && - (op2_array == NULL || PyArray_NDIM(op2_array) == 1) && + (op2_array == NULL || PyArray_NDIM(op2_array) <= 1) && (iter->subspace_iter == NULL) && (iter->numiter == 1)) { res = trivial_at_loop(ufuncimpl, flags, iter, op1_array, op2_array, &context); diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py index 58f9ef0b5..b0c435f33 100644 --- a/numpy/core/tests/test_ufunc.py +++ b/numpy/core/tests/test_ufunc.py @@ -2045,6 +2045,15 @@ class TestUfunc: np.add.at(arr, index, values) assert arr[0] == len(values) + @pytest.mark.parametrize("value", [ + np.ones(1), np.ones(()), np.float64(1.), 1.]) + def test_ufunc_at_scalar_value_fastpath(self, value): + arr = np.zeros(1000) + # index must be cast, which may be buffered in chunks: + index = np.repeat(np.arange(1000), 2) + np.add.at(arr, index, value) + assert_array_equal(arr, np.full_like(arr, 2 * value)) + 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]]) |