python - 从 numba vectorize 获取多个返回值

Question

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

我正在为一个函数使用矢量化装饰器，该函数接受一系列参数，执行迭代求根例程，然后返回解决方案。矢量化函数工作得很好，因为所有的参数和返回的结果都是相同的维度。问题是，解决方案并不总是收敛。我想返回一个 True/False 数组，告诉我解是否已经收敛。

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

Answer 1

快速的回答当然是……不。

来自 numba 文档:

While numba.vectorize() will produce a simple ufunc whose core functionality (the function you are decorating) operates on scalar operands and returns a scalar value...

但是，如果感兴趣的是创建一个编译函数，该函数将以矢量化方式处理计算但​​返回多个值，这仍然可以使用 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.]]