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

1.5.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.5.2. Compilation options

The generated_jit() decorator supports the same keyword-only arguments as the jit() decorator, for example the nopython and cache options.