python - numba @vectorize 目标 ='parallel' TypeError

Question

如果我定义

import numba as nb
import numpy as np
@nb.vectorize
def nb_vec(x):
    if x>0:
        x=x+100
    return x

然后

x=np.random.random(1000000)
nb_vec(x)

运行没有问题

但是如果我添加像这样的目标选项

@nb.vectorize(target='parallel')
def nb_vec(x):
    if x>0:
        x=x+100
    return x

然后

x=np.random.random(1000000)
nb_vec(x)

输出错误信息

--------------------------------------------------------------------------- TypeError Traceback (most recent call last) in 1 x=np.random.random(1000000) ----> 2 nb_vec(x)

TypeError: ufunc 'nb_vec' not supported for the input types, and the inputs could not be safely coerced to any supported types according to the casting rule ''safe''

怎么了？

Answer 1

在 numba 0.46 中，没有签名的 numba.vectorize 装饰器将创建一个动态通用函数，这意味着它会在调用时根据类型编译代码。因此您无需提供签名。

import numpy as np
import numba as nb

@nb.vectorize()
def nb_vec(x):
    if x > 0:
        x = x + 100
    return x

>>> nb_vec
<numba._DUFunc 'nb_vec'>
>>> nb_vec.types
[]
>>> nb_vec(np.ones(5))
array([101., 101., 101., 101., 101.])
>>> nb_vec.types
['d->d']

但是，如果您指定 target='parallel'，那么它将创建一个普通的通用函数。所以它只支持提供的 签名。在您的情况下，您省略了签名，因此它实际上不支持任何输入。

import numpy as np
import numba as nb

@nb.vectorize(target='parallel')
def nb_vec(x):
    if x > 0:
        x = x + 100
    return x

>>> nb_vec
<ufunc 'nb_vec'>
>>> nb_vec.types
[]

这里的解决方案是在使用并行向量化时指定具有适当类型的签名:

import numpy as np
import numba as nb

@nb.vectorize(
    [nb.int32(nb.int32), 
     nb.int64(nb.int64), 
     nb.float32(nb.float32), 
     nb.float64(nb.float64)], 
    target='parallel')
def nb_vec(x):
    if x > 0:
        x = x + 100
    return x