python - 如何使用 Numba 对具有多个输出的函数进行矢量化？

Question

例如，我想矢量化以下函数:

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

但是，当我像这样使用 @numba.vectorize 时:

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

它会显示:

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

如果我像这样重写这个函数:

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

它仍然会显示:

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

我如何矢量化这个函数？是否可以对具有多个输出的函数进行矢量化？

Answer 1

vectorize仅适用于单个标量输出(广播到输入向量的维度)。作为解决方法，您可以使用 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)