Numba

Question

Numba 似乎支持 jit-ed 函数中可变数量的参数：

@numba.jit(nopython=True)
def func(*arrays):
    return np.vstack(arrays)

但是，我想知道为此函数提供签名的正确语法是什么，以便 numba 知道元组 arrays 的成员期望哪种类型。我在文档中找不到对此的引用。

Answer 1

我不确定您是否应该尝试手动提供类型签名，因为 Numba 能够进行类型推断。引用 advanced types 上的 Numba 文档：

Warning

The APIs documented here are not guaranteed to be stable. Unless necessary, it is recommended to let Numba infer argument types by using the signature-less variant of @jit.

为了演示，观察我们可以在 nopython 模式下编译此函数，Numba 将在使用时构建一个签名列表：

import numba
import numpy as np


@numba.jit(nopython=True)
def func(*arrays):
    return np.vstack(arrays)


print("signatures before call", func.nopython_signatures)

arr1 = np.array([1, 2])
arr2 = np.array([3, 4])
arr3 = func(arr1, arr2)
print(arr3)

print("signatures after call", func.nopython_signatures)

结果输出：

signatures before call []
[[1 2]
 [3 4]]
signatures after call [(UniTuple(array(int64, 1d, C) x 2),) -> array(int64, 2d, C)]

Numba - 如何为可变数量的函数参数声明签名？

Numba - how to declare signature for variable number of function arguments?

python