Numba

1. https://numba.pydata.org/
2. https://numba.readthedocs.io/en/stable/index.html

Numba is an open source JIT compiler that translates a subset of Python and NumPy code into fast machine code.

有点类似 taichi ?

# 本代码由 AI 生成
 
from numba import njit, prange
import numpy as np
 
# 示例：计算数组的平方和
@njit(parallel=True)
def parallel_sum_squares(arr):
    n = 0.0
    for i in prange(len(arr)):  # 使用 prange 而不是 range
        n += arr[i] ** 2
    return n
 
# 测试
arr = np.random.rand(10_000_000)
result = parallel_sum_squares(arr)
print("结果:", result)

在 Numba 中，@jit 和 @njit 都是用于加速 Python 函数的装饰器，它们的功能有许多相似之处，但也存在一些关键区别。

主要区别

详细对比

1.

• @jit：

  • 默认情况下，如果 Numba 无法将函数编译为机器代码，它会回退到 Python 的解释模式执行。这可能会导致性能下降，但程序不会因为编译失败而崩溃。
  • 可以通过设置 nopython=True 来禁用回退模式。

  • 示例：

    from numba import jit
     
    @jit
    def add(a, b):
        return a + b  # 如果无法编译，会回退到解释模式

• @njit：

  • 等价于 @jit(nopython=True)，强制 Numba 使用 nopython 模式。
  • 如果编译失败，会直接抛出错误，而不是回退到解释模式。

  • 示例：

    from numba import njit
     
    @njit
    def add(a, b):
        return a + b  # 必须编译成功，否则抛出错误

2.

• @jit：

  • 在默认情况下，@jit 的性能可能稍差，因为它支持回退到解释模式。
  • 如果设置 nopython=True，则性能和 @njit 相当。
• @njit：

  • 始终启用 nopython 模式，专注于性能。

3.

• @jit：

  • 提供了更多的参数配置选项，例如 nopython, cache, parallel 等。
  • 示例：

    @jit(nopython=True, parallel=True)
    def compute(arr):
        return arr ** 2

• @njit：

  • 更简洁，直接启用高性能模式，适合追求性能的代码。
  • 示例：

    @njit
    def compute(arr):
        return arr ** 2

4.

• @jit：

  • 在没有启用 nopython=True 时，如果编译失败，可以回退到解释模式运行，便于调试。
• @njit：

  • 不支持回退模式，直接抛出错误，因此调试时需要确保代码是完全支持 Numba 的。

什么时候使用哪种？

1. 开发和调试阶段：

   • 使用 @jit，因为它支持回退模式，可以帮助快速找到不支持 Numba 的代码部分。

2. 追求性能的生产环境：

   • 使用 @njit，因为它专注于高性能执行，且不允许回退，确保代码已优化。

3. 高性能并行计算：

   • 配合 parallel=True 使用，可以通过 prange 进行多线程加速。

总结

• @njit 是 @jit 的简化形式，等价于 @jit(nopython=True)。
• 如果您需要最高性能并确保代码完全支持 Numba，使用 @njit。
• 如果需要灵活性和调试方便性，使用 @jit。