diff --git a/numba/np/ufunc/gufunc.py b/numba/np/ufunc/gufunc.py index 54fbfc1c617..a089906e901 100644 --- a/numba/np/ufunc/gufunc.py +++ b/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) diff --git a/numba/tests/npyufunc/test_gufunc.py b/numba/tests/npyufunc/test_gufunc.py index e367580fd5c..8fde667f595 100644 --- a/numba/tests/npyufunc/test_gufunc.py +++ b/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): """