使用具有可变数量输入和输出参数的 numba 实现逐元素 ufunc

Question

我正在用 Python 开发一个系统模拟器。我试图允许使用逐元素函数的自定义系统块。 Numpy 提供了一个函数“frompyfunc”来从 python 函数创建通用函数。我想做同样的事情，但使用 numba。问题是我不知道如何使用可变数量的输入和输出参数来做到这一点。

这项工作背后的动机是 numba 的测试速度非常快。下面的例子：

def nb_frompyfunc(func):
    nb_ufunc = nb.njit(func)
    @nb.njit
    def inner(x, y):
        out = np.zeros_like(x)
        for i in range(x.shape[0]):
            for j in range(x.shape[1]):
                out[i, j] = nb_ufunc(x[i, j], y[i, j])
        return out
    return inner

func = lambda x, y: x + y

np_func = np.frompyfunc(func, 2, 1)
nb_func = nb_frompyfunc(func)

x = np.random.rand(1000, 1000)
y = np.random.rand(1000, 1000)

# Force compile
nb_func(x, y)

print('Numpy ufunc time:')
%timeit np_func(x, y)
print('Numba ufunc time:')
%timeit nb_func(x, y)

产生以下结果：

Numpy ufunc time:
78.7 ms ± 516 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
Numba ufunc time:
1.05 ms ± 15.6 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Numba 在处理 1000 x 1000 数据块时速度提高约 80 倍。

不幸的是，上面的 numba 解决方案假设了输入和输出参数的数量。我希望 numba 函数生成器将输入和输出参数的数量作为输入，就像 numpy 解决方案一样。该函数将具有以下签名：

nb_frompyfunc(func, nargin, nargout)

不幸的是，numba 的“nopython”模式不支持使用 varargin 的 Pythonic 解决方案。

我尝试了动态代码生成，但看来我必须将新代码保存到临时 python 文件并将其作为模块导入。这个解决方案应该可行，但我不喜欢生成临时文件的想法。如果我可以在不保存到文件的情况下让动态代码工作，那么我真的很喜欢这种方法。

我提出的另一个解决方案是根据参数数量选择多个生成器，但这不是一个非常优雅的解决方案。它还根据提供的预先编写的生成器数量对参数施加限制。

有什么好的方法可以解决这个问题？我不依赖于我目前的方法。只要创建的矢量化函数也非常快，我就会很高兴。我的后备计划是只使用 numpy 实现，这可能已经足够好了。我主要关心的是，如果这些系统部署到生产中，它们需要能够尽快处理数据。

编辑：添加说明。

Answer 1

Numba 的矢量化正是我一直在寻找的。

例如。

ufunc = nb.vectorize(lambda x, y: x + y)

我不知道 numba 的矢量化装饰器。谢谢 Ken 在评论中向我展示了这一点。