numba 二维数组

numba 是一个用于加速 Python 程序的库，它可以将 Python 代码转换为机器码，从而提高程序的运行速度。对于二维数组，可以使用 numba 提供的 jit 装饰器来加速其运算。下面是一个简单的例子：

import numpy as np
from numba import jit
@jit(nopython=True)
def sum_2d_array(arr):
    m, n = arr.shape
    res = 0
    for i in range(m):
        for j in range(n):
            res += arr[i, j]
    return res

arr = np.random.rand(1000, 1000)
print(sum_2d_array(arr))

在上面的例子中，我们定义了一个函数 `sum_2d_array`，它接受一个二维数组作为参数，并返回数组中所有元素的和。使用 `@jit(nopython=True)` 装饰器可以让 numba 对该函数进行优化。在函数内部，我们使用了两个 for 循环来遍历数组中的所有元素，并将它们累加到变量 `res` 中。最后，我们输出了数组中所有元素的和。