从numba vectorize获取多个返回值

Question

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

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

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

Answer 1

快速回答当然是……不。

来自 numba 文档：

whilenumba.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.]]