核心学习路径（第一遍）
共 119 处必读补充，243 处选读可以在第一遍跳过。
怎么用：按顺序学习每本 notebook 的课程正文（没有标记的讲解和代码都属于主线），再加上下面列出的「🟢 必读」补充；标为「⚪ 选读」的内容第一遍可以跳过，以后遇到相关问题再回头看。每本 notebook 的练习建议至少做前两道：在「✍️ 作答区」写代码，用「🧪 自检」检查，参考答案在本章末尾的附录。
第一阶段：入门课程
入门/01_JAX简介 （必读 3 处 · 选读 4 处可跳过）
入门/02_JIT编译 （必读 5 处 · 选读 9 处可跳过）
- 🔍 补充：
jit的几种写法与一个高频陷阱 - 🔍 补充：逐元素条件优先用
jnp.where，但要小心梯度里的 NaN - 🔍 补充：静态参数的硬性要求与细节
- 🔍 补充：调试打印用
jax.debug.print，调试逻辑用jax.disable_jit() - 🔍 补充：全局变量的“新旧值并存”，以及大数组也要作为参数传入
入门/03_自动微分grad （必读 4 处 · 选读 7 处可跳过）
- 🔍 补充：亲眼看看 jaxpr，以及
jvp/vjp与两种模式的关系 - 🔍 补充：
has_aux=True——函数还要返回其他东西时 - 🔍 补充：与
has_aux组合时的返回结构 - 🔍 补充（专题）：NaN 梯度与「双 where」技巧
入门/04_自动向量化vmap （必读 7 处 · 选读 13 处可跳过）
- 🔍 补充：手动向量化最容易错在哪里
- 🔍 补充：Python
print与jax.debug.print的区别 - 🔍 补充：vmap 内不能用依赖“批数据”的 Python
if - 🔍 专题：逐样本梯度（per-example gradients）
- 🔍 补充：vmap 与随机数——对 key 做映射
- 🔍 补充（环境实测）：大数组最好作为参数传入，不要被 jit 函数闭包捕获
- 🔍 补充：批处理下的
lax.cond会变成select，两个分支都会执行
入门/05_多设备并行pmap （必读 2 处 · 选读 3 处可跳过）
入门/06_状态管理 （必读 4 处 · 选读 6 处可跳过）
- 🔍 补充：jit 中修改/读取全局变量的"只在追踪时执行一次"陷阱
- 🔍 补充：
lax.scan的 carry 就是状态 —— 用逐步输出记录轨迹 - 🔍 补充：随机数 key 本身就是一种需要显式传递的状态
- 🔍 补充：用
out_axes避免共享状态被无谓复制
入门/07_实战专题 （必读 16 处 · 选读 41 处可跳过）
- 专题 1 重新编译：触发条件、诊断与治理
- 1.2 原理：jit 缓存键里有什么
- 1.5 ✅ 治理手段
- 💡 实战规则
- 2.2 双
where：原理和通用写法 - 2.3 logsumexp / softmax / log_softmax 的稳定写法
- 2.4 用
jax_debug_nans定位 NaN（0.12 实测行为） - 专题 3 jacfwd / jacrev / jvp / vjp 选型
- 3.4 ❌ 常见错误：只需要 $Jv$ 却先构造 $J$
- 4.3 ❌ 常见错误写法
- 5.4 ❌ 常见错误写法与实测
- 6.5 ❌ 常见错误与实测
- 6.6 ✅ 训练 step 的标准模板
- 7.4 ❌ 常见错误与实测
- 7.5 TrainState：
register_dataclass与NamedTuple - ✔ 实战规则清单
第二阶段：排错与随机数（进阶 08、10，建议学完入门后立即看）
进阶/08_调试工具链 （必读 14 处 · 选读 36 处可跳过）
- 0 症状速查表
- 1 看值：在编译后的代码里看到运行时数值
- 1.4
jax.disable_jit()：既是调试手段，也是二分工具 - 2 看程序：JAX 到底追踪、编译出了什么
- 3 数值问题：NaN、inf、越界、除零和错误的梯度
- 3.2
jax.experimental.checkify - 4.1
TracerBoolConversionError - 4.5 布尔掩码索引、形状和切片里的 tracer
- 4.6
UnexpectedTracerError：tracer 泄漏 - 5.2
jax_explain_cache_misses：三种解释，一种「失灵」 - 5.5 速度慢：排查顺序
- 7.3
jax.live_arrays()：找出没释放的数组 - 9.4 演练小结
- 本章小结
进阶/10_随机数与分片 （必读 12 处 · 选读 20 处可跳过）
- 1 设计理念：为什么不用全局状态
- 4.5 「随机数不保证跨版本一致」，以及第 1 章的 10.81
- 5 key 管理模式
- 5.2
split(key, n)一次产生多个 - 5.3
fold_in(key, step)：给训练步 / 数据下标派生 key - 5.4 在
lax.scan中把 key 放进 carry - 5.5
vmap时对 key 数组做映射 - 6 常见陷阱
- 6.2 在 jit 里闭包捕获 key：随机数被固化成常量
- 6.4 Python
random/ NumPy 随机数在 jit 里只在追踪时执行一次 - 8 实战规则清单
- 本章小结
第三阶段：进阶课程第 1–6 章
进阶/01_高级变换与控制流 （必读 6 处 · 选读 10 处可跳过）
- 🔍 补充：读懂 carry 类型不一致的报错
- 🔍 补充：
lax.switch与lax.select_n - 🔍 补充：
cond在vmap下会变成select - 1.6 高级掩码技术
- 课程代码：注意力分数的掩码
- 🔍 补充：闭包捕获大数组的四个问题
进阶/02_JAX代码性能优化 （必读 2 处 · 选读 6 处可跳过）
进阶/03_JAX分布式计算 （必读 6 处 · 选读 11 处可跳过）
- 🔍 补充：三种并行的内存账和通信账
- 设备数据驻留
- 🔍 补充：同一个 SPMD 程序的三种「视角」
- 其他集合操作
- 🔍 补充：二维网格与
PartitionSpec的组合 - 🔍 补充：
AxisType—— Explicit / Auto / Manual
进阶/04_进阶自动微分 （必读 6 处 · 选读 9 处可跳过）
- 🔍 补充：反向模式 = 线性化 + 转置
- 🔍 补充：签名细节、residuals 与
nondiff_argnums - 🔍 补充：用
jax.test_util.check_grads验证自定义规则 - 🔍 补充：JVP 规则必须关于切向量线性
- 🔍 补充：各控制流原语的可微性一览（实测）
- 🔍 补充：常见「不可微」操作的实测梯度行为
进阶/05_互操作性与自定义操作 （必读 1 处 · 选读 7 处可跳过）
进阶/06_大规模模型训练 （必读 5 处 · 选读 8 处可跳过）
- 🔍 补充：每个参数到底要多少字节？激活又有多大？
- 🔍 补充：Flax 0.12 中 NNX 的几个"新规矩"
- 🔍 补充：累积什么时候不等于大 batch？
- 🔍 补充：用
lax.scan堆叠层 + 检查点是大模型的标准结构 - 注意事项
第四阶段：按需选学（进阶 07、09、11）
进阶/09_导出与部署 （必读 6 处 · 选读 12 处可跳过）
- §2
jax.export：把计算图导出成可移植的程序 - §2.2 序列化到文件，在全新子进程中运行
- §2.3 形状多态：一次导出，适配任意 batch
- §2.6 版本兼容性：三层版本号
- §3.2
CheckpointManager：多步保存与保留策略 - §3.4 与 Flax NNX 配合：保存、恢复、继续训练
进阶/07_Pallas自定义内核 （必读 9 处 · 选读 17 处可跳过）
- §1 为什么需要 Pallas
- §2 编程模型
- 2.2
grid、BlockSpec与pl.program_id - 3.2 越界的块：读到的是「垃圾」，interpret 模式下是 NaN
- §4 进阶 kernel
- 4.3 带掩码的边界处理
- 4.4 跨 grid 的归约
- 5.3
grad：默认不可微，用custom_vjp+ Pallas 反向 kernel - §6 GPU/TPU 上的注意事项