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:
jit
, vectorize
, cfunc
, etc.) are not impacted by jit_module
.jit_module
is called are not automatically jit
-wrapped.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.
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:
inc
and add
functions will be replaced with their
jit
-wrapped equivalents with compilation options
nopython=True
and error_model="numpy"
.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>
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 . |
---|