"签名中的错误令牌"与 numba 和 @guvectorize

Question

我正在尝试decorator@guvectorize与输出[m,2]数组的函数。(m取决于输入的大小)

当使用以下函数签名时，我得到一个错误。

import numpy as np
from numba import guvectorize
@guvectorize('void(float64[:], float64, uint32[:,2])', '(m),()->(m,2)', nopython=True)
def array_copy(data, delta, result):
for i, val in np.ndenumerate(data):
# Open i.
i, = i
result[i,:] = [i, int(val)+delta]
data = np.arange(3, dtype='float64')
res = np.zeros((data.shape[0], 2), dtype='uint32')
array_copy(data, 3, res)

此函数仅用于支持示例。下面是错误信息。

Traceback (most recent call last):
File "<ipython-input-17-63f8983bbf61>", line 5, in <module>
def array_copy(data, delta, result):
File "/home/pierre/anaconda3/lib/python3.8/site-packages/numba/np/ufunc/decorators.py", line 177, in wrap
guvec = GUVectorize(func, signature, **kws)
File "/home/pierre/anaconda3/lib/python3.8/site-packages/numba/np/ufunc/decorators.py", line 49, in __new__
return imp(func, signature, identity=identity, cache=cache,
File "/home/pierre/anaconda3/lib/python3.8/site-packages/numba/np/ufunc/ufuncbuilder.py", line 298, in __init__
self.sin, self.sout = parse_signature(signature)
File "/home/pierre/anaconda3/lib/python3.8/site-packages/numba/np/ufunc/sigparse.py", line 49, in parse_signature
outputs = list(parse(outs))
File "/home/pierre/anaconda3/lib/python3.8/site-packages/numba/np/ufunc/sigparse.py", line 35, in parse
raise ValueError('bad token in signature "%s"' % tok[1])
ValueError: bad token in signature "2"

请告诉我怎么了?我如何指定result将有两列?

Answer 1

当前的Numba实现并没有实现Numpy的GUFunc签名中定义的所有功能，特别是固定维度。

技巧是传递结果数组两次，因此结果的维数是从该形参推断出来的。

下面的示例向量化了一个函数，该函数计算3个元素的每个行向量的和和积:

import numpy as np
from numba import guvectorize
@guvectorize('void(float64[:], float64[:], float64[:])',
'(m),(n)->(n)',
nopython=True)
def sum_and_product(data, dummy, result):
result[0] = data[0] + data[1] + data[2]
result[1] = data[0] * data[1] * data[2]

:

data = np.array([(1, 2, 3), (4, 5, 6), (7, 8, 9)], dtype=float)
data
array([[1., 2., 3.],
[4., 5., 6.],
[7., 8., 9.]])
res = np.empty((3, 2), dtype=float)
res = sum_and_product(data, res)
res
array([[  6.,   6.],
[ 15., 120.],
[ 24., 504.]])

注意，签名指定了单个向量的维数，因此，函数中的代码操作的是单个向量。