Merge pull request #9106 from mhvk/array_ufunc_reduce_out_tuple

BUG: Array ufunc reduce out tuple

10 files changed, 197 insertions, 70 deletions

diff --git a/doc/neps/ufunc-overrides.rst b/doc/neps/ufunc-overrides.rst
index fadc20440..90869e1ac 100644
--- a/doc/neps/ufunc-overrides.rst
+++ b/doc/neps/ufunc-overrides.rst
@@ -166,7 +166,8 @@ Hence, the arguments are normalized: only the required input arguments
 passed on as a dict of keyword arguments (``kwargs``). In particular, if
 there are output arguments, positional are otherwise, that are not
 :obj:`None`, they are passed on as a tuple in the ``out`` keyword
-argument.
+argument (even for the ``reduce``, ``accumulate``, and ``reduceat`` methods
+where in all current cases only a single output makes sense).
 
 The function dispatch proceeds as follows:
 
diff --git a/doc/release/1.13.0-notes.rst b/doc/release/1.13.0-notes.rst
index 84d1106d6..faa0f0d27 100644
--- a/doc/release/1.13.0-notes.rst
+++ b/doc/release/1.13.0-notes.rst
@@ -544,3 +544,11 @@ The ABCPolyBase class, from which the convenience classes are derived, sets
 ``__array_ufun__ = None`` in order of opt out of ufuncs. If a polynomial
 convenience class instance is passed as an argument to a ufunc, a ``TypeError``
 will now be raised.
+
+Output arguments to ufuncs can be tuples also for ufunc methods
+---------------------------------------------------------------
+For calls to ufuncs, it was already possible, and recommended, to use an
+``out`` argument with a tuple for ufuncs with multiple outputs. This has now
+been extended to output arguments in the ``reduce``, ``accumulate``, and
+``reduceat`` methods. This is mostly for compatibility with ``__array_ufunc``;
+there are no ufuncs yet that have more than one output.
diff --git a/doc/source/reference/ufuncs.rst b/doc/source/reference/ufuncs.rst
index b3fb4d384..9a8d8e20e 100644
--- a/doc/source/reference/ufuncs.rst
+++ b/doc/source/reference/ufuncs.rst
@@ -426,8 +426,8 @@ Methods
 All ufuncs have four methods. However, these methods only make sense on
 ufuncs that take two input arguments and return one output argument.
 Attempting to call these methods on other ufuncs will cause a
-:exc:`ValueError`. The reduce-like methods all take an *axis* keyword
-and a *dtype* keyword, and the arrays must all have dimension >= 1.
+:exc:`ValueError`. The reduce-like methods all take an *axis* keyword, a *dtype*
+keyword, and an *out* keyword, and the arrays must all have dimension >= 1.
 The *axis* keyword specifies the axis of the array over which the reduction
 will take place and may be negative, but must be an integer. The
 *dtype* keyword allows you to manage a very common problem that arises
@@ -443,7 +443,10 @@ mostly up to you. There is one exception: if no *dtype* is given for a
 reduction on the "add" or "multiply" operations, then if the input type is
 an integer (or Boolean) data-type and smaller than the size of the
 :class:`int_` data type, it will be internally upcast to the :class:`int_`
-(or :class:`uint`) data-type.
+(or :class:`uint`) data-type. Finally, the *out* keyword allows you to provide
+an output array (for single-output ufuncs, which are currently the only ones
+supported; for future extension, however, a tuple with a single argument
+can be passed in). If *out* is given, the *dtype* argument is ignored.
 
 Ufuncs also have a fifth method that allows in place operations to be
 performed using fancy indexing. No buffering is used on the dimensions where
diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
index 8a016c966..05fb73df4 100644
--- a/numpy/add_newdocs.py
+++ b/numpy/add_newdocs.py
@@ -5444,9 +5444,11 @@ add_newdoc('numpy.core', 'ufunc',
     ----------
     *x : array_like
         Input arrays.
-    out : ndarray or tuple of ndarray, optional
+    out : ndarray, None, or tuple of ndarray and None, optional
         Alternate array object(s) in which to put the result; if provided, it
-        must have a shape that the inputs broadcast to.
+        must have a shape that the inputs broadcast to. A tuple of arrays
+        (possible only as a keyword argument) must have length equal to the
+        number of outputs; use `None` for outputs to be allocated by the ufunc.
     where : array_like, optional
         Values of True indicate to calculate the ufunc at that position, values
         of False indicate to leave the value in the output alone.
@@ -5667,9 +5669,14 @@ add_newdoc('numpy.core', 'ufunc', ('reduce',
         The type used to represent the intermediate results. Defaults
         to the data-type of the output array if this is provided, or
         the data-type of the input array if no output array is provided.
-    out : ndarray, optional
-        A location into which the result is stored. If not provided, a
-        freshly-allocated array is returned.
+    out : ndarray, None, or tuple of ndarray and None, optional
+        A location into which the result is stored. If not provided or `None`,
+        a freshly-allocated array is returned. For consistency with
+        :ref:`ufunc.__call__`, if given as a keyword, this may be wrapped in a
+        1-element tuple.
+
+        .. versionchanged:: 1.13.0
+           Tuples are allowed for keyword argument.
     keepdims : bool, optional
         If this is set to True, the axes which are reduced are left
         in the result as dimensions with size one. With this option,
@@ -5741,9 +5748,14 @@ add_newdoc('numpy.core', 'ufunc', ('accumulate',
         The data-type used to represent the intermediate results. Defaults
         to the data-type of the output array if such is provided, or the
         the data-type of the input array if no output array is provided.
-    out : ndarray, optional
-        A location into which the result is stored. If not provided a
-        freshly-allocated array is returned.
+    out : ndarray, None, or tuple of ndarray and None, optional
+        A location into which the result is stored. If not provided or `None`,
+        a freshly-allocated array is returned. For consistency with
+        :ref:`ufunc.__call__`, if given as a keyword, this may be wrapped in a
+        1-element tuple.
+
+        .. versionchanged:: 1.13.0
+           Tuples are allowed for keyword argument.
     keepdims : bool
         Has no effect. Deprecated, and will be removed in future.
 
@@ -5820,9 +5832,14 @@ add_newdoc('numpy.core', 'ufunc', ('reduceat',
         The type used to represent the intermediate results. Defaults
         to the data type of the output array if this is provided, or
         the data type of the input array if no output array is provided.
-    out : ndarray, optional
-        A location into which the result is stored. If not provided a
-        freshly-allocated array is returned.
+    out : ndarray, None, or tuple of ndarray and None, optional
+        A location into which the result is stored. If not provided or `None`,
+        a freshly-allocated array is returned. For consistency with
+        :ref:`ufunc.__call__`, if given as a keyword, this may be wrapped in a
+        1-element tuple.
+
+        .. versionchanged:: 1.13.0
+           Tuples are allowed for keyword argument.
 
     Returns
     -------
diff --git a/numpy/core/code_generators/ufunc_docstrings.py b/numpy/core/code_generators/ufunc_docstrings.py
index 7beda59f2..c5172f6a8 100644
--- a/numpy/core/code_generators/ufunc_docstrings.py
+++ b/numpy/core/code_generators/ufunc_docstrings.py
@@ -19,9 +19,11 @@ def get(name):
 
 # common parameter text to all ufuncs
 _params_text = textwrap.dedent("""
-    out : ndarray or tuple of ndarray, optional
-        Alternate array object(s) in which to put the result; if provided, it
-        must have a shape that the inputs broadcast to.
+    out : ndarray, None, or tuple of ndarray and None, optional
+        A location into which the result is stored. If provided, it must have
+        a shape that the inputs broadcast to. If not provided or `None`,
+        a freshly-allocated array is returned. A tuple (possible only as a
+        keyword argument) must have length equal to the number of outputs.
     where : array_like, optional
         Values of True indicate to calculate the ufunc at that position, values
         of False indicate to leave the value in the output alone.
diff --git a/numpy/core/src/umath/override.c b/numpy/core/src/umath/override.c
index 1157cae66..bdb992c28 100644
--- a/numpy/core/src/umath/override.c
+++ b/numpy/core/src/umath/override.c
@@ -144,14 +144,16 @@ normalize_reduce_args(PyUFuncObject *ufunc, PyObject *args,
             return -1;
         }
         obj = PyTuple_GET_ITEM(args, i);
-        if (obj != Py_None) {
-            if (i == 3) {
-                obj = PyTuple_GetSlice(args, 3, 4);
-            }
-            PyDict_SetItemString(*normal_kwds, kwlist[i], obj);
-            if (i == 3) {
-                Py_DECREF(obj);
+        if (i == 3) {
+            /* remove out=None */
+            if (obj == Py_None) {
+                continue;
             }
+            obj = PyTuple_GetSlice(args, 3, 4);
+        }
+        PyDict_SetItemString(*normal_kwds, kwlist[i], obj);
+        if (i == 3) {
+            Py_DECREF(obj);
         }
     }
     return 0;
@@ -188,14 +190,16 @@ normalize_accumulate_args(PyUFuncObject *ufunc, PyObject *args,
             return -1;
         }
         obj = PyTuple_GET_ITEM(args, i);
-        if (obj != Py_None) {
-            if (i == 3) {
-                obj = PyTuple_GetSlice(args, 3, 4);
-            }
-            PyDict_SetItemString(*normal_kwds, kwlist[i], obj);
-            if (i == 3) {
-                Py_DECREF(obj);
+        if (i == 3) {
+            /* remove out=None */
+            if (obj == Py_None) {
+                continue;
             }
+            obj = PyTuple_GetSlice(args, 3, 4);
+        }
+        PyDict_SetItemString(*normal_kwds, kwlist[i], obj);
+        if (i == 3) {
+            Py_DECREF(obj);
         }
     }
     return 0;
@@ -234,14 +238,16 @@ normalize_reduceat_args(PyUFuncObject *ufunc, PyObject *args,
             return -1;
         }
         obj = PyTuple_GET_ITEM(args, i);
-        if (obj != Py_None) {
-            if (i == 4) {
-                obj = PyTuple_GetSlice(args, 4, 5);
-            }
-            PyDict_SetItemString(*normal_kwds, kwlist[i], obj);
-            if (i == 4) {
-                Py_DECREF(obj);
+        if (i == 4) {
+            /* remove out=None */
+            if (obj == Py_None) {
+                continue;
             }
+            obj = PyTuple_GetSlice(args, 4, 5);
+        }
+        PyDict_SetItemString(*normal_kwds, kwlist[i], obj);
+        if (i == 4) {
+            Py_DECREF(obj);
         }
     }
     return 0;
@@ -360,11 +366,11 @@ PyUFunc_CheckOverride(PyUFuncObject *ufunc, char *method,
         if (out != NULL) {
             int nout = ufunc->nout;
 
-            if (PyTuple_Check(out)) {
+            if (PyTuple_CheckExact(out)) {
                 int all_none = 1;
 
                 if (PyTuple_GET_SIZE(out) != nout) {
-                    PyErr_Format(PyExc_TypeError,
+                    PyErr_Format(PyExc_ValueError,
                                  "The 'out' tuple must have exactly "
                                  "%d entries: one per ufunc output", nout);
                     goto fail;
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index e1219039c..ea99560ea 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -1084,7 +1084,7 @@ get_ufunc_arguments(PyUFuncObject *ufunc,
                                     "positional and keyword argument");
                             goto fail;
                         }
-                        if (PyTuple_Check(value)) {
+                        if (PyTuple_CheckExact(value)) {
                             if (PyTuple_GET_SIZE(value) != nout) {
                                 PyErr_SetString(PyExc_ValueError,
                                         "The 'out' tuple must have exactly "
@@ -3894,6 +3894,7 @@ PyUFunc_GenericReduction(PyUFuncObject *ufunc, PyObject *args,
     PyObject *obj_ind, *context;
     PyArrayObject *indices = NULL;
     PyArray_Descr *otype = NULL;
+    PyObject *out_obj = NULL;
     PyArrayObject *out = NULL;
     int keepdims = 0;
     static char *reduce_kwlist[] = {
@@ -3927,7 +3928,20 @@ PyUFunc_GenericReduction(PyUFuncObject *ufunc, PyObject *args,
                      _reduce_type[operation]);
         return NULL;
     }
-
+    /* if there is a tuple of 1 for `out` in kwds, unpack it */
+    if (kwds != NULL) {
+        PyObject *out_obj = PyDict_GetItem(kwds, npy_um_str_out);
+        if (out_obj != NULL && PyTuple_CheckExact(out_obj)) {
+            if (PyTuple_GET_SIZE(out_obj) != 1) {
+                PyErr_SetString(PyExc_ValueError,
+                                "The 'out' tuple must have exactly one entry");
+                return NULL;
+            }
+            out_obj = PyTuple_GET_ITEM(out_obj, 0);
+            PyDict_SetItem(kwds, npy_um_str_out, out_obj);
+        }
+    }
+            
     if (operation == UFUNC_REDUCEAT) {
         PyArray_Descr *indtype;
         indtype = PyArray_DescrFromType(NPY_INTP);
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index 30bda20de..ca2fb35fb 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -3040,7 +3040,7 @@ class TestBinop(object):
             warnings.filterwarnings('always', '', DeprecationWarning)
             assert_equal(np.modf(dummy, out=a), (0,))
             assert_(w[0].category is DeprecationWarning)
-        assert_raises(TypeError, np.modf, dummy, out=(a,))
+        assert_raises(ValueError, np.modf, dummy, out=(a,))
 
         # 2 inputs, 1 output
         assert_equal(np.add(a, dummy), 0)
diff --git a/numpy/core/tests/test_umath.py b/numpy/core/tests/test_umath.py
index 9e849df92..3572e273d 100644
--- a/numpy/core/tests/test_umath.py
+++ b/numpy/core/tests/test_umath.py
@@ -1755,15 +1755,18 @@ class TestSpecialMethods(TestCase):
                               'keepdims': 'keep0',
                               'axis': 'axis0'})
 
-        # reduce, output equal to None removed.
-        res = np.multiply.reduce(a, out=None)
-        assert_equal(res[4], {})
-        res = np.multiply.reduce(a, out=(None,))
-        assert_equal(res[4], {})
+        # reduce, output equal to None removed, but not other explicit ones,
+        # even if they are at their default value.
+        res = np.multiply.reduce(a, 0, None, None, False)
+        assert_equal(res[4], {'axis': 0, 'dtype': None, 'keepdims': False})
+        res = np.multiply.reduce(a, out=None, axis=0, keepdims=True)
+        assert_equal(res[4], {'axis': 0, 'keepdims': True})
+        res = np.multiply.reduce(a, None, out=(None,), dtype=None)
+        assert_equal(res[4], {'axis': None, 'dtype': None})
 
         # reduce, wrong args
-        assert_raises(TypeError, np.multiply.reduce, a, out=())
-        assert_raises(TypeError, np.multiply.reduce, a, out=('out0', 'out1'))
+        assert_raises(ValueError, np.multiply.reduce, a, out=())
+        assert_raises(ValueError, np.multiply.reduce, a, out=('out0', 'out1'))
         assert_raises(TypeError, np.multiply.reduce, a, 'axis0', axis='axis0')
 
         # accumulate, pos args
@@ -1788,14 +1791,16 @@ class TestSpecialMethods(TestCase):
                               'axis': 'axis0'})
 
         # accumulate, output equal to None removed.
-        res = np.multiply.accumulate(a, out=None)
-        assert_equal(res[4], {})
-        res = np.multiply.accumulate(a, out=(None,))
-        assert_equal(res[4], {})
+        res = np.multiply.accumulate(a, 0, None, None)
+        assert_equal(res[4], {'axis': 0, 'dtype': None})
+        res = np.multiply.accumulate(a, out=None, axis=0, dtype='dtype1')
+        assert_equal(res[4], {'axis': 0, 'dtype': 'dtype1'})
+        res = np.multiply.accumulate(a, None, out=(None,), dtype=None)
+        assert_equal(res[4], {'axis': None, 'dtype': None})
 
         # accumulate, wrong args
-        assert_raises(TypeError, np.multiply.accumulate, a, out=())
-        assert_raises(TypeError, np.multiply.accumulate, a,
+        assert_raises(ValueError, np.multiply.accumulate, a, out=())
+        assert_raises(ValueError, np.multiply.accumulate, a,
                       out=('out0', 'out1'))
         assert_raises(TypeError, np.multiply.accumulate, a,
                       'axis0', axis='axis0')
@@ -1822,14 +1827,16 @@ class TestSpecialMethods(TestCase):
                               'axis': 'axis0'})
 
         # reduceat, output equal to None removed.
-        res = np.multiply.reduceat(a, [4, 2], out=None)
-        assert_equal(res[4], {})
-        res = np.multiply.reduceat(a, [4, 2], out=(None,))
-        assert_equal(res[4], {})
+        res = np.multiply.reduceat(a, [4, 2], 0, None, None)
+        assert_equal(res[4], {'axis': 0, 'dtype': None})
+        res = np.multiply.reduceat(a, [4, 2], axis=None, out=None, dtype='dt')
+        assert_equal(res[4], {'axis': None, 'dtype': 'dt'})
+        res = np.multiply.reduceat(a, [4, 2], None, None, out=(None,))
+        assert_equal(res[4], {'axis': None, 'dtype': None})
 
         # reduceat, wrong args
-        assert_raises(TypeError, np.multiply.reduce, a, [4, 2], out=())
-        assert_raises(TypeError, np.multiply.reduce, a, [4, 2],
+        assert_raises(ValueError, np.multiply.reduce, a, [4, 2], out=())
+        assert_raises(ValueError, np.multiply.reduce, a, [4, 2],
                       out=('out0', 'out1'))
         assert_raises(TypeError, np.multiply.reduce, a, [4, 2],
                       'axis0', axis='axis0')
@@ -1907,12 +1914,12 @@ class TestSpecialMethods(TestCase):
         # wrong number of arguments in the tuple is an error too.
         assert_raises(TypeError, np.multiply, a, b, 'one', out='two')
         assert_raises(TypeError, np.multiply, a, b, 'one', 'two')
-        assert_raises(TypeError, np.multiply, a, b, out=('one', 'two'))
-        assert_raises(TypeError, np.multiply, a, out=())
+        assert_raises(ValueError, np.multiply, a, b, out=('one', 'two'))
+        assert_raises(ValueError, np.multiply, a, out=())
         assert_raises(TypeError, np.modf, a, 'one', out=('two', 'three'))
         assert_raises(TypeError, np.modf, a, 'one', 'two', 'three')
-        assert_raises(TypeError, np.modf, a, out=('one', 'two', 'three'))
-        assert_raises(TypeError, np.modf, a, out=('one',))
+        assert_raises(ValueError, np.modf, a, out=('one', 'two', 'three'))
+        assert_raises(ValueError, np.modf, a, out=('one',))
 
     def test_ufunc_override_exception(self):
 
@@ -1999,11 +2006,12 @@ class TestSpecialMethods(TestCase):
         assert_raises(TypeError, inner1d, a, out='two')
         assert_raises(TypeError, inner1d, a, a, 'one', out='two')
         assert_raises(TypeError, inner1d, a, a, 'one', 'two')
-        assert_raises(TypeError, inner1d, a, a, out=('one', 'two'))
-        assert_raises(TypeError, inner1d, a, a, out=())
+        assert_raises(ValueError, inner1d, a, a, out=('one', 'two'))
+        assert_raises(ValueError, inner1d, a, a, out=())
 
     def test_ufunc_override_with_super(self):
-
+        # NOTE: this class is given as an example in doc/subclassing.py;
+        # if you make any changes here, do update it there too.
         class A(np.ndarray):
             def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
                 args = []
@@ -2041,6 +2049,8 @@ class TestSpecialMethods(TestCase):
                     return NotImplemented
 
                 if method == 'at':
+                    if isinstance(inputs[0], A):
+                        inputs[0].info = info
                     return
 
                 if ufunc.nout == 1:
@@ -2107,6 +2117,70 @@ class TestSpecialMethods(TestCase):
         assert_(a.__array_ufunc__(np.add, '__call__', a, b) is NotImplemented)
         assert_(b.__array_ufunc__(np.add, '__call__', a, b) == "A!")
         assert_(np.add(a, b) == "A!")
+        # regression check for gh-9102 -- tests ufunc.reduce implicitly.
+        d = np.array([[1, 2, 3], [1, 2, 3]])
+        a = d.view(A)
+        c = a.any()
+        check = d.any()
+        assert_equal(c, check)
+        assert_(c.info, {'inputs': [0]})
+        c = a.max()
+        check = d.max()
+        assert_equal(c, check)
+        assert_(c.info, {'inputs': [0]})
+        b = np.array(0).view(A)
+        c = a.max(out=b)
+        assert_equal(c, check)
+        assert_(c is b)
+        assert_(c.info, {'inputs': [0], 'outputs': [0]})
+        check = a.max(axis=0)
+        b = np.zeros_like(check).view(A)
+        c = a.max(axis=0, out=b)
+        assert_equal(c, check)
+        assert_(c is b)
+        assert_(c.info, {'inputs': [0], 'outputs': [0]})
+        # simple explicit tests of reduce, accumulate, reduceat
+        check = np.add.reduce(d, axis=1)
+        c = np.add.reduce(a, axis=1)
+        assert_equal(c, check)
+        assert_(c.info, {'inputs': [0]})
+        b = np.zeros_like(c)
+        c = np.add.reduce(a, 1, None, b)
+        assert_equal(c, check)
+        assert_(c is b)
+        assert_(c.info, {'inputs': [0], 'outputs': [0]})
+        check = np.add.accumulate(d, axis=0)
+        c = np.add.accumulate(a, axis=0)
+        assert_equal(c, check)
+        assert_(c.info, {'inputs': [0]})
+        b = np.zeros_like(c)
+        c = np.add.accumulate(a, 0, None, b)
+        assert_equal(c, check)
+        assert_(c is b)
+        assert_(c.info, {'inputs': [0], 'outputs': [0]})
+        indices = [0, 2, 1]
+        check = np.add.reduceat(d, indices, axis=1)
+        c = np.add.reduceat(a, indices, axis=1)
+        assert_equal(c, check)
+        assert_(c.info, {'inputs': [0]})
+        b = np.zeros_like(c)
+        c = np.add.reduceat(a, indices, 1, None, b)
+        assert_equal(c, check)
+        assert_(c is b)
+        assert_(c.info, {'inputs': [0], 'outputs': [0]})
+        # and a few tests for at
+        d = np.array([[1, 2, 3], [1, 2, 3]])
+        check = d.copy()
+        a = d.copy().view(A)
+        np.add.at(check, ([0, 1], [0, 2]), 1.)
+        np.add.at(a, ([0, 1], [0, 2]), 1.)
+        assert_equal(a, check)
+        assert_(a.info, {'inputs': [0]})
+        b = np.array(1.).view(A)
+        a = d.copy().view(A)
+        np.add.at(a, ([0, 1], [0, 2]), b)
+        assert_equal(a, check)
+        assert_(a.info, {'inputs': [0, 2]})
 
 
 class TestChoose(TestCase):
diff --git a/numpy/doc/subclassing.py b/numpy/doc/subclassing.py
index 36d8ff97d..51d9dc120 100644
--- a/numpy/doc/subclassing.py
+++ b/numpy/doc/subclassing.py
@@ -489,6 +489,8 @@ following.
                 return NotImplemented
 
             if method == 'at':
+                if isinstance(inputs[0], A):
+                    inputs[0].info = info
                 return
 
             if ufunc.nout == 1: