Python：使用 Numba 签名

Question

我有一个函数正在做一些计算，并在某个时刻调用另一个函数。例如，主函数是这样的：

import numba

@numba.njit(some signature here)
def my_funct():
    ...
    value = cosd(angle)

由于函数 cosd 位于另一个用 numba.njit 修饰的函数中，它也必须被修饰，在我的例子中是：

from numba import float64

@numba.njit(float64(float64))
def cosd(angle):
    return np.cos(np.radians(angle))

我现在的问题是，在另一个函数中，输入值angle是一个数组，相关的输出也是一个数组。我知道我可以将我的函数装饰为 @numba.njit(float64[:](float64[:])) 但这样做该函数将不再接受标量。我怎样才能知道 numba 输入类似于 Union[float64, float64[:]]？当然这也适用于输出。非常感谢！

Answer 1

我终于自己找到了答案。解决方案是创建一个签名列表，对于我的示例，它将是：

from numba import float64

@njit([float64(float64), float64[:](float64[:])])
def cosd(angle):
    return np.cos(np.radians(angle))

希望对其他人有所帮助。

Python：使用 Numba 签名

Python: signature with Numba

python-3.x

numba