1.4. 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 fulling retaining runtime execution speed of a JIT function.
1.4.1. 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.
1.4.2. Compilation options¶
The generated_jit()
decorator supports the same keyword-only
arguments as the jit()
decorator, for example the nopython
and cache
options.