Python硬件加速: JIT vs JAX

1. JIT定义

JIT（Just-In-Time，即时编译）是一种优化技术，它在程序运行时动态编译代码（如Python函数），并将编译结果缓存以避免重复编译，从而加速后续执行。

核心机制

• 编译阶段：首次运行时，将高阶语言（如Python）代码编译为底层机器码（如CPU/GPU指令）。

• 缓存阶段：保存编译结果，当输入参数的结构（如数组形状、类型）相同时直接复用缓存。

应用场景

• 数值计算：如NumPy、JAX中的数组运算。

• 深度学习：PyTorch/TensorFlow的模型推理优化。

示例（Python装饰器）

from functools import lru_cache
import numba

@numba.jit  # JIT编译装饰器
def fast_sum(x):
    return x.sum()

@lru_cache  # 缓存装饰器（纯Python函数）
def cached_func(x):
    return x * 2

2.