使用带有 numba 的 python 类型提示

Posted 2023-02-16

【中文标题】使用带有 numba 的 python 类型提示

【英文标题】：Using python type hints with numba

【发布时间】：2017-07-07 17:00:26

【问题描述】：

来自 numba 网站：

from numba import jit

@jit
def f(x, y):
    # A somewhat trivial example
    return x + y

有没有办法让 numba 使用 python 类型提示（如果提供）？

【参考方案1】：

是和不是。您可以简单地使用普通的 Python 语法进行注释（jit 装饰器将保留它们）。以您的简单示例为基础：

from numba import jit

@jit
def f(x: int, y: int) -> int:
    # A somewhat trivial example
    return x + y

>>> f.__annotations__
'return': int, 'x': int, 'y': int
>>> f.signatures  # they are not recognized as signatures for jit
[]

但是要显式（强制执行）签名，它必须在jit-decorator 中给出：

from numba import int_

@jit(int_(int_, int_))
def f(x: int, y: int) -> int:
    # A somewhat trivial example
    return x + y

>>> f.signatures
[(int32, int32)]  # may be different on other machines

据我所知，jit 没有自动方式来理解注释并从中构建其签名。

【讨论】：

似乎应该很简单地推出自己的装饰器来提取类型注释，将它们转换为 numba，然后通过使用翻译后的类型调用 jit 来包装函数。

【参考方案2】：

由于是即时编译，所以必须执行生成签名的函数

In [119]: f(1.0,1.0)
Out[119]: 2.0

In [120]: f(1j,1)
Out[120]: (1+1j)

In [121]: f.signatures
Out[121]: [(float64, float64), (complex128, int64)]

每次前一个签名不适合数据时都会生成一个新签名。