Flexible specializations with @generated_jit
¶
While the jit()
decorator is useful for many situations,
sometimes you want to write a function that has different implementations
depending on its input types. The generated_jit()
decorator
allows the user to control the selection of a specialization at compile-time,
while fully retaining runtime execution speed of a JIT function.
Example¶
Suppose you want to write a function which returns whether a given value is a “missing” value according to certain conventions. For the sake of the example, let’s adopt the following definition:
- for floating-point arguments, a missing value is a
NaN
- for Numpy datetime64 and timedelta64 arguments, a missing value is a
NaT
- other types don’t have the concept of a missing value.
That compile-time logic is easily implemented using the
generated_jit()
decorator:
import numpy as np
from numba import generated_jit, types
@generated_jit(nopython=True)
def is_missing(x):
"""
Return True if the value is missing, False otherwise.
"""
if isinstance(x, types.Float):
return lambda x: np.isnan(x)
elif isinstance(x, (types.NPDatetime, types.NPTimedelta)):
# The corresponding Not-a-Time value
missing = x('NaT')
return lambda x: x == missing
else:
return lambda x: False
There are several things to note here:
- The decorated function is called with the Numba types of the arguments, not their values.
- The decorated function doesn’t actually compute a result, it returns a callable implementing the actual definition of the function for the given types.
- It is possible to pre-compute some data at compile-time (the
missing
variable above) to have them reused inside the compiled implementation. - The function definitions use the same names for arguments as in the decorated function, this is required to ensure passing arguments by name works as expected.
Compilation options¶
The generated_jit()
decorator supports the same keyword-only
arguments as the jit()
decorator, for example the nopython
and cache
options.