CODEWITHSUNDEEP
pythonnumpynumba
DennisDai
DennisDai

Reputation: 105

How do I vectorize a function which has multiple outputs with Numba?

For example, I want to vectorize the following function:

@nb.njit(nb.types.UniTuple(nb.float64,2)(nb.float64, nb.float64))
def add_subtract(x,y):
    return x+y, x-y

However, when I use @numba.vectorize like this:

@nb.vectorize([nb.types.UniTuple(nb.float64,2)(nb.float64, nb.float64)])
def add_subtract(x,y):
    return x+y, x-y

It will show:

NotImplementedError                       Traceback (most recent call last)
<ipython-input-80-4e6510dce1de> in <module>
      1 @nb.vectorize([nb.types.UniTuple(nb.float64,2)(nb.float64, nb.float64)])
----> 2 def add_subtract(x,y):
      3     return x+y, x-y

E:\anaconda3\lib\site-packages\numba\np\ufunc\decorators.py in wrap(func)
    118         vec = Vectorize(func, **kws)
    119         for sig in ftylist:
--> 120             vec.add(sig)
    121         if len(ftylist) > 0:
    122             vec.disable_compile()

E:\anaconda3\lib\site-packages\numba\np\ufunc\dufunc.py in add(self, sig)
    168         """
    169         args, return_type = sigutils.normalize_signature(sig)
--> 170         return self._compile_for_argtys(args, return_type)
    171 
    172     def _compile_for_args(self, *args, **kws):

E:\anaconda3\lib\site-packages\numba\np\ufunc\dufunc.py in _compile_for_argtys(self, argtys, return_type)
    220         actual_sig = ufuncbuilder._finalize_ufunc_signature(
    221             cres, argtys, return_type)
--> 222         dtypenums, ptr, env = ufuncbuilder._build_element_wise_ufunc_wrapper(
    223             cres, actual_sig)
    224         self._add_loop(utils.longint(ptr), dtypenums)

E:\anaconda3\lib\site-packages\numba\np\ufunc\ufuncbuilder.py in _build_element_wise_ufunc_wrapper(cres, signature)
    177     # Get dtypes
    178     dtypenums = [as_dtype(a).num for a in signature.args]
--> 179     dtypenums.append(as_dtype(signature.return_type).num)
    180     return dtypenums, ptr, cres.environment
    181 

E:\anaconda3\lib\site-packages\numba\np\numpy_support.py in as_dtype(nbtype)
    149     if isinstance(nbtype, types.PyObject):
    150         return np.dtype(object)
--> 151     raise NotImplementedError("%r cannot be represented as a Numpy dtype"
    152                               % (nbtype,))
    153 

NotImplementedError: UniTuple(float64 x 2) cannot be represented as a Numpy dtype

If I rewrite this function like this:

@nb.vectorize([nb.float64[:](nb.float64, nb.float64)])
def add_subtract(x,y):
    return np.array([x+y, x-y])

It will still show:

NotImplementedError: array(float64, 1d, A) cannot be represented as a Numpy dtype

How do I vectorize this function? Is it possible to vectorize a function which has multiple outputs?

Upvotes: 4

Views: 638

Answers (1)

Bob
Bob

Reputation: 508

vectorize only works on a single scalar output (broadcasted to the dimensions of your input vector). As a workaround you can use guvectorize:

import numpy as np
from numba import guvectorize


@guvectorize(
    ["void(float64[:], float64[:] , float64[:], float64[:])"], "(),()->(),()"
)
def add_subtract(x, y, s, d):
    s[:] = x + y
    d[:] = x - y


dim = (2, 3, 4)
x = np.random.random_sample(dim)
y = np.random.random_sample(dim)
s, d = add_subtract(x, y)

Upvotes: 2

Related Questions