1. 课程导入与环境搭建

JAX 是 Google 推出的高性能数值计算库，它把 NumPy 风格的数组 API 与函数式变换（function transformation）结合起来，让同一段 Python 代码既能自动求导，又能向量化、并行化，还能被即时编译（Just-In-Time，JIT）成高效的机器码。理解 JAX，关键不是记住一堆 API，而是理解它的设计哲学：纯函数 + 可组合变换。JAX 把“计算什么”和“如何变换这段计算”分开，jit、grad、vmap、pmap 就是四种作用于函数的变换器，它们可以任意嵌套组合。

一、JAX 的设计哲学：函数式、可追踪、可组合

JAX 的核心假设是：被变换的函数应当是纯函数——相同输入总得到相同输出，不依赖也不修改外部可变状态。原因在于 JAX 的变换依赖“追踪”（tracing）：当你把函数交给 jit 或 grad 时，JAX 会用抽象的 tracer 值代替真实数组执行一遍函数，记录下所有算子及其依赖关系，形成一张计算图（jaxpr），再对这张图做编译或求导。如果函数内部有随机的 Python 副作用（如修改全局列表、依赖 print 的时机、用 id() 判断对象），追踪阶段看到的行为可能与真实执行不一致，导致结果错误或编译失败。

因此，JAX 中的随机数、状态更新都采用显式传递的方式：随机数用 jax.random.PRNGKey 作为显式参数，状态更新用函数式返回新值而非原地修改。这与 NumPy 的“就地修改”习惯不同，是初学者最容易踩坑的地方。另一个重要概念是静态形状：JIT 编译时数组形状被固定，若形状变化会触发重新编译。实践中应尽量保持批大小、序列长度等维度稳定，或用 jax.lax 中的动态形状算子处理变长场景。

二、四大核心组件及其分工

四个变换器各司其职，又常组合使用：

• jax.jit：把函数编译为 XLA 可执行程序，减少 Python 解释开销并融合算子。适合计算密集、被反复调用的函数，如训练步、前向传播。
• jax.grad：对纯函数做自动微分，默认返回梯度函数。支持 argnums 指定对第几个参数求导，has_aux 同时返回辅助输出（如损失值）。
• jax.vmap：把只处理单样本的函数自动扩展为批处理版本，通过 in_axes 指定各参数的批维度，避免手写循环。
• jax.pmap：把函数映射到多个设备（如多 GPU/TPU）上并行执行，通过 axis_name 配合 jax.lax.pmean 等集合通信算子实现跨设备梯度同步。单机单卡环境下可先不深入，但需知道它是数据并行的入口。

组合顺序通常遵循“由内到外”的语义：jit(grad(f)) 表示先求导再编译整个梯度函数；vmap(grad(f)) 表示对每个样本分别求导再批量化。理解嵌套顺序，是写出正确高效 JAX 代码的关键。

三、安装与环境配置

JAX 的安装需要区分 CPU 与加速器版本。CPU 版最简单，直接安装即可；GPU 版需要匹配 CUDA 与 cuDNN 版本，官方推荐用 pip 指定对应 wheel 源。以下命令以 CPU 环境为例，GPU 环境请按官方文档替换索引地址：

# 建议使用独立虚拟环境
python -m venv jaxenv
source jaxenv/bin/activate        # Windows 用 jaxenv\Scripts\activate

# 安装 CPU 版 JAX（含 jaxlib）
pip install --upgrade "jax[cpu]"

# 验证安装与设备识别
python -c "import jax; print(jax.__version__); print(jax.devices())"

若输出类似 [CpuDevice(id=0)]，说明环境正常；GPU 环境应看到 GpuDevice。安装后建议同时确认 jax.numpy 可正常导入，并注意 JAX 默认使用 32 位浮点，若需 64 位精度要显式开启 jax.config.update("jax_enable_x64", True)，否则容易出现精度相关的“意外”结果。

四、首个 JIT 示例：从 NumPy 习惯到 JAX 习惯

下面用一个带参数的简单函数演示 jit 与 grad 的组合。示例中的数值仅用于说明流程，不代表任何实测性能：

import jax
import jax.numpy as jnp

def loss_fn(w, x):
    # 纯函数：只依赖输入，不修改外部状态
    pred = jnp.dot(x, w)
    return jnp.sum((pred - 1.0) ** 2)

x = jnp.array([1.0, 2.0, 3.0])
w = jnp.array([0.1, 0.2, 0.3])

# 先求导，再编译整个梯度函数
grad_fn = jax.jit(jax.grad(loss_fn))
g = grad_fn(w, x)
print(g)          # 形状与 w 相同，为损失对 w 的梯度

# 首次调用触发编译，后续调用复用编译结果
g2 = grad_fn(w, x)
print(g2)

运行这段代码时，第一次调用 grad_fn 会有明显的编译延迟，第二次调用则快得多——这正是 JIT 的典型特征。判断 JIT 是否生效，可以观察首次调用耗时明显高于后续调用，或用 grad_fn.lower(w, x).compile() 查看编译产物。需要注意的是，若传入的 x 形状改变，JAX 会重新编译，因此实际训练中应固定批形状。