Question

我想知道是否可以从使用numba.vectorize编译的函数中获取多个返回值。

我正在使用vectorize装饰器来获取一系列参数的函数，并执行迭代根查找例程，然后返回解决方案。 vectorize函数完美地工作，因为所有参数和返回的结果都是相同的维度。问题是，解决方案并不总是趋同。我想返回一个True / False数组，告诉我解决方案是否已收敛。

我意识到这可以通过guvectorize来实现，但是vectorize的标量语法很好，而且由于我的所有数组都是相同的维度，因此需要引入guvectorize似乎很愚蠢。有没有人知道是否可以从numba.vectorize编译函数中返回2个或更多数组？

Answer 1

快速答案当然是……不。

从numba文档中获取：

numba.vectorize()会产生一个简单的ufunc，其核心功能（您正在装饰的功能）按标量运行操作数并返回标量值...

但是，如果感兴趣的是创建一个将以矢量化方式处理计算但返回多个值的编译函数，那么使用guvectorize仍然可以实现。在下面的示例中，函数签名仅在签名中定义了一个维度，但是如果它是多维的，它将很高兴地在整个数组上运行。

我知道我在问题中说过：“我知道我可以使用guvectorize做到这一点”，但实际上我并不理解该功能可以在未指定尺寸的数组上工作，实际上是作为扁平数组工作。我以为您需要准确指定尺寸，而且我还没有看到这种现象的记录（尽管可能是这样）。

from numba import guvectorize, float64
import numpy as np

ones = np.ones((3, 3))
twos = ones * 2

@guvectorize(
    [(float64[:], float64[:], float64[:], float64[:])],
    "(n),(n)->(n),(n)",
    nopython=True)
def add_guvectorize(a, b, c, d):
    for i in range(len(a)):
        c[i] = a[i] + b[i]
        d[i] = a[i] + c[i]
threes, fours = add_guvectorize(ones, twos)
print(threes)
print(fours)

打印：

[[3. 3. 3.]
 [3. 3. 3.]
 [3. 3. 3.]]
[[4. 4. 4.]
 [4. 4. 4.]
 [4. 4. 4.]]

从numba vectorize

1 个答案: