diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/core/src/umath/ufunc_object.c | 21 | ||||
| -rw-r--r-- | numpy/core/tests/test_ufunc.py | 23 |
2 files changed, 41 insertions, 3 deletions
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c index 186f18a62..bb951b0b6 100644 --- a/numpy/core/src/umath/ufunc_object.c +++ b/numpy/core/src/umath/ufunc_object.c @@ -2717,11 +2717,11 @@ reducelike_promote_and_resolve(PyUFuncObject *ufunc, char *method) { /* - * Note that the `ops` is not realy correct. But legacy resolution + * Note that the `ops` is not really correct. But legacy resolution * cannot quite handle the correct ops (e.g. a NULL first item if `out` - * is NULL), and it should only matter in very strange cases. + * is NULL) so we pass `arr` instead in that case. */ - PyArrayObject *ops[3] = {arr, arr, NULL}; + PyArrayObject *ops[3] = {out ? out : arr, arr, out}; /* * TODO: If `out` is not provided, arguably `initial` could define * the first DType (and maybe also the out one), that way @@ -2738,6 +2738,21 @@ reducelike_promote_and_resolve(PyUFuncObject *ufunc, Py_INCREF(operation_DTypes[0]); operation_DTypes[2] = operation_DTypes[0]; Py_INCREF(operation_DTypes[2]); + /* + * We have to force the dtype, because otherwise value based casting + * may ignore the request entirely. This means that `out=` the same + * as `dtype=`, which should probably not be forced normally, so we + * only do it for legacy DTypes... + */ + if (NPY_DT_is_legacy(operation_DTypes[0]) + && NPY_DT_is_legacy(operation_DTypes[1])) { + if (signature[0] == NULL) { + Py_INCREF(operation_DTypes[0]); + signature[0] = operation_DTypes[0]; + Py_INCREF(operation_DTypes[0]); + signature[2] = operation_DTypes[0]; + } + } } PyArrayMethodObject *ufuncimpl = promote_and_get_ufuncimpl(ufunc, diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py index ef0bac957..398bf88db 100644 --- a/numpy/core/tests/test_ufunc.py +++ b/numpy/core/tests/test_ufunc.py @@ -2134,6 +2134,29 @@ class TestUfunc: # It would be safe, but not equiv casting: ufunc(a, c, out=out, casting="equiv") + def test_reducelike_out_promotes(self): + # Check that the out argument to reductions is considered for + # promotion. See also gh-20455. + # Note that these paths could prefer `initial=` in the future and + # do not up-cast to the default integer for add and prod + arr = np.ones(1000, dtype=np.uint8) + out = np.zeros((), dtype=np.uint16) + assert np.add.reduce(arr, out=out) == 1000 + arr[:10] = 2 + assert np.multiply.reduce(arr, out=out) == 2**10 + + # For legacy dtypes, the signature currently has to be forced if `out=` + # is passed. The two paths below should differ, without `dtype=` the + # expected result should be: `np.prod(arr.astype("f8")).astype("f4")`! + arr = np.full(5, 2**25-1, dtype=np.int64) + + # float32 and int64 promote to float64: + res = np.zeros((), dtype=np.float32) + # If `dtype=` is passed, the calculation is forced to float32: + single_res = np.zeros((), dtype=np.float32) + np.multiply.reduce(arr, out=single_res, dtype=np.float32) + assert single_res != res + def test_reduce_noncontig_output(self): # Check that reduction deals with non-contiguous output arrays # appropriately. |