使用带有numba的python类型提示

Question

From the numba website:

from numba import jit

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

Is there a way to have numba use python type hints (if provided)?

Answer 1

Yes and no. You can simply use the normal python syntax for annotations (the jit decorator will preserve them). Building on your simple example:

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
[]

However to explicitly (enforce) the signature it must be given in the 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

There is as far as I know no automatic way for jit to understand the annotations and build its signature from these.

Answer 2

Since it is Just In Time compilation, you must execute the function to generate the signatures

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

A new signature is generated each time the previous doesn't fit the data.