趋近智
JAX 高级编程：性能、优化与规模化
章节 1: 高级 JAX 变换与控制流
核心转换回顾：jit, grad, vmap
精通 lax.scan 处理序列操作
使用 lax.cond 进行条件执行
使用 lax.while_loop 进行循环
结合控制流与变换
高级掩码技术
了解闭包和JAX暂存
实践：实现复杂的循环逻辑
章节 2: JAX 代码的性能优化
在 CPU、GPU 和 TPU 上分析 JAX 代码
理解 JAX 计算图 (jaxpr)
XLA 编译的作用
内存布局及其对性能的影响
避免重复编译
算子合并与操作优化
异步调度
实践：优化数值计算
章节 3: JAX 分布式计算
并行处理原理介绍
JAX 中的设备管理
使用 pmap 进行单一程序多数据 (SPMD) 运算
使用 pmap 实现数据并行
集合通信原语 (psum, pmean 等)
pmap 中轴名称的处理
嵌套 pmap 与进阶分区
多主机编程简介
实践：分布式数据并行训练
章节 4: 进阶自动微分方法
前向和反向模式自动微分回顾
雅可比向量积 (JVPs) 与 jax.jvp
向量-雅可比积 (VJPs) 与 jax.vjp
高阶导数
计算完整的雅可比矩阵和海森矩阵
使用 jax.custom_vjp 的自定义微分规则
jax.custom_jvp 自定义求导规则
通过控制流原语求导
处理不可微分函数
实践：实现自定义梯度
章节 5: 互操作性与自定义操作
JAX 与 NumPy 的集成
使用 DLPack 实现零拷贝数据共享
使用 jax.experimental.host_callback 调用外部 CPU/GPU 代码
使用 jax.pure_callback 进行无副作用调用
JAX 原语简介
定义自定义原语
实现抽象求值规则
为后端（CPU/GPU/TPU）实现转换规则
制定自定义原语的求导规则
实践：整合 C++ 函数
章节 6: 大规模模型训练方法
大型模型训练中的挑战概述
JAX 生态系统库（Flax, Haiku）简介
模型参数和状态的处理
将 pmap 与训练框架结合使用
梯度累积
梯度检查点（再物化）
混合精度训练
模型并行策略
大规模优化算法
实践：实现梯度检查点
JAX 与 NumPy 的集成
这部分内容有帮助吗？
- JAX NumPy, JAX developers, current - 解释了 jax.numpy API、其与 NumPy 的相似性以及底层 jax.Array 类型。
- Device and memory considerations, JAX developers, current - 详细说明了 jax.Array 设备放置、jax.device_put、异步调度和同步，与性能考量直接相关。
- JAX: Composable transformations of Python+NumPy programs, James Bradbury, Roy Frostig, Peter Hawkins, Matthew J. Johnson, Chris Leary, Skye Wanderman-Milne, and Adam Paszke, 2018 - 描述 JAX 架构的奠基性技术报告，包括其类 NumPy API 和加速器集成机制。