将 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

结果甚至都不正确。更糟糕的是，在某些情况下它会抱怨 "Invalid usage of [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 看起来很相关，但是这里的函数只在单个线程上运行，对吧...

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 数组一样，空切片 [:] 指的是整个数组。

将 NumPy 数组操作与 Numba 的 @guvectorize 结合使用

Using NumPy array operations with Numba's @guvectorize

python

numpy

numba