使用带有numba的python类型提示

Question

来自numba网站:

from numba import jit

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

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

Answer 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可以理解注释并从中构建其签名。