在 numba 中,njit、vectorize和guvectorize都是用于对 Python 代码进行优化的装饰器。其中,njit用于将 Python 代码以LLVM IR代码的形式编译为本地机器码,提高代码的执行效率;而vectorize和guvectorize则是用于将函数向量化,可以提高并行计算的效率。
下面以一个简单的例子来说明这三个装饰器的使用:
from numba import njit, vectorize, guvectorize import numpy as np
@njit def sum_njit(a): s = 0 for i in range(len(a)): s += a[i] return s
a = np.array([1, 2, 3, 4, 5]) print(sum_njit(a))
@vectorize(['float64(float64, float64)']) def add_vectorize(x, y): return x + y
a = np.array([1, 2, 3, 4, 5], dtype=np.float64) b = np.array([6, 7, 8, 9, 10], dtype=np.float64) print(add_vectorize(a, b))
@guvectorize(['void(float64[:], float64[:], float64[:])'], '(n),(n)->(n)') def add_guvectorize(a, b, out): for i in range(len(a)): out[i] = a[i] + b[i]
a = np.array([1, 2, 3, 4, 5], dtype=np.float64) b = np.array([6, 7, 8, 9, 10], dtype=np.float64) c = np.zeros(5, dtype=np.float64) add_guvectorize(a, b, c) print(c)
在上述示例中,我们分别使用了njit、vectorize和guvectorize来进行代码优化。其中njit