使用 jit_module 自动进行模块 JIT 编译

一种常见的用法模式是让整个模块包含所有需要进行 JIT 编译的用户定义函数。实现这一点的一种选择是手动将 @jit 装饰器应用于每个函数定义。这种方法有效,并且在许多情况下都很好。然而,对于包含许多函数的大型模块,手动对每个函数定义进行 jit 封装可能会很繁琐。对于这些情况,Numba 提供了另一种选择,即 jit_module 函数,它可以自动将模块中声明的函数替换为其经过 jit 封装的等效函数。

需要注意的是,在以下条件下,jit_module不会影响函数:

  1. 已使用 Numba 装饰器(例如 jitvectorizecfunc 等)封装的函数不受 jit_module 的影响。

  2. 在调用 jit_module 的模块之外声明的函数不会自动进行 jit 封装。

  3. 在调用 jit_module 之后在逻辑上出现的函数声明不受影响。

模块中的所有其他函数将自动应用 @jit 装饰器。请参阅以下部分以获取示例用例。

注意

此功能供模块作者使用。jit_module 不应在包含要 JIT 编译的函数的模块上下文之外调用。

示例用法

假设我们创建了一个 Python 模块 mymodule.py(如下所示),其中包含几个函数。其中一些函数在 mymodule.py 中定义,而另一些则从其他模块导入。我们希望使用 jit_modulemymodule.py 中定义的所有函数进行 JIT 编译。

# 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

在上面的示例中有几点需要注意:

  • incadd 函数都将替换为其经过 jit 封装的等效函数,并带有 编译选项 nopython=Trueerror_model="numpy"

  • 由于 mean 函数定义在 NumPy 中,不在 mymodule.py 内,因此不会被修改。

  • mul 不会被修改,因为它已手动使用 jit 装饰。

  • div 不会自动进行 jit 封装,因为它是在调用 jit_module 之后声明的。

当导入上述模块时,我们得到:

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

API

警告

此功能是实验性的。支持的功能可能会更改,恕不另行通知。

numba.jit_module(**kwargs)

自动对 Python 模块中定义的函数进行 jit 封装

请注意,jit_module 应该只在要进行 JIT 编译的模块的末尾调用。此外,只有从调用 jit_module 的模块中定义的函数才会被考虑进行自动 JIT 封装。有关哪些可以/不能进行 JIT 编译的更多信息,请参阅 Numba 文档。

参数

kwargs – 要传递给 jit 的关键字参数,例如 nopythonerror_model