summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorSebastian Berg <sebastian@sipsolutions.net>2020-06-02 18:39:42 -0500
committerSebastian Berg <sebastian@sipsolutions.net>2020-06-02 18:39:42 -0500
commit72a4e301d7b8b8da16a035f9a720308e9a4f2d60 (patch)
treee1c91929318300c928ccb2e40992e28ee57d1a2e
parentcbde47818a821ef951ca73e851579ca75d23a82b (diff)
downloadnumpy-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.c10
-rw-r--r--numpy/core/tests/test_ufunc.py15
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))