Merge pull request #5242 from juliantaylor/fix-ufunc-subok-out

BUG: fix not returning out array from ufuncs with subok=False set

4 files changed, 55 insertions, 6 deletions

diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 466cc2b22..dc5065f14 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -3971,18 +3971,19 @@ _find_array_wrap(PyObject *args, PyObject *kwds,
     PyObject *with_wrap[NPY_MAXARGS], *wraps[NPY_MAXARGS];
     PyObject *obj, *wrap = NULL;
 
-    /* If a 'subok' parameter is passed and isn't True, don't wrap */
+    /*
+     * If a 'subok' parameter is passed and isn't True, don't wrap but put None
+     * into slots with out arguments which means return the out argument
+     */
     if (kwds != NULL && (obj = PyDict_GetItem(kwds,
                                               npy_um_str_subok)) != NULL) {
         if (obj != Py_True) {
-            for (i = 0; i < nout; i++) {
-                output_wrap[i] = NULL;
-            }
-            return;
+            /* skip search for wrap members */
+            goto handle_out;
         }
     }
 
-    nargs = PyTuple_GET_SIZE(args);
+
     for (i = 0; i < nin; i++) {
         obj = PyTuple_GET_ITEM(args, i);
         if (PyArray_CheckExact(obj) || PyArray_IsAnyScalar(obj)) {
@@ -4040,6 +4041,8 @@ _find_array_wrap(PyObject *args, PyObject *kwds,
      * exact ndarray so that no PyArray_Return is
      * done in that case.
      */
+handle_out:
+    nargs = PyTuple_GET_SIZE(args);
     for (i = 0; i < nout; i++) {
         int j = nin + i;
         int incref = 1;
diff --git a/numpy/core/tests/test_numeric.py b/numpy/core/tests/test_numeric.py
index 46e864495..ea145ef81 100644
--- a/numpy/core/tests/test_numeric.py
+++ b/numpy/core/tests/test_numeric.py
@@ -1697,6 +1697,20 @@ class TestStdVar(TestCase):
         assert_almost_equal(std(self.A, ddof=2)**2,
                             self.real_var*len(self.A)/float(len(self.A)-2))
 
+    def test_out_scalar(self):
+        d = np.arange(10)
+        out = np.array(0.)
+        r = np.std(d, out=out)
+        assert_(r is out)
+        assert_array_equal(r, out)
+        r = np.var(d, out=out)
+        assert_(r is out)
+        assert_array_equal(r, out)
+        r = np.mean(d, out=out)
+        assert_(r is out)
+        assert_array_equal(r, out)
+
+
 class TestStdVarComplex(TestCase):
     def test_basic(self):
         A = array([1, 1.j, -1, -1.j])
diff --git a/numpy/core/tests/test_umath.py b/numpy/core/tests/test_umath.py
index fa0ac3fde..c71b7b658 100644
--- a/numpy/core/tests/test_umath.py
+++ b/numpy/core/tests/test_umath.py
@@ -36,6 +36,36 @@ class TestConstants(TestCase):
     def test_euler_gamma(self):
         assert_allclose(ncu.euler_gamma, 0.5772156649015329, 1e-15)
 
+class TestOut(TestCase):
+    def test_out_subok(self):
+        for b in (True, False):
+            aout = np.array(0.5)
+
+            r = np.add(aout, 2, out=aout)
+            assert_(r is aout)
+            assert_array_equal(r, aout)
+
+            r = np.add(aout, 2, out=aout, subok=b)
+            assert_(r is aout)
+            assert_array_equal(r, aout)
+
+            r = np.add(aout, 2, aout, subok=False)
+            assert_(r is aout)
+            assert_array_equal(r, aout)
+
+            d = np.ones(5)
+            o1 = np.zeros(5)
+            o2 = np.zeros(5, dtype=np.int32)
+            r1, r2 = np.frexp(d, o1, o2, subok=b)
+            assert_(r1 is o1)
+            assert_array_equal(r1, o1)
+            assert_(r2 is o2)
+            assert_array_equal(r2, o2)
+
+            r1, r2 = np.frexp(d, out=o1, subok=b)
+            assert_(r1 is o1)
+            assert_array_equal(r1, o1)
+
 
 class TestDivision(TestCase):
     def test_division_int(self):
diff --git a/numpy/ma/core.py b/numpy/ma/core.py
index 501177bad..34e52d86e 100644
--- a/numpy/ma/core.py
+++ b/numpy/ma/core.py
@@ -780,6 +780,8 @@ class _DomainSafeDivide:
         # component of numpy's import time.
         if self.tolerance is None:
             self.tolerance = np.finfo(float).tiny
+        # don't call ma ufuncs from __array_wrap__ which would fail for scalars
+        a, b = np.asarray(a), np.asarray(b)
         return umath.absolute(a) * self.tolerance >= umath.absolute(b)