Merge pull request #6791 from sklam/misc/release0.53.0rc3

Preparing for 0.53.0rc3

numba/np/ufunc/gufunc.py

@@ -116,11 +116,14 @@ def _get_signature(self, *args):
 
         return types.none(*l)
 
-    def __call__(self, *args):
+    def __call__(self, *args, **kwargs):
         # If compilation is disabled OR it is NOT a dynamic gufunc
         # call the underlying gufunc
         if self._frozen or not self.is_dynamic:
-            return self.ufunc(*args)
+            return self.ufunc(*args, **kwargs)
+        elif "out" in kwargs:
+            # If "out" argument is supplied
+            args += (kwargs.pop("out"),)
 
         if self._num_args_match(*args) is False:
             # It is not allowed to call a dynamic gufunc without
@@ -136,8 +139,7 @@ def __call__(self, *args):
         # at this point we know the gufunc is a dynamic one
         ewise = self._get_ewise_dtypes(args)
         if not (self.ufunc and ufunc_find_matching_loop(self.ufunc, ewise)):
-            self._is_dynamic = True
             sig = self._get_signature(*args)
             self.add(sig)
             self.build_ufunc()
-        return self.ufunc(*args)
+        return self.ufunc(*args, **kwargs)

numba/tests/npyufunc/test_gufunc.py

@@ -62,6 +62,25 @@ def test_ufunc_like(self):
 
         np.testing.assert_equal(out, x * x + x)
 
+    def test_axis(self):
+        # issue https://github.com/numba/numba/issues/6773
+        @guvectorize(["f8[:],f8[:]"], "(n)->(n)")
+        def my_cumsum(x, res):
+            acc = 0
+            for i in range(x.shape[0]):
+                acc += x[i]
+                res[i] = acc
+
+        x = np.ones((20, 30))
+        # Check regular call
+        y = my_cumsum(x, axis=0)
+        expected = np.cumsum(x, axis=0)
+        np.testing.assert_equal(y, expected)
+        # Check "out" kw
+        out_kw = np.zeros_like(y)
+        my_cumsum(x, out=out_kw, axis=0)
+        np.testing.assert_equal(out_kw, expected)
+
 
 class TestGUFuncParallel(TestGUFunc):
     _numba_parallel_test_ = False
@@ -74,8 +93,8 @@ class TestDynamicGUFunc(TestCase):
     def test_dynamic_matmul(self):
 
         def check_matmul_gufunc(gufunc, A, B, C):
-            gufunc(A, B, C)
             Gold = ut.matrix_multiply(A, B)
+            gufunc(A, B, C)
             np.testing.assert_allclose(C, Gold, rtol=1e-5, atol=1e-8)
 
         gufunc = GUVectorize(matmulcore, '(m,n),(n,p)->(m,p)',
@@ -98,8 +117,12 @@ def test_dynamic_ufunc_like(self):
 
         def check_ufunc_output(gufunc, x):
             out = np.zeros(10, dtype=x.dtype)
+            out_kw = np.zeros(10, dtype=x.dtype)
             gufunc(x, x, x, out)
-            np.testing.assert_equal(out, x * x + x)
+            gufunc(x, x, x, out=out_kw)
+            golden = x * x + x
+            np.testing.assert_equal(out, golden)
+            np.testing.assert_equal(out_kw, golden)
 
         # Test problem that the stride of "scalar" gufunc argument not properly
         # handled when the actual argument is an array,
@@ -140,6 +163,25 @@ def sum_row(inp, out):
         with self.assertRaisesRegex(TypeError, msg):
             sum_row(inp)
 
+    def test_axis(self):
+        # issue https://github.com/numba/numba/issues/6773
+        @guvectorize("(n)->(n)")
+        def my_cumsum(x, res):
+            acc = 0
+            for i in range(x.shape[0]):
+                acc += x[i]
+                res[i] = acc
+
+        x = np.ones((20, 30))
+        expected = np.cumsum(x, axis=0)
+        # Check regular call
+        y = np.zeros_like(expected)
+        my_cumsum(x, y, axis=0)
+        np.testing.assert_equal(y, expected)
+        # Check "out" kw
+        out_kw = np.zeros_like(y)
+        my_cumsum(x, out=out_kw, axis=0)
+        np.testing.assert_equal(out_kw, expected)
 
 class TestGUVectorizeScalar(TestCase):
     """