Automatic module jitting with
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,
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
It’s important to note the conditions under which
jit_module will not
impact a function:
- Functions which have already been wrapped with a Numba decorator (e.g.
cfunc, etc.) are not impacted by
- Functions which are declared outside the module from which
jit_moduleis called are not automatically
- Function declarations which occur logically after calling
jit_moduleare 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.
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.
Let’s assume we have a Python module we’ve created,
below), which contains several functions. Some of these functions are defined
mymodule.py while others are imported from other modules. We wish to have
all the functions which are defined in
mymodule.py jitted using
# 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
addfunctions will be replaced with their
jit-wrapped equivalents with compilation options
meanfunction, because it’s defined outside of
mymodule.pyin NumPy, will not be modified.
mulwill not be modified because it has been manually decorated with
divwill not be automatically
jit-wrapped because it is declared after
When the above module is imported, we have:
>>> import mymodule >>> mymodule.inc CPUDispatcher(<function inc at 0x1032f86a8>) >>> mymodule.mean <function mean at 0x1096b8950>
This feature is experimental. The supported features may change with or without notice.
jit-wraps functions defined in a Python module
jit_moduleshould only be called at the end of the module to be jitted. In addition, only functions which are defined in the module
jit_moduleis 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