diff options
author | Sebastian Berg <sebastian@sipsolutions.net> | 2020-06-02 18:39:42 -0500 |
---|---|---|
committer | Sebastian Berg <sebastian@sipsolutions.net> | 2020-06-02 18:39:42 -0500 |
commit | 72a4e301d7b8b8da16a035f9a720308e9a4f2d60 (patch) | |
tree | e1c91929318300c928ccb2e40992e28ee57d1a2e | |
parent | cbde47818a821ef951ca73e851579ca75d23a82b (diff) | |
download | numpy-72a4e301d7b8b8da16a035f9a720308e9a4f2d60.tar.gz |
BUG: Fix result when a gufunc output broadcasts the inputs.
In this case no error was given, but additional dimensions in the
output simply ignored. Thus the returned array was only partially
set.
Now it will be fully set, with the calculation being repeated as
often as necessary (typical broadcasting logic). This is consistent
with how normal ufuncs work.
Closes gh-16484
-rw-r--r-- | numpy/core/src/umath/ufunc_object.c | 10 | ||||
-rw-r--r-- | numpy/core/tests/test_ufunc.py | 15 |
2 files changed, 22 insertions, 3 deletions
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c index c57199c79..591fbd7ad 100644 --- a/numpy/core/src/umath/ufunc_object.c +++ b/numpy/core/src/umath/ufunc_object.c @@ -2610,11 +2610,15 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc, } /* * Figure out the number of iteration dimensions, which - * is the broadcast result of all the input non-core - * dimensions. + * is the broadcast result of all the non-core dimensions. + * (We do allow outputs to broadcast inputs currently, if they are given. + * This is in line with what normal ufuncs do.) */ broadcast_ndim = 0; - for (i = 0; i < nin; ++i) { + for (i = 0; i < nop; ++i) { + if (op[i] == NULL) { + continue; + } int n = PyArray_NDIM(op[i]) - op_core_num_dims[i]; if (n > broadcast_ndim) { broadcast_ndim = n; diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py index abdaeeb93..642308d3b 100644 --- a/numpy/core/tests/test_ufunc.py +++ b/numpy/core/tests/test_ufunc.py @@ -612,6 +612,21 @@ class TestUfunc: assert_equal(len(w), 1) assert_(x[0,0] != u[0, 0]) + def test_out_broadcasts(self): + # For ufuncs and gufuncs (not for reductions), we currently allow + # the output to cause broadcasting of the input arrays. + # both along dimensions with shape 1 and dimensions which do not + # exist at all in the inputs. + arr = np.arange(3).reshape(1, 3) + out = np.empty((5, 4, 3)) + np.add(arr, arr, out=out) + assert (out == np.arange(3) * 2).all() + + # The same holds for gufuncs (gh-16484) + umt.inner1d(arr, arr, out=out) + # the result would be just a scalar `5`, but is broadcast fully: + assert (out == 5).all() + def test_type_cast(self): msg = "type cast" a = np.arange(6, dtype='short').reshape((2, 3)) |