python - 在 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 没有自动方式来理解注释并从中构建其签名。