1.13. Automatic module jitting with jit_module

A common usage pattern is to have an entire module containing user-defined functions that all need to be jitted. One option to accomplish this is to manually apply the @jit decorator to each function definition. This approach works and is great in many cases. However, for large modules with many functions, manually jit-wrapping each function definition can be tedious. For these situations, Numba provides another option, the jit_module function, to automatically replace functions declared in a module with their jit-wrapped equivalents.

It’s important to note the conditions under which jit_module will not impact a function:

  1. Functions which have already been wrapped with a Numba decorator (e.g. jit, vectorize, cfunc, etc.) are not impacted by jit_module.
  2. Functions which are declared outside the module from which jit_module is called are not automatically jit-wrapped.
  3. Function declarations which occur logically after calling jit_module are not impacted.

All other functions in a module will have the @jit decorator automatically applied to them. See the following section for an example use case.

Note

This feature is for use by module authors. jit_module should not be called outside the context of a module containing functions to be jitted.

1.13.1. Example usage

Let’s assume we have a Python module we’ve created, mymodule.py (shown below), which contains several functions. Some of these functions are defined in mymodule.py while others are imported from other modules. We wish to have all the functions which are defined in mymodule.py jitted using jit_module.

# mymodule.py

from numba import jit, jit_module

def inc(x):
   return x + 1

def add(x, y):
   return x + y

import numpy as np
# Use NumPy's mean function
mean = np.mean

@jit(nogil=True)
def mul(a, b):
   return a * b

jit_module(nopython=True, error_model="numpy")

def div(a, b):
    return a / b

There are several things to note in the above example:

  • Both the inc and add functions will be replaced with their jit-wrapped equivalents with compilation options nopython=True and error_model="numpy".
  • The mean function, because it’s defined outside of mymodule.py in NumPy, will not be modified.
  • mul will not be modified because it has been manually decorated with jit.
  • div will not be automatically jit-wrapped because it is declared after jit_module is called.

When the above module is imported, we have:

>>> import mymodule
>>> mymodule.inc
CPUDispatcher(<function inc at 0x1032f86a8>)
>>> mymodule.mean
<function mean at 0x1096b8950>

1.13.2. API

Warning

This feature is experimental. The supported features may change with or without notice.

numba.jit_module(**kwargs)

Automatically jit-wraps functions defined in a Python module

Note that jit_module should only be called at the end of the module to be jitted. In addition, only functions which are defined in the module jit_module is called from are considered for automatic jit-wrapping. See the Numba documentation for more information about what can/cannot be jitted.

Parameters:kwargs – Keyword arguments to pass to jit such as nopython or error_model.