1. 课程导入与JAX环境搭建

JAX 是 Google 推出的高性能数值计算库，它把 NumPy 风格的数组接口与函数式自动微分、即时编译（JIT）和硬件加速结合在一起。理解 JAX 的第一步，不是记住 API，而是接受它的设计哲学：函数是纯的、变换是可组合的、随机数需要显式管理状态。本节将围绕这三个要点展开，并完成环境搭建与第一个梯度计算示例。

JAX 的设计哲学：变换优先

JAX 的核心抽象是“函数变换”。你写一个普通的 Python 函数，JAX 提供 jax.grad、jax.jit、jax.vmap 等变换器，把原函数“包装”成求导版本、编译版本或向量化版本。这种组合方式让代码保持简洁，同时获得接近手写 C++ 的性能。

与 PyTorch 的动态图不同，JAX 采用函数式风格：参数不保存在模块对象里，而是作为函数输入显式传入。这意味着梯度计算本质上是“对输入求导”，而不是“对模型参数求导”。初学时容易困惑，但一旦适应，会发现它对高阶导数、批处理和并行化的支持非常自然。

另一个关键点是纯函数约束。JAX 的变换要求函数没有副作用：不能修改全局变量，不能依赖外部可变状态。随机数因此需要显式传入 PRNGKey，而不是像 NumPy 那样隐式使用全局种子。这不是限制，而是保证变换可复现、可并行的必要设计。

自动微分的核心概念

自动微分（Automatic Differentiation，AD）不是数值差分，也不是符号求导。它利用链式法则，在程序执行过程中记录基本运算，从而精确计算导数。JAX 采用反向模式（reverse-mode）自动微分，适合“多输入、少输出”的场景，例如神经网络的损失函数。

JAX 提供两个基础求导函数：jax.grad 计算标量函数对输入的梯度；jax.value_and_grad 同时返回函数值和梯度，避免重复计算。对于高阶导数，只需嵌套调用，例如 jax.grad(jax.grad(f)) 得到二阶导数。这种“变换套变换”的能力是 JAX 区别于多数框架的显著特征。

需要区分两个概念：梯度是标量对向量的导数，结果是同形状向量；雅可比矩阵是向量对向量的导数，JAX 用 jax.jacobian 处理。本节示例聚焦标量梯度，后续章节再展开雅可比与海森矩阵。

安装与环境配置

JAX 的安装需要根据硬件选择对应版本。纯 CPU 环境最简单：

pip install --upgrade pip
pip install --upgrade "jax[cpu]"

如果使用 NVIDIA GPU，需要安装 CUDA 对应版本，例如 CUDA 12：

pip install --upgrade "jax[cuda12]"

安装完成后，用以下代码验证环境。注意 jax.devices() 返回可用设备列表，jax.default_backend() 返回当前默认后端：

import jax
import jax.numpy as jnp

print("JAX 版本:", jax.__version__)
print("可用设备:", jax.devices())
print("默认后端:", jax.default_backend())

若输出中包含 CpuDevice 或 GpuDevice，说明环境正常。常见问题是 CUDA 版本与 jaxlib 不匹配，表现为导入时报错或设备列表为空。此时应核对驱动版本，并参考官方安装矩阵选择对应 wheel。

第一个梯度计算示例

下面计算函数 f(x) = x² + 3x + 1 在 x = 2.0 处的导数。手工求导得 f'(x) = 2x + 3，代入得 7.0。用 JAX 验证：

import jax
import jax.numpy as jnp

def f(x):
    return x**2 + 3*x + 1

grad_f = jax.grad(f)
value, grad = jax.value_and_grad(f)(2.0)

print("f(2.0) =", value)      # 11.0
print("f'(2.0) =", grad)      # 7.0

这里 jax.grad(f) 返回一个新函数，输入 2.0 输出梯度。注意输入必须是浮点类型，整数会导致求导报错，因为整数不可微。若函数有多个参数，jax.grad 默认对第一个参数求导，可用 argnums 指定：

def g(x, y):
    return x*y + y**2

# 对第一个参数 x 求偏导
print(jax.grad(g, argnums=0)(2.0, 3.0))  # 3.0
# 对第二个参数 y 求偏导
print(jax.grad(g, argnums=1)(2.0, 3.0))  # 8.0

验证：∂g/∂x = y = 3.0，∂g/∂y = x + 2y = 8.0，结果一致。这说明 JAX 的梯度计算是精确的，而非近似。

理解计算图与变换组合

下图展示了本节的核心关系：原始函数经过 grad 变换得到梯度函数，再经过 value_and_grad 同时获得函数值与梯度；jit 可以包裹任意一层以加速执行。图中还标出了输入类型约束与设备后端，帮助建立“函数—变换—执行”的整体认知。

常见问题与判断依据

初学者最容易遇到三类问题。第一，整数输入报错：jax.grad 要求输入为浮点类型，若传入 2 而非 2.0，会提示“grad requires real-valued inputs”。解决方法是显式转换或使用 jnp.float32。

第二，副作用导致变换失败：在函数内修改列表、打印调试信息或使用全局变量，可能在 jit 下报错或行为异常。判断依据是：若去掉 jit 后结果正确、加上后报错，通常就是副作用问题。

第三，设备未识别：jax.devices() 只返回 CPU，说明 GPU 版 jaxlib 未正确安装。此时应检查 pip list | grep jax 的版本，并确认 CUDA 驱动版本与安装矩阵匹配。

完成本节后，你应当能够：解释 JAX 的函数式变换设计；说明反向模式自动微分与数值差分的区别；独立完成 CPU/GPU 环境安装与验证；写出并运行第一个梯度计算示例，并能判断输入类型、副作用和设备识别三类常见问题。这些是后续高阶导数、jit 编译优化与神经网络训练的基础。