diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/core/src/umath/ufunc_object.c | 11 | ||||
| -rw-r--r-- | numpy/core/tests/test_ufunc.py | 11 |
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 79af0837b..62e06f281 100644 --- a/numpy/core/src/umath/ufunc_object.c +++ b/numpy/core/src/umath/ufunc_object.c @@ -6323,7 +6323,6 @@ ufunc_at(PyUFuncObject *ufunc, PyObject *args) if ((iter2 = (PyArrayIterObject *)\ PyArray_BroadcastToShape((PyObject *)op2_array, iter->dimensions, iter->nd))==NULL) { - /* only on memory allocation failures */ goto fail; } } @@ -6417,17 +6416,21 @@ ufunc_at(PyUFuncObject *ufunc, PyObject *args) int fast_path = 1; /* check no casting, alignment */ if (PyArray_DESCR(op1_array) != operation_descrs[0]) { - fast_path = 0; + fast_path = 0; + } + if (PyArray_DESCR(op1_array) != operation_descrs[nop - 1]) { + /* output casting */ + fast_path = 0; } if (!PyArray_ISALIGNED(op1_array)) { - fast_path = 0; + fast_path = 0; } if (nop >2) { if (PyArray_DESCR(op2_array) != operation_descrs[1]) { fast_path = 0; } if (!PyArray_ISALIGNED(op2_array)) { - fast_path = 0; + fast_path = 0; } } if (fast_path == 1) { diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py index ff20ad35c..85c90c129 100644 --- a/numpy/core/tests/test_ufunc.py +++ b/numpy/core/tests/test_ufunc.py @@ -2102,6 +2102,17 @@ class TestUfunc: with pytest.raises(np.core._exceptions._UFuncNoLoopError): np.add.at(arr, [0, 1], [0, 1]) + def test_at_output_casting(self): + arr = np.array([-1]) + np.equal.at(arr, [0], [0]) + assert arr[0] == 0 + + def test_at_broadcast_failure(self): + arr = np.arange(5) + with pytest.raises(ValueError): + np.add.at(arr, [0, 1], [1, 2, 3]) + + def test_reduce_arguments(self): f = np.add.reduce d = np.ones((5,2), dtype=int) |