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

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