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:
Functions which have already been wrapped with a Numba decorator (e.g.
jit
,vectorize
,cfunc
, etc.) are not impacted byjit_module
.Functions which are declared outside the module from which
jit_module
is called are not automaticallyjit
-wrapped.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.
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
andadd
functions will be replaced with theirjit
-wrapped equivalents with compilation optionsnopython=True
anderror_model="numpy"
.The
mean
function, because it’s defined outside ofmymodule.py
in NumPy, will not be modified.mul
will not be modified because it has been manually decorated withjit
.div
will not be automaticallyjit
-wrapped because it is declared afterjit_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>
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 moduleNote 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 modulejit_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 asnopython
orerror_model
.