Get multiple return values from numba vectorize

Question

I am wondering if it is possible to get multiple return values from a function compiled with numba.vectorize.

I am using the vectorize decorator for a function that takes in a series of arguments, and performs and iterative root finding routine and then returns the solution. The vectorize functions works perfectly because all of the arguments and the returned result are all the same dimension. The problem is, the solution doesn't always converge. I'd like to return a True/False array that tells me whether the solution has converged or not.

I realize this can be achieved with guvectorize, but the scalar syntax of vectorize is nice, and since all of my arrays are the same dimension, it seems silly to need to bring in guvectorize. Does anyone know if its possible to return 2 or more arrays from a numba.vectorize compiled function?

Answer 1

The quick answer is of course... No.

From the numba docs:

While numba.vectorize() will produce a simple ufunc whose core functionality (the function you are decorating) operates on scalar operands and returns a scalar value...

However, if the interest is in creating a compiled function that will handle computations in a vectorized manner but return multiple values, this is still possible with guvectorize. In the below example, a function signature is defined with only one dimension in the signature, but it will happily operate on the entire array if it is multidimensional.

I know I said in the question, "I know I can do this with guvectorize" but I did not actually understand this fact that the function would work on an array of unspecified dimensions, essentially as a flattened array. I thought you needed to specify the dimensions exactly, and I haven't yet seen this behavior documented (though it might be).

from numba import guvectorize, float64
import numpy as np

ones = np.ones((3, 3))
twos = ones * 2

@guvectorize(
    [(float64[:], float64[:], float64[:], float64[:])],
    "(n),(n)->(n),(n)",
    nopython=True)
def add_guvectorize(a, b, c, d):
    for i in range(len(a)):
        c[i] = a[i] + b[i]
        d[i] = a[i] + c[i]
threes, fours = add_guvectorize(ones, twos)
print(threes)
print(fours)

prints:

[[3. 3. 3.]
 [3. 3. 3.]
 [3. 3. 3.]]
[[4. 4. 4.]
 [4. 4. 4.]
 [4. 4. 4.]]