1. 课程导入与环境搭建

JAX 是 Google 推出的一个面向高性能数值计算与机器学习的 Python 库，它的核心可以概括为一句话：用 NumPy 的写法，获得自动微分、向量化和即时编译的能力。理解这句话，是理解整个 JAX 生态的起点。JAX 并不试图重新发明数组运算的语法，而是把 jax.numpy 设计成与 numpy 高度兼容的接口，让熟悉 NumPy 的人几乎零成本迁移；在此之上，它通过函数变换（function transformation）的方式，把自动微分、批处理和编译优化叠加到普通 Python 函数上。

这种“函数变换”的设计哲学是 JAX 与 PyTorch、TensorFlow 最本质的区别。JAX 不维护计算图对象，也不依赖可变的模型状态，而是把函数当作一等公民：grad 接收一个函数返回另一个函数，jit 接收一个函数返回编译后的函数，vmap 接收一个函数返回自动批处理的函数。这些变换可以自由组合，例如 jit(grad(f)) 或 vmap(grad(f))，这正是 JAX 在科研与高性能优化场景中备受青睐的原因。相应地，JAX 要求函数尽量是纯函数——相同输入产生相同输出、不依赖外部可变状态，这样才能被安全地追踪、编译和变换。

从生态定位看，JAX 更像“数值计算内核”而非完整框架。它本身只提供数组运算、自动微分、编译和并行原语，而神经网络层、优化器、训练循环通常由上层库承担：Flax 与 Equinox 提供模块化的网络定义，Optax 提供梯度变换与优化器组合，Haiku 则是另一种函数式网络写法。这种分层让 JAX 在科研中灵活、在工程中可控，但也意味着学习者需要自己理解训练循环的组装方式，而不是调用一个 model.fit()。

安装 CPU 与 GPU 版本

JAX 的安装有一个容易踩坑的地方：CPU 版与 GPU 版是不同的发行包，不能简单 pip install jax 就指望用上 GPU。官方推荐通过 pip 指定索引源安装。CPU 版最省事：

pip install --upgrade "jax[cpu]"

GPU 版需要匹配 CUDA 与 cuDNN 版本。以 CUDA 12 为例，常用命令为：

pip install --upgrade "jax[cuda12]" \
  -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

安装完成后，务必验证后端是否生效。执行下面的代码，观察输出中的 backend 字段：

import jax
print(jax.devices())
# CPU 环境示例输出: [CpuDevice(id=0)]
# GPU 环境示例输出: [CudaDevice(id=0)]

如果安装了 GPU 版却仍显示 CpuDevice，常见原因是 CUDA 驱动版本不匹配、jaxlib 与 jax 版本不一致，或环境变量 JAX_PLATFORMS 被错误设置。此时可先执行 pip show jax jaxlib 核对版本，再检查 nvidia-smi 显示的驱动是否满足要求。

验证 jax.numpy 与自动微分

环境就绪后，先确认 jax.numpy 的基本行为。它与 NumPy 接口一致，但数组类型是 jax.Array，且默认不可变：

import jax.numpy as jnp
x = jnp.array([1.0, 2.0, 3.0])
print(x, x.dtype, type(x))
# 输出示例: [1. 2. 3.] float32 <class 'jax.Array'>

注意默认浮点精度是 float32，这与 NumPy 默认的 float64 不同，是 JAX 为加速计算做的取舍；如需双精度，可在导入前设置 jax.config.update("jax_enable_x64", True)。接着验证自动微分。以 f(x) = x³ + 2x 为例，其导数为 3x² + 2，在 x = 2 处应为 14：

import jax
def f(x):
    return x**3 + 2*x

df = jax.grad(f)
print(df(2.0))   # 期望 14.0
print(jax.grad(df)(2.0))  # 二阶导 6x = 12.0

若输出为 14.0 与 12.0，说明自动微分链路正常。这里 jax.grad 只对返回标量的函数生效，若函数返回向量，需要配合 jax.jacobian 或指定 argnums 与 has_aux 等参数。

编写第一个 JAX 程序

把上面的能力组合起来，写一个最小但完整的梯度下降示例：用 JAX 求 g(w) = (w - 3)² 的最小值，理论最优解为 w = 3。示例参数为学习率 0.1、迭代 20 次，属于教学假设，实际运行结果会因浮点精度略有差异：

import jax, jax.numpy as jnp

def loss(w):
    return (w - 3.0) ** 2

grad_loss = jax.jit(jax.grad(loss))

w = 0.0
for i in range(20):
    g = grad_loss(w)
    w = w - 0.1 * g
print(w)   # 期望接近 3.0

这段代码体现了 JAX 的典型写法：先定义纯函数，再用 grad 求导、用 jit 编译，最后在 Python 循环中更新参数。需要注意，jit 编译发生在首次调用时，循环中的 w 是普通 Python 标量，每次迭代都会触发一次编译后的函数调用；若把整个循环写进 jax.lax.fori_loop 或 jax.lax.scan，则能把循环也纳入编译，获得更稳定的性能。这是后续章节会反复出现的模式。