Extending Numba with @overload

Libraries built on Numba, and teams with shared kernels, need typed implementations of their own APIs.

When a function is not supported in nopython mode, you can register a compiled implementation for it with numba.extending.overload. The pattern has two layers:

@overload(target_func)
def ol_target(x, y):          # typing time: receives Numba TYPES
    if isinstance(x, types.Array):
        def impl(x, y):       # run time: receives values
            ...
        return impl

The outer function runs during type inference with the Numba types of the arguments, not values. It inspects them and returns a Python function implementing the operation for those types, or returns None to say "not supported for these types", which lets other overloads be tried. The returned implementation is compiled like any @njit function. This gives you type-based dispatch at compile time with zero run-time cost.

Related APIs:

@overload_method(types.Array, "name") adds methods to Numba types.
@overload_attribute adds attributes.
@register_jitable is a shortcut for a function that works the same in Python and compiled code.
@intrinsic is the low-level layer: you write the LLVM IR generation yourself, for operations not expressible in Python.
The low-level extension API (typeof_impl, models, lower_builtin) adds entirely new types, which is how libraries teach Numba about their objects.

Overloads are registered globally when the module is imported, so import the module that defines them before compiling code that uses them.

Example

import numpy as np
from numba import njit, types
from numba.extending import overload

def clamp(x, lo, hi):                     # plain Python version
    return min(max(x, lo), hi)

@overload(clamp)
def ol_clamp(x, lo, hi):                  # receives Numba types
    if isinstance(x, types.Array):
        def impl(x, lo, hi):
            out = np.empty_like(x)
            for i in range(x.size):
                out.flat[i] = min(max(x.flat[i], lo), hi)
            return out
        return impl
    if isinstance(x, (types.Integer, types.Float)):
        return lambda x, lo, hi: min(max(x, lo), hi)
    return None                           # unsupported types

@njit
def use(a):
    return clamp(a, 0.0, 1.0), clamp(2.5, 0.0, 1.0)

print(use(np.array([-1.0, 0.5, 3.0])))    # (array([0. , 0.5, 1. ]), 1.0)

Remember

@overload maps a Python function to typed compiled implementations
The outer function sees types and returns an implementation (or None)
overload_method, intrinsic and the low-level API go further

Watch out for

Inspecting argument values inside the outer overload function, which only receives types.

Check yourself

What does the outer function decorated with @overload receive as arguments?

Official docs: numba.readthedocs.io/en/stable/extending/high-level.html

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks