关键词索引

概念索引

每个概念列出相关度最高的小节（标题命中优先），最多 8 处。

追踪与 jaxpr

• 进阶 02_JAX代码性能优化 › 2.2 理解 JAX 计算图 (jaxpr)
• 入门 02_JIT编译 › 2.4 Python 控制流与 jit
• 入门 02_JIT编译 › 2.3 JIT 工作原理：追踪与编译
• 进阶 08_调试工具链 › 4.6 UnexpectedTracerError：tracer 泄漏
• 入门 04_自动向量化vmap › 4.2 介绍 jax.vmap
• 入门 03_自动微分grad › 3.3 自动微分的工作方式：反向模式
• 进阶 01_高级变换与控制流 › 1.7 了解闭包和 JAX 暂存
• 入门 02_JIT编译 › 2.6 jit 的常见问题

重新编译与编译缓存

• 进阶 02_JAX代码性能优化 › 2.5 避免重复编译
• 入门 02_JIT编译 › 2.3 JIT 工作原理：追踪与编译
• 入门 02_JIT编译 › 2.5 静态值与跟踪值
• 入门 07_实战专题 › 1.1 问题动机
• 入门 07_实战专题 › ★.4 训练：编译预算 = 1
• 入门 07_实战专题 › 练习 1：给服务函数写编译预算测试
• 入门 02_JIT编译 › 练习 3：重新编译侦探
• 入门 07_实战专题 › 1.6 编译预算工具

静态参数与弱类型

• 入门 02_JIT编译 › 2.5 静态值与跟踪值
• 入门 07_实战专题 › 1.5 ✅ 治理手段
• 进阶 02_JAX代码性能优化 › 2.5 避免重复编译
• 入门 02_JIT编译 › 2.3 JIT 工作原理：追踪与编译
• 入门 02_JIT编译 › 2.7 动手实践：应用 jit
• 入门 06_状态管理 › 6.8 实践：实现有状态函数
• 入门 06_状态管理 › 6.5 示例：有状态计数器
• 入门 06_状态管理 › 6.6 例子：简单的优化器状态

控制流（cond / scan / while）

• 入门 02_JIT编译 › 2.4 Python 控制流与 jit
• 进阶 01_高级变换与控制流 › 1.4 使用 lax.while_loop 进行循环
• 进阶 04_进阶自动微分 › 4.8 通过控制流原语求导
• 入门 03_自动微分grad › 3.7 求导与控制流
• 入门 04_自动向量化vmap › 4.7 vmap 的性能考量
• 进阶 01_高级变换与控制流 › 1.3 使用 lax.cond 进行条件执行
• 进阶 02_JAX代码性能优化 › 2.2 理解 JAX 计算图 (jaxpr)
• 入门 06_状态管理 › 🏋️ 补充练习（含参考答案）

掩码与变长序列

• 进阶 01_高级变换与控制流 › 1.6 高级掩码技术
• 入门 07_实战专题 › 1.5 ✅ 治理手段
• 入门 02_JIT编译 › 2.6 jit 的常见问题
• 进阶 01_高级变换与控制流 › 1.8 实践：实现复杂的循环逻辑
• 进阶 07_Pallas自定义内核 › 4.3 带掩码的边界处理
• 入门 07_实战专题 › 练习 6：冻结部分层 + 捐赠状态的训练 step
• 入门 07_实战专题 › 2.3 logsumexp / softmax / log_softmax 的稳定写法
• 进阶 08_调试工具链 › 4.5 布尔掩码索引、形状和切片里的 tracer

NaN 与数值稳定

• 入门 03_自动微分grad › 3.8 局限性与注意事项
• 入门 07_实战专题 › 2.4 用 jax_debug_nans 定位 NaN（0.12 实测行为）
• 进阶 08_调试工具链 › 3.1 jax_debug_nans / jax_debug_infs
• 入门 07_实战专题 › 2.1 问题动机
• 入门 07_实战专题 › 2.3 logsumexp / softmax / log_softmax 的稳定写法
• 进阶 11_科学计算与可微分仿真 › 7.4 势垒函数的数值稳定写法：双 where
• 进阶 08_调试工具链 › 9.1 症状 1：NaN —— 定位
• 进阶 08_调试工具链 › 练习 4：自动定位第一个 NaN

自动微分模式（jvp / vjp）

• 入门 03_自动微分grad › 3.3 自动微分的工作方式：反向模式
• 进阶 04_进阶自动微分 › 4.7 jax.custom_jvp 自定义求导规则
• 进阶 04_进阶自动微分 › 4.3 向量-雅可比积 (VJPs) 与 jax.vjp
• 入门 07_实战专题 › 3.3 计时实验：扫描输入/输出维度
• 进阶 04_进阶自动微分 › 4.1 前向和反向模式自动微分回顾
• 进阶 07_Pallas自定义内核 › 5.3 grad：默认不可微，用 custom_vjp + Pallas 反向 kernel
• 进阶 05_互操作性与自定义操作 › 5.9 制定自定义原语的求导规则
• 进阶 09_导出与部署 › §2.5 导出可求导的函数：vjp_order

Jacobian / Hessian

• 进阶 04_进阶自动微分 › 4.5 计算完整的雅可比矩阵和海森矩阵
• 入门 03_自动微分grad › 3.3 自动微分的工作方式：反向模式
• 入门 07_实战专题 › 3.3 计时实验：扫描输入/输出维度
• 入门 03_自动微分grad › 3.5 高阶导数（grad 的 grad）
• 进阶 11_科学计算与可微分仿真 › 练习 4：用 jacfwd 验证角点 Jacobian
• 入门 07_实战专题 › 3.5 Hessian-向量积（HVP）的三种实现
• 入门 03_自动微分grad › 补充练习
• 进阶 04_进阶自动微分 › 4.2 雅可比向量积 (JVPs) 与 jax.jvp

自定义求导（custom_jvp / custom_vjp）

• 入门 03_自动微分grad › 3.8 局限性与注意事项
• 进阶 04_进阶自动微分 › 4.6 使用 jax.custom_vjp 的自定义微分规则
• 进阶 04_进阶自动微分 › 4.7 jax.custom_jvp 自定义求导规则
• 进阶 04_进阶自动微分 › 4.9 处理不可微分函数
• 进阶 07_Pallas自定义内核 › 5.3 grad：默认不可微，用 custom_vjp + Pallas 反向 kernel
• 进阶 04_进阶自动微分 › 4.10 实践：实现自定义梯度
• 进阶 04_进阶自动微分 › 补充练习 D：在 0 处安全的向量范数（custom_jvp）
• 入门 07_实战专题 › 2.5 用 custom_jvp 修复梯度

逐样本梯度

• 入门 04_自动向量化vmap › 4.6 结合 vmap 与 jit 和 grad
• 入门 04_自动向量化vmap › 4.8 动手实践：函数向量化
• 入门 07_实战专题 › 4.2 vmap(grad) 与循环实现对拍
• 入门 07_实战专题 › 4.5 DP-SGD：裁剪 + 加噪的完整实现
• 入门 07_实战专题 › ★.6 逐样本梯度：找出 top-5「难样本」
• 入门 07_实战专题 › 练习 4：用逐样本梯度找出标签错误的样本
• 入门 07_实战专题 › 4.4 内存：$B \times P$，以及分块计算
• 入门 07_实战专题 › 学习目标

向量化 vmap

• 入门 04_自动向量化vmap › 4.7 vmap 的性能考量
• 入门 04_自动向量化vmap › 4.8 动手实践：函数向量化
• 入门 04_自动向量化vmap › 4.6 结合 vmap 与 jit 和 grad
• 进阶 01_高级变换与控制流 › 1.1 核心转换回顾：jit, grad, vmap
• 入门 04_自动向量化vmap › 4.2 介绍 jax.vmap
• 入门 04_自动向量化vmap › 4.5 嵌套 vmap
• 入门 04_自动向量化vmap › 4.1 向量化的原理
• 入门 04_自动向量化vmap › 4.3 对特定参数进行映射（in_axes，out_axes）

多设备与分片

• 进阶 06_大规模模型训练 › 6.8 模型并行策略
• 进阶 03_JAX分布式计算 › 3.7 嵌套 pmap 与进阶分区
• 进阶 06_大规模模型训练 › 6.4 将 pmap 与训练框架结合使用
• 入门 05_多设备并行pmap › 5.4 设备网格与轴名称
• 入门 05_多设备并行pmap › ⚠️ 专题：shard_map 中梯度变成 N 倍（pmap → shard_map 迁移头号陷阱）
• 进阶 09_导出与部署 › §2.7 含分片、shard_map、pmap 的函数怎么导出
• 入门 05_多设备并行pmap › 5.2 介绍 jax.pmap
• 进阶 03_JAX分布式计算 › 动手练习

pmap 与集合通信

• 入门 05_多设备并行pmap › 5.5 集体操作（lax.psum、lax.pmean 等）
• 进阶 03_JAX分布式计算 › 3.5 集合通信原语 (psum, pmean 等)
• 入门 05_多设备并行pmap › 5.2 介绍 jax.pmap
• 入门 05_多设备并行pmap › 5.7 调试 pmap 化的函数
• 进阶 06_大规模模型训练 › 6.4 将 pmap 与训练框架结合使用
• 入门 05_多设备并行pmap › 5.6 将 pmap 与其他变换结合使用
• 入门 05_多设备并行pmap › ⚠️ 专题：shard_map 中梯度变成 N 倍（pmap → shard_map 迁移头号陷阱）
• 入门 05_多设备并行pmap › 5.8 动手实践：并行计算

模型并行（FSDP / 张量并行）

• 进阶 06_大规模模型训练 › 6.8 模型并行策略
• 进阶 03_JAX分布式计算 › 3.1 并行处理原理介绍
• 进阶 03_JAX分布式计算 › 3.7 嵌套 pmap 与进阶分区
• 进阶 03_JAX分布式计算 › 动手练习
• 进阶 03_JAX分布式计算 › 3.5 集合通信原语 (psum, pmean 等)
• 入门 05_多设备并行pmap › 5.8 动手实践：并行计算
• 进阶 06_大规模模型训练 › 6.1 大型模型训练中的挑战概述
• 进阶 06_大规模模型训练 › 6.9 大规模优化算法

多主机

• 进阶 03_JAX分布式计算 › 3.8 多主机编程简介
• 进阶 10_随机数与分片 › 5.6 多进程：为每个进程派生 key
• 进阶 03_JAX分布式计算 › 3.2 JAX 中的设备管理
• 进阶 06_大规模模型训练 › 🧭 后续学习路线
• 入门 05_多设备并行pmap › 5.8 动手实践：并行计算
• 进阶 09_导出与部署 › §2.7 含分片、shard_map、pmap 的函数怎么导出
• 进阶 10_随机数与分片 › 练习

PyTree

• 入门 06_状态管理 › 6.4 使用 PyTree 管理分层状态
• 入门 07_实战专题 › 7.2 旧 → 新迁移表（0.12 实测）
• 入门 06_状态管理 › 6.6 例子：简单的优化器状态
• 入门 06_状态管理 › 6.7 将状态管理与变换结合
• 进阶 01_高级变换与控制流 › 1.2 精通 lax.scan 处理序列操作
• 入门 04_自动向量化vmap › 4.3 对特定参数进行映射（in_axes，out_axes）
• 进阶 06_大规模模型训练 › 6.3 模型参数和状态的处理
• 入门 05_多设备并行pmap › 5.3 将数据映射到设备（in_axes、out_axes）

随机数与 key

• 入门 01_JAX简介 › 1.5 使用 JAX 数组
• 进阶 10_随机数与分片 › 3　PRNG 算法：threefry、philox、rbg、unsafe_rbg
• 入门 04_自动向量化vmap › 4.6 结合 vmap 与 jit 和 grad
• 进阶 10_随机数与分片 › 练习
• 进阶 10_随机数与分片 › 4.1 默认值
• 入门 01_JAX简介 › 1.7 动手练习：基本数组操作
• 入门 06_状态管理 › 6.5 示例：有状态计数器
• 进阶 10_随机数与分片 › 4.5 「随机数不保证跨版本一致」，以及第 1 章的 10.81

状态管理与优化器

• 进阶 06_大规模模型训练 › 6.9 大规模优化算法
• 入门 06_状态管理 › 6.6 例子：简单的优化器状态
• 进阶 06_大规模模型训练 › 6.3 模型参数和状态的处理
• 入门 06_状态管理 › 🔍 补充：生态库如何管理状态 —— Optax 与 Flax
• 入门 06_状态管理 › 6.7 将状态管理与变换结合
• 入门 07_实战专题 › 5.2 原理：GradientTransformation 与各个组件
• 进阶 06_大规模模型训练 › 6.5 梯度累积
• 进阶 11_科学计算与可微分仿真 › 5.2 用 Optax 的 L-BFGS 优化

梯度检查点（remat）

• 进阶 06_大规模模型训练 › 6.6 梯度检查点（再物化）
• 进阶 06_大规模模型训练 › 6.10 实践：实现梯度检查点
• 进阶 06_大规模模型训练 › 6.3 模型参数和状态的处理
• 进阶 04_进阶自动微分 › 4.8 通过控制流原语求导
• 进阶 11_科学计算与可微分仿真 › 6.3 长时间积分的内存：checkpoint 与 $O(\sqrt{T})$ 分段
• 进阶 01_高级变换与控制流 › 1.5 结合控制流与变换
• 进阶 06_大规模模型训练 › 🧭 后续学习路线
• 进阶 09_导出与部署 › §3.4 与 Flax NNX 配合：保存、恢复、继续训练

梯度累积

• 进阶 06_大规模模型训练 › 6.5 梯度累积
• 进阶 03_JAX分布式计算 › 3.1 并行处理原理介绍
• 进阶 06_大规模模型训练 › 6.8 模型并行策略
• 进阶 06_大规模模型训练 › ✍️ 动手练习
• 入门 06_状态管理 › 6.6 例子：简单的优化器状态
• 进阶 06_大规模模型训练 › 6.3 模型参数和状态的处理
• 进阶 06_大规模模型训练 › 6.9 大规模优化算法

混合精度

• 进阶 06_大规模模型训练 › 6.7 混合精度训练
• 进阶 07_Pallas自定义内核 › 4.1′ 用 scratch 累加器：bf16 输入、float32 累加
• 进阶 06_大规模模型训练 › 6.1 大型模型训练中的挑战概述
• 进阶 06_大规模模型训练 › ✍️ 动手练习
• 进阶 03_JAX分布式计算 › 3.1 并行处理原理介绍
• 进阶 06_大规模模型训练 › 6.9 大规模优化算法
• 进阶 07_Pallas自定义内核 › 6.2 compiler_params：后端专属参数
• 进阶 07_Pallas自定义内核 › 6.3 调优思路：块大小、算术强度与 VMEM 预算

内存与 donate

• 入门 06_状态管理 › 6.6 例子：简单的优化器状态
• 入门 04_自动向量化vmap › 4.7 vmap 的性能考量
• 入门 07_实战专题 › 6.2 原理：输入-输出别名（aliasing）
• 进阶 08_调试工具链 › 7.1 memory_analysis()：运行之前就知道要多少内存
• 进阶 02_JAX代码性能优化 › 2.4 内存布局及其对性能的影响
• 入门 07_实战专题 › 6.5 ❌ 常见错误与实测
• 进阶 06_大规模模型训练 › 6.1 大型模型训练中的挑战概述
• 进阶 06_大规模模型训练 › 6.6 梯度检查点（再物化）

性能分析与融合

• 进阶 02_JAX代码性能优化 › 2.3 XLA 编译的作用
• 进阶 02_JAX代码性能优化 › 2.6 算子合并与操作优化
• 进阶 02_JAX代码性能优化 › 2.1 在 CPU、GPU 和 TPU 上分析 JAX 代码
• 入门 02_JIT编译 › 2.1 速度提升：为何需要编译？
• 进阶 07_Pallas自定义内核 › 1.3 三条「自定义操作」路线的对照
• 进阶 09_导出与部署 › §4.1 StableHLO 文本 / 字节码 → C++ 运行时（IREE、PJRT）
• 进阶 02_JAX代码性能优化 › 2.8 实践：优化数值计算——两两距离矩阵
• 进阶 02_JAX代码性能优化 › 2.7 异步调度

调试工具

• 进阶 08_调试工具链 › 3.2 jax.experimental.checkify
• 入门 02_JIT编译 › 2.6 jit 的常见问题
• 入门 05_多设备并行pmap › 5.7 调试 pmap 化的函数
• 进阶 08_调试工具链 › 1.4 jax.disable_jit()：既是调试手段，也是二分工具
• 进阶 08_调试工具链 › 1.1 jax.debug.print
• 进阶 08_调试工具链 › 1.3 jax.debug.breakpoint
• 入门 04_自动向量化vmap › 4.2 介绍 jax.vmap
• 进阶 05_互操作性与自定义操作 › 5.3 使用 jax.experimental.host_callback 调用外部 CPU/GPU 代码

回调与外部代码

• 进阶 05_互操作性与自定义操作 › 5.3 使用 jax.experimental.host_callback 调用外部 CPU/GPU 代码
• 进阶 05_互操作性与自定义操作 › 5.4 使用 jax.pure_callback 进行无副作用调用
• 进阶 05_互操作性与自定义操作 › 5.10 实践：整合 C++ 函数
• 进阶 08_调试工具链 › 1.2 jax.debug.callback
• 进阶 09_导出与部署 › §4.4 用 pure_callback 把模型包成服务函数：取舍
• 进阶 05_互操作性与自定义操作 › 与入门课的关联
• 进阶 08_调试工具链 › 1.1 jax.debug.print
• 进阶 04_进阶自动微分 › 4.7 jax.custom_jvp 自定义求导规则

自定义原语与 FFI

• 进阶 05_互操作性与自定义操作 › 5.10 实践：整合 C++ 函数
• 进阶 05_互操作性与自定义操作 › 5.5 JAX 原语简介
• 进阶 05_互操作性与自定义操作 › 5.6 定义自定义原语
• 进阶 05_互操作性与自定义操作 › 5.9 制定自定义原语的求导规则
• 进阶 02_JAX代码性能优化 › 2.2 理解 JAX 计算图 (jaxpr)
• 进阶 05_互操作性与自定义操作 › 5.8 为后端实现转换规则
• 进阶 04_进阶自动微分 › 4.8 通过控制流原语求导
• 进阶 05_互操作性与自定义操作 › 与入门课的关联

DLPack 与零拷贝

• 进阶 05_互操作性与自定义操作 › 5.2 使用 DLPack 实现零拷贝数据共享
• 进阶 05_互操作性与自定义操作 › 与入门课的关联
• 进阶 05_互操作性与自定义操作 › 5.1 JAX 与 NumPy 的集成
• 入门 07_实战专题 › 6.5 ❌ 常见错误与实测
• 入门 07_实战专题 › 💡 实战规则
• 进阶 02_JAX代码性能优化 › 2.4 内存布局及其对性能的影响

Pallas 内核

• 进阶 07_Pallas自定义内核 › 1.3 三条「自定义操作」路线的对照
• 进阶 07_Pallas自定义内核 › 5.3 grad：默认不可微，用 custom_vjp + Pallas 反向 kernel
• 进阶 07_Pallas自定义内核 › 与已学章节的关系
• 进阶 07_Pallas自定义内核 › §0 探针：确认 0.12 的 Pallas API
• 进阶 07_Pallas自定义内核 › 2.2 grid、BlockSpec 与 pl.program_id
• 进阶 07_Pallas自定义内核 › 6.3 调优思路：块大小、算术强度与 VMEM 预算
• 进阶 09_导出与部署 › §2.6 版本兼容性：三层版本号
• 进阶 07_Pallas自定义内核 › 2.1 kernel、Ref 与 out_shape

导出与部署

• 进阶 09_导出与部署 › §5 综合实践：训练 → checkpoint → 导出 → 全新子进程独立推理
• 进阶 09_导出与部署 › §2.3 形状多态：一次导出，适配任意 batch
• 进阶 09_导出与部署 › §2.6 版本兼容性：三层版本号
• 进阶 09_导出与部署 › §2.1 基本用法：Exported 对象里有什么
• 进阶 09_导出与部署 › §2.7 含分片、shard_map、pmap 的函数怎么导出
• 进阶 09_导出与部署 › §2.5 导出可求导的函数：vjp_order
• 进阶 09_导出与部署 › §4.4 用 pure_callback 把模型包成服务函数：取舍
• 进阶 09_导出与部署 › §2.2 序列化到文件，在全新子进程中运行

Flax NNX

• 进阶 06_大规模模型训练 › 6.2 JAX 生态系统库（Flax, Haiku）简介
• 入门 06_状态管理 › 🔍 补充：生态库如何管理状态 —— Optax 与 Flax
• 进阶 06_大规模模型训练 › 6.3 模型参数和状态的处理
• 进阶 09_导出与部署 › §3.4 与 Flax NNX 配合：保存、恢复、继续训练
• 进阶 06_大规模模型训练 › 6.4 将 pmap 与训练框架结合使用
• 进阶 10_随机数与分片 › 5.7 dropout 的 key 管理，以及 Flax nnx.Rngs 的设计
• 进阶 06_大规模模型训练 › 6.6 梯度检查点（再物化）
• 进阶 06_大规模模型训练 › 6.8 模型并行策略

科学计算与 PDE

• 入门 05_多设备并行pmap › 5.4 设备网格与轴名称
• 进阶 11_科学计算与可微分仿真 › 7.3 两种平滑能量
• 进阶 11_科学计算与可微分仿真 › 3.1 用三元组组装五点差分矩阵，与 SciPy CSR 对拍
• 进阶 11_科学计算与可微分仿真 › 3.3 BCOO + CG，以及稀疏直接求解器的现状
• 进阶 11_科学计算与可微分仿真 › 3.2 sparse.sparsify：让「稠密写法」的函数接受稀疏矩阵
• 进阶 03_JAX分布式计算 › 3.7 嵌套 pmap 与进阶分区
• 进阶 11_科学计算与可微分仿真 › 2.2 无矩阵（matrix-free）算子
• 进阶 11_科学计算与可微分仿真 › 7.5 Winslow 优化：从翻转的 TFI 网格出发

本环境已知问题

• 进阶 02_JAX代码性能优化 › 2.3 XLA 编译的作用
• 进阶 11_科学计算与可微分仿真 › 9.1 先实测：闭包常量 bug 在 float64 下是否存在
• 入门 03_自动微分grad › 3.3 自动微分的工作方式：反向模式
• 入门 02_JIT编译 › 2.6 jit 的常见问题
• 入门 07_实战专题 › 0　公共设置
• 入门 07_实战专题 › 2.4 用 jax_debug_nans 定位 NaN（0.12 实测行为）
• 进阶 08_调试工具链 › 3.1 jax_debug_nans / jax_debug_infs
• 入门 04_自动向量化vmap › 4.7 vmap 的性能考量

API 索引

按命名空间分组；每个 API 列出讲到它的 notebook（最多 6 个，链接到首次出现的小节）。

jax.*

lax.*

jnp.*

nnx.*

optax.*

ocp.*

pl.*

checkify.*

sparse.*