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))

我希望这对其他人有帮助。