diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/core/src/umath/reduction.c | 23 | ||||
| -rw-r--r-- | numpy/core/tests/test_ufunc.py | 17 |
2 files changed, 38 insertions, 2 deletions
diff --git a/numpy/core/src/umath/reduction.c b/numpy/core/src/umath/reduction.c index 4ce8d8ab7..79c302755 100644 --- a/numpy/core/src/umath/reduction.c +++ b/numpy/core/src/umath/reduction.c @@ -84,10 +84,12 @@ allocate_reduce_result(PyArrayObject *arr, const npy_bool *axis_flags, * The return value is a view into 'out'. */ static PyArrayObject * -conform_reduce_result(int ndim, const npy_bool *axis_flags, +conform_reduce_result(PyArrayObject *in, const npy_bool *axis_flags, PyArrayObject *out, int keepdims, const char *funcname, int need_copy) { + int ndim = PyArray_NDIM(in); + npy_intp *shape_in = PyArray_DIMS(in); npy_intp strides[NPY_MAXDIMS], shape[NPY_MAXDIMS]; npy_intp *strides_out = PyArray_STRIDES(out); npy_intp *shape_out = PyArray_DIMS(out); @@ -118,6 +120,16 @@ conform_reduce_result(int ndim, const npy_bool *axis_flags, return NULL; } } + else { + if (shape_out[idim] != shape_in[idim]) { + PyErr_Format(PyExc_ValueError, + "output parameter for reduction operation %s " + "has a non-reduction dimension not equal to " + "the input one.", funcname); + return NULL; + } + } + } Py_INCREF(out); @@ -138,6 +150,13 @@ conform_reduce_result(int ndim, const npy_bool *axis_flags, "does not have enough dimensions", funcname); return NULL; } + if (shape_out[idim_out] != shape_in[idim]) { + PyErr_Format(PyExc_ValueError, + "output parameter for reduction operation %s " + "has a non-reduction dimension not equal to " + "the input one.", funcname); + return NULL; + } strides[idim] = strides_out[idim_out]; shape[idim] = shape_out[idim_out]; ++idim_out; @@ -240,7 +259,7 @@ PyArray_CreateReduceResult(PyArrayObject *operand, PyArrayObject *out, /* Steal the dtype reference */ Py_XDECREF(dtype); - result = conform_reduce_result(PyArray_NDIM(operand), axis_flags, + result = conform_reduce_result(operand, axis_flags, out, keepdims, funcname, need_copy); } diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py index 679bea96a..01e3c5087 100644 --- a/numpy/core/tests/test_ufunc.py +++ b/numpy/core/tests/test_ufunc.py @@ -1886,6 +1886,23 @@ class TestUfunc: assert_equal(y_base[1,:], y_base_copy[1,:]) assert_equal(y_base[3,:], y_base_copy[3,:]) + @pytest.mark.parametrize('output_shape', + [(), (1,), (1, 1), (1, 3), (4, 3)]) + @pytest.mark.parametrize('f_reduce', [np.add.reduce, np.minimum.reduce]) + def test_reduce_wrong_dimension_output(self, f_reduce, output_shape): + # Test that we're not incorrectly broadcasting dimensions. + # See gh-15144 (failed for np.add.reduce previously). + a = np.arange(12.).reshape(4, 3) + out = np.empty(output_shape, a.dtype) + assert_raises(ValueError, f_reduce, a, axis=0, out=out) + if output_shape != (1, 3): + assert_raises(ValueError, f_reduce, a, axis=0, out=out, + keepdims=True) + else: + check = f_reduce(a, axis=0, out=out, keepdims=True) + assert_(check is out) + assert_array_equal(check, f_reduce(a, axis=0, keepdims=True)) + def test_no_doc_string(self): # gh-9337 assert_('\n' not in umt.inner1d_no_doc.__doc__) |