numba 加速 python 例子

下面是一个使用 Numba 加速 Python 代码的例子：将一个二维 NumPy 数组中的所有元素相加。

import numpy as np
from numba import jit

# 定义一个用于计算数组元素和的函数
@jit(nopython=True)
def sum_array(arr):
    s = 0.0
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            s += arr[i][j]
    return s

# 创建一个二维数组
arr = np.random.rand(1000, 1000)

# 计算数组元素的和
result = sum_array(arr)

# 输出结果
print(result)

在上面的代码中，我们首先定义了一个函数 `sum_array`，用于计算一个二维 NumPy 数组中所有元素的和。然后，我们使用 `@jit` 装饰器将该函数编译为本地机器码。最后，我们创建一个随机的二维数组，并使用 `sum_array` 函数计算其元素的和。

由于我们使用了 Numba 进行 JIT 编译，因此该代码的性能得到了显著提升。如果您将上面的代码与不使用 Numba 的代码进行比较，您将会看到巨大的性能差异。