numba guvectorize的奇怪行为

Question

我写了一个函数来测试numba.guvectorize。这个函数取两个numpy数组的乘积，并计算第一个轴之后的和，如下所示:

from numba import guvectorize, float64
import numpy as np
@guvectorize([(float64[:], float64[:], float64)], '(n),(n)->()')
def g(x, y, res):
res = np.sum(x * y)

但是，上面的guvectorize函数返回错误的结果，如下所示:

>>> a = np.random.randn(3,4)
>>> b = np.random.randn(3,4)
>>> np.sum(a * b, axis=1)
array([-0.83053829, -0.15221319, -2.27825015])
>>> g(a, b)
array([4.67406747e-310, 0.00000000e+000, 1.58101007e-322])

是什么导致了这个问题?

Answer 1

函数g()通过res参数接收一个未初始化的数组。给它赋一个新值不会修改传递给函数的原始数组。

您需要替换res的内容(并将其声明为数组):

@guvectorize([(float64[:], float64[:], float64[:])], '(n),(n)->()')
def g(x, y, res):
res[:] = np.sum(x * y)

该函数操作一维向量并返回标量(因此签名(n),(n)->())， guvectorize处理二维输入并返回一维输出。

>>> a = np.random.randn(3,4)
>>> b = np.random.randn(3,4)
>>> np.sum(a * b, axis=1)
array([-3.1756397 ,  5.72632531,  0.45359806])
>>> g(a, b)
array([-3.1756397 ,  5.72632531,  0.45359806])

但是原来的Numpy函数np.sum已经被矢量化和编译了，所以在这个特定的情况下使用guvectorize几乎没有速度增益。

Answer 2

您的a和b数组是二维的，而您的guvectorized函数具有接受1D数组并返回0D标量的签名。您必须修改它以接受2D并返回1D。

在一种情况下，你做np.sum和axis = 1，在另一种情况下，没有它，你必须在两种情况下做同样的事情。

用res[...] = ...代替res = ...。也许在guvectorize的情况下这不是问题，但它可能是Numpy代码中的一个普遍问题，因为你必须分配值而不是变量引用。

在我的情况下，我添加了cache = True参数到guvectorize装饰器，它只通过缓存/重用Numba编译的代码来加速运行，而不是在每次运行时重新编译它。它只是加快了速度。

完整修改后的更正代码见下文:

上网试试!

from numba import guvectorize, float64
import numpy as np
@guvectorize([(float64[:, :], float64[:, :], float64[:])], '(n, m),(n, m)->(n)', cache = True)
def g(x, y, res):
res[...] = np.sum(x * y, axis = 1)
# Test
np.random.seed(0)
a = np.random.randn(3, 4)
b = np.random.randn(3, 4)
print(np.sum(a * b, axis = 1))
print(g(a, b))

输出:

[ 2.57335386  3.41749149 -0.42290296]
[ 2.57335386  3.41749149 -0.42290296]