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函数returns错误的结果如下图：

>>> 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)

该函数对一维向量和 returns 标量（因此签名 (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 函数具有接受一维数组和 returning 0D 标量的签名。您必须修改它以接受 2D 和 return 1D。

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

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

在我的例子中，我将 cache = True 参数添加到 guvectorize 装饰器，它只会通过 caching/re-using Numba 编译代码加速运行ning，而不是在每个 [=36] 上重新编译它=].它只会加快速度。

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

Try it online!

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]

numba guvectorize 的奇怪行为

weird behavior of numba guvectorize

python

numpy

numba