1.4. Flexible specializations with @generated_jit — Anaconda Enterprise 4.4 documentation

Anaconda Enterprise: Documentation

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