Repository: numba/numba

Issue:
cpu `@guvectorize` doesn't support multiple outputs
<!--

Thanks for opening an issue! To help the Numba team handle your information
efficiently, please first ensure that there is no other issue present that
already describes the issue you have
(search at https://github.com/numba/numba/issues?&q=is%3Aissue).

-->

## Reporting a bug

<!--

Before submitting a bug report please ensure that you can check off these boxes:

-->

- [x] I have tried using the latest released version of Numba (most recent is
 visible in the change log (https://github.com/numba/numba/blob/main/CHANGE_LOG).
- [x] I have included a self contained code sample to reproduce the problem.
  i.e. it's possible to run as 'python bug.py'.

<!--

Please include details of the bug here, including, if applicable, what you
expected to happen!

-->

```python
from numba import guvectorize
import numpy as np

@guvectorize('(n),(n) -> (n),(n)')
def divmod(x, y, div, mod):
    for i in range(x.shape[0]):
        div[i], mod[i] = np.divmod(x[i], y[i])

x = np.arange(5, dtype='i4') + 5
y = np.full_like(x, fill_value=2)
div = np.zeros_like(x)
mod = np.zeros_like(x)
divmod(x, y, div, mod)
```

```bash
$ python reproducer.py
Traceback (most recent call last):
  File "/Users/guilhermeleobas/git/numba/reproducer.py", line 25, in <module>
    divmod(x, y, div, mod)
  File "/Users/guilhermeleobas/git/numba/numba/np/ufunc/gufunc.py", line 357, in __call__
    raise TypeError(msg)
TypeError: Too few arguments for function 'divmod'. Note that the pattern `out = gufunc(Arg1, Arg2, ..., ArgN)` is not allowed. Use `gufunc(Arg1, Arg2, ..., ArgN, out) instead.
```

This is supported in NumPy, so, it is a bug in Numba side:

```python
$ python
Python 3.10.12 | packaged by conda-forge | (main, Jun 23 2023, 22:41:52) [Clang 15.0.7 ] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import numpy as np
>>> np.linalg._umath_linalg.svd_n_s.signature
'(m,n)->(m,n),(n),(n,n)'
```


Discussion:

Resolving patch:
diff --git a/numba/np/ufunc/gufunc.py b/numba/np/ufunc/gufunc.py
--- a/numba/np/ufunc/gufunc.py
+++ b/numba/np/ufunc/gufunc.py
@@ -139,8 +139,7 @@ def _get_ewise_dtypes(self, args):
 
     def _num_args_match(self, *args):
         parsed_sig = parse_signature(self.gufunc_builder.signature)
-        # parsed_sig[1] has always length 1
-        return len(args) == len(parsed_sig[0]) + 1
+        return len(args) == len(parsed_sig[0]) + len(parsed_sig[1])
 
     def _get_signature(self, *args):
         parsed_sig = parse_signature(self.gufunc_builder.signature)
@@ -157,10 +156,12 @@ def _get_signature(self, *args):
             else:
                 l.append(types.Array(ewise_types[idx], ndim, 'A'))
 
+        offset = len(parsed_sig[0])
         # add return type to signature
-        retty = ewise_types[-1]
-        ret_ndim = len(parsed_sig[-1][0]) or 1  # small hack to return scalar
-        l.append(types.Array(retty, ret_ndim, 'A'))
+        for idx, sig_dim in enumerate(parsed_sig[1]):
+            retty = ewise_types[idx + offset]
+            ret_ndim = len(sig_dim) or 1  # small hack to return scalars
+            l.append(types.Array(retty, ret_ndim, 'A'))
 
         return types.none(*l)
 


Test patch:
diff --git a/numba/tests/npyufunc/test_gufunc.py b/numba/tests/npyufunc/test_gufunc.py
--- a/numba/tests/npyufunc/test_gufunc.py
+++ b/numba/tests/npyufunc/test_gufunc.py
@@ -4,7 +4,7 @@
 import numpy as np
 import numpy.core.umath_tests as ut
 
-from numba import void, float32, int64, jit, guvectorize
+from numba import void, float32, float64, int32, int64, jit, guvectorize
 from numba.np.ufunc import GUVectorize
 from numba.tests.support import tag, TestCase
 
@@ -97,6 +97,75 @@ def gufunc(x, y, res):
         self.assertEqual("docstring for gufunc", gufunc.__doc__)
 
 
+class TestMultipleOutputs(TestCase):
+    target = 'cpu'
+
+    def test_multiple_outputs_same_type_passed_in(self):
+        @guvectorize('(x)->(x),(x)',
+                     target=self.target)
+        def copy(A, B, C):
+            for i in range(B.size):
+                B[i] = A[i]
+                C[i] = A[i]
+
+        A = np.arange(10, dtype=np.float32) + 1
+        B = np.zeros_like(A)
+        C = np.zeros_like(A)
+        copy(A, B, C)
+        np.testing.assert_allclose(A, B)
+        np.testing.assert_allclose(A, C)
+
+    def test_multiple_outputs_distinct_values(self):
+
+        @guvectorize('(x)->(x),(x)',
+                     target=self.target)
+        def copy_and_double(A, B, C):
+            for i in range(B.size):
+                B[i] = A[i]
+                C[i] = A[i] * 2
+
+        A = np.arange(10, dtype=np.float32) + 1
+        B = np.zeros_like(A)
+        C = np.zeros_like(A)
+        copy_and_double(A, B, C)
+        np.testing.assert_allclose(A, B)
+        np.testing.assert_allclose(A * 2, C)
+
+    def test_multiple_output_dtypes(self):
+
+        @guvectorize('(x)->(x),(x)',
+                     target=self.target)
+        def copy_and_multiply(A, B, C):
+            for i in range(B.size):
+                B[i] = A[i]
+                C[i] = A[i] * 1.5
+
+        A = np.arange(10, dtype=np.int32) + 1
+        B = np.zeros_like(A)
+        C = np.zeros_like(A, dtype=np.float64)
+        copy_and_multiply(A, B, C)
+        np.testing.assert_allclose(A, B)
+        np.testing.assert_allclose(A * np.float64(1.5), C)
+
+    def test_incorrect_number_of_pos_args(self):
+        @guvectorize('(m),(m)->(m),(m)', target=self.target)
+        def f(x, y, z, w):
+            pass
+
+        arr = np.arange(5, dtype=np.int32)
+
+        # Inputs only, too few
+        msg = "Too few arguments for function 'f'"
+        with self.assertRaises(TypeError) as te:
+            f(arr)
+        self.assertIn(msg, str(te.exception))
+
+        # Inputs and outputs, too many
+        with self.assertRaises(TypeError) as te:
+            f(arr, arr, arr, arr, arr)
+        self.assertIn(msg, str(te.exception))
+
+
 class TestGUFuncParallel(TestGUFunc):
     _numba_parallel_test_ = False
     target = 'parallel'