Using python type hints with numba

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

MSeifert · Accepted Answer · 2017-02-18T02:32:16.953

16

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.

edited Feb 18 '17 at 02:32

answered Feb 18 '17 at 02:26

MSeifert

145,886
38
333
352

4

Seems like it should be pretty simple to roll your own decorator that extracts the type annotations, translates them to numba, and then wraps the function by calling jit with the translated types. – Mike Aug 31 '20 at 13:58
@Mike If it was that simple, why doesn't numba have it already? It is a pretty huge project. Therefore I assume this is very hard to safely do in edge-cases. – julaine May 30 '23 at 07:16

score 2 · Answer 2 · answered Feb 19 '17 at 16:27

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.

Using python type hints with numba

2 Answers2