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.