将 NumPy 数组操作与 Numba 的 @guvectorize 一起使用

Question

我最近一直在尝试使用 Numba，但我仍然无法理解：

在带有 NumPy 数组的普通 Python 函数中，您可以执行以下操作：

# Subtracts two NumPy arrays and returns an array as the result
def sub(a, b):
    res = a - b
    return res

但是，当你像这样使用 Numba 的 @guvectorize 装饰器时：

# Subtracts two NumPy arrays and returns an array as the result
@guvectorize(['void(float32[:], float32[:], float32[:])'],'(n),(n)->(n)')
def subT(a, b, res):
    res = a - b

结果甚至不正确。更糟糕的是，在某些情况下它会抱怨“[math operator] with [parameters] 的无效使用”

我很困惑。即使我试试这个：

# Subtracts two NumPy arrays and returns an array as the result
@guvectorize(['void(float32[:], float32[:], float32[:])'],'(n),(n)->(n)')
def subTt(a, b, res):
    res = np.subtract(a,b)

结果仍然不正确。考虑到this is supposed to be a supported Math operation，我不明白为什么它不起作用。

我知道标准的方式是这样的：

# Subtracts two NumPy arrays and returns an array as the result
@guvectorize(['void(float32[:], float32[:], float32[:])'],'(n),(n)->(n)')
def subTtt(a, b, res):
    for i in range(a.shape[0]):
        res[i] = a[i] - b[i]

这确实按预期工作。

但是我的方式有什么问题？

P/S 这只是一个简单的例子来解释我的问题，我实际上并不打算使用 @guvectorize 来减去数组：P P/P/S 我怀疑它与如何将数组复制到 gpu 内存有关，但我不确定...... P/P/P/S This 看起来很相关，但这里的函数只在单个线程上运行...

Answer 1

正确的写法是：

@guvectorize(['void(float32[:], float32[:], float32[:])'],'(n),(n)->(n)')
def subT(a, b, res):
    res[:] = a - b

您尝试的方法不起作用的原因是 python 语法的限制不是 numba 特有的。

name = expr 将 name 的值重新绑定到 expr，它永远不会改变 name 的原始值，就像你可以使用的那样，例如c++ 参考。

name[] = expr 调用（本质上），name.__setitem__ 可用于修改名称，就像 numpy 数组一样，空切片 [:] 指的是整个数组。