进阶 · 总目录

第 1–6 章跟着 apxml.com 的课程 JAX 高级编程：性能、优化与规模化整理；第 7–11 章是补充章节，讲的是实际用 JAX 时很重要的几个主题，写法和前面各章一样。所有 notebook 都在 JAX 0.12.0.1 / Python 3.14 / CPU 上完整跑过。

先学什么：入门部分 00_目录与导读，尤其是 02 JIT、03 grad、05 并行和 07 实战专题。进阶各章会直接引用这些内容，不再重复讲。

阅读约定（和入门部分一样）

正文代码都用 JAX 0.12 推荐的写法。已经过时的写法放在 # ❌ 旧写法 单元里，后面紧跟 # ✅ 推荐写法。
需要特别留意的地方用 > ⚠️ 注意 标出来，重要结论都配了 assert。
讲解部分按课程要点用中文重新组织，每节开头附有课程原文链接；测验题是根据章节内容出的。
做练习：在题目下面的「✍️ 作答区」写代码，运行「🧪 自检」检查；参考答案在每章最后的附录里。

章节索引¶

章	Notebook	本章亮点	更新说明
1	01_高级变换与控制流	scan 参数细节与 `associative_scan`；掩码技术与 `segment_*`；闭包与暂存（constvars）；综合实践：变长掩码 LSTM、RK4 ODE 积分器	22 条
2	02_JAX代码性能优化	在 CPU 上生成真实的 profiler trace；jaxpr → StableHLO → HLO 全链路；持久化编译缓存（跨进程命中实测）；用 HLO 定量分析融合（访存 3N 对 7N）；定位 CPU bug 的根因	19 条
3	03_JAX分布式计算	每个 pmap 示例都配现代写法；集合通信精确语义；Megatron MLP；FSDP 每设备内存 26144 B → 5956 B；本机 2 进程真实多主机演示	26 条
4	04_进阶自动微分	`linearize` + `linear_transpose`；`jet` 高阶 Taylor 展开；Jacobian 着色；隐函数定理与 custom_vjp；DEQ 隐式层、ODE 伴随	18 条
5	05_互操作性与自定义操作	三种 callback 的实测对比表；DLPack 零拷贝；自定义原语全套规则（jit/grad/vmap）；用 clang++ 编译 C++ 内核，经 `jax.ffi` 接入	29 条
6	06_大规模模型训练	Flax NNX + Optax + Orbax；梯度累积的三种不等价情形；checkpoint 策略定量对比（Transformer 临时内存减少 84%）；bf16/fp16 与 loss scaling；张量并行、FSDP、Adafactor	19 条

补充章节

章	Notebook	内容	建议先读
7	07_Pallas自定义内核	Pallas 编程模型（Ref、grid、BlockSpec）；CPU 上用 interpret 模式验证正确性；分块 matmul、softmax、LayerNorm、掩码；配合 custom_vjp 求导；交叉 lowering 到 TPU 检查块形状规则	第 2 章（融合边界）、第 5 章（自定义原语 / FFI）
8	08_调试工具链	按「症状 → 工具」组织的排错手册：debug.print / callback、jaxpr 与 HLO、debug_nans 与 checkify、各种 Tracer 错误、编译诊断、内存排查；综合演练：修复带 3 个隐藏 bug 的训练脚本	入门 02、07，第 2 章
9	09_导出与部署	参数用 Orbax、程序用 `jax.export`；形状多态、多平台、`vjp_order`、版本兼容性；全新子进程只加载导出文件和参数完成推理	第 2 章（编译缓存）、第 6 章（Orbax）
10	10_随机数与分片	类型化 key 与原始 key；6 种 PRNG 实现；`jax_threefry_partitionable` 对分片生成和跨版本复现的影响；key 管理模式与常见陷阱	入门 01、07，第 3 章
11	11_科学计算与可微分仿真	float64；无矩阵 Poisson 求解（验证二阶收敛）；BCOO 稀疏矩阵；对求解器做隐式微分（内存约省 290 倍）；系数反问题；热方程灵敏度分析；Winslow 网格优化，翻转单元从 23 个降到 0	第 1 章（scan / ODE）、第 4 章（隐函数定理）

补充章节的建议顺序：08 调试工具链（最常用，学完入门就可以看）→ 10 随机数（也解释了入门第 1 章示例输出为什么复现不出来）→ 09 导出与部署 → 07 Pallas（需要手写高性能内核时）→ 11 科学计算（做数值仿真或网格方向时最有价值）。

第 3、6、10 章开头会设置 XLA_FLAGS 来模拟 8 个设备，所以要重启内核，从第一格开始运行。第 5 章需要本机装有 C++ 编译器（macOS 自带 clang++ 即可）。第 9 章的序列化需要 flatbuffers（已写入 requirements.txt）。第 11 章全程使用 float64。

环境自检¶

import sys, shutil, importlib
import jax
print("Python :", sys.version.split()[0])
print("JAX    :", jax.__version__)
for name in ["optax", "flax", "orbax.checkpoint", "flatbuffers"]:
    try:
        m = importlib.import_module(name)
        print(f"{name:7s}:", getattr(m, "__version__", "已安装"))
    except ImportError:
        print(f"{name:7s}: 未安装（第 6、9 章需要）")
print("clang++:", shutil.which("clang++") or shutil.which("c++") or "未找到（第 5 章 C++ FFI 需要）")
print("设备   :", jax.devices())

Python : 3.14.7
JAX    : 0.12.0.1
optax  : 0.2.9
flax   : 0.12.12

orbax.checkpoint: 0.12.10
flatbuffers: 25.12.19
clang++: /usr/bin/clang++
设备   : [CpuDevice(id=0)]

进阶各章要点¶

1. 已经删除或过时的 API（以 0.12 为准）

旧写法	在 0.12 里会怎样	现在的写法	章
`jax.experimental.host_callback`	模块已删除	`jax.debug.callback` / `jax.experimental.io_callback` / `jax.pure_callback`	5
`jax.dlpack.to_dlpack`、裸 DLPack 胶囊	已移除 / TypeError	直接 `jax.dlpack.from_dlpack(obj)`、`np.from_dlpack(x)`	5
`jax.core.Primitive`	不存在	`jax.extend.core.Primitive`	5
`xla_client.register_translation`、`CustomCall`	不可用	`mlir.register_lowering`；C++ 接入用 XLA FFI（`jax.ffi`）	5
`pure_callback(..., vectorized=True)`	参数被当作 kwargs 传给回调	`vmap_method="sequential"` / `"expand_dims"` 等	5
`device_put_replicated` / `device_put_sharded`	AttributeError	`jax.device_put(x, NamedSharding(mesh, P(...)))`	3, 6
嵌套 pmap、`jax.experimental.maps`	报错 / 已删除	二维 `jax.make_mesh` + `jax.shard_map`	3
Haiku、Flax Linen	Haiku 进入维护模式；Linen 是旧 API	Flax NNX（`.value` 已弃用，改用 `var[...]`）	6

2. 容易想错的几点

while_loop 只支持前向求导，反向模式（grad）会报错（第 1、4 章）。
不能直接对 pure_callback 求 grad，要用 custom_jvp/vjp 补上导数（第 5 章）。
闭包捕获的常量变了，jit 不会重新编译，而是继续用旧值（第 1 章）。
给整个模型套一个检查点，峰值内存一点都不降；要逐层做检查点才有用（第 6 章）。
混合精度下，对 fp32 参数求导得到的梯度是 fp32，不是 bf16（第 6 章）。
Optax 本身不管分片，优化器状态放在哪儿跟着参数走（第 6 章）。
pmean 和 psum/N 编译出来是一样的，没有谁更快（第 3 章）。

每章最后的「本章更新说明」里有完整列表，共 133 条。

后续学习路线¶

训练基础设施：Grain（数据加载）、参考实现 MaxText。
自定义内核进阶：在真实 GPU/TPU 上运行第 7 章的 Pallas 内核（见 GPU 验证指南）。
规模化理论：官方 How to Scale Your Model。