从 numba vectorize 获取多个 return 值

Question

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

我正在为一个函数使用向量化装饰器，该函数接受一系列参数，执行迭代求根例程，然后 return 得出解决方案。矢量化函数完美运行，因为所有参数和 returned 结果都是相同的维度。问题是，解决方案并不总是收敛。我想要 return 一个 True/False 数组来告诉我解是否收敛。

我知道这可以通过 guvectorize 来实现，但是 vectorize 的标量语法很好，而且由于我所有的数组都是相同的维度，所以需要引入 guvectorize 似乎很愚蠢。有谁知道是否可以从 numba.vectorize 编译函数中 return 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...

但是，如果感兴趣的是创建一个编译函数，该函数将以矢量化方式处理计算，但 return 多个值，这仍然可以使用 guvectorize。在下面的例子中，一个函数签名被定义为只有一个维度的签名，但如果它是多维的，它会愉快地对整个数组进行操作。

我知道我在问题中说过，"I know I can do this with 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 获取多个 return 值

Get multiple return values from numba vectorize

python

numba