Flexible specializations with
jit() decorator is useful for many situations,
sometimes you want to write a function that has different implementations
depending on its input types. The
allows the user to control the selection of a specialization at compile-time,
while fully 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
for Numpy datetime64 and timedelta64 arguments, a missing value is a
other types don’t have the concept of a missing value.
That compile-time logic is easily implemented using the
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
missingvariable 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.