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.
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.