趋近智
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暂存
这部分内容有帮助吗？
- Fluent Python: Clear, Concise, and Effective Programming, Luciano Ramalho, 2022 (O'Reilly Media) - 对高级 Python 特性（包括闭包和词法作用域）进行了全面解释。
- JAX: Autograd and XLA, Accelerated, Zachary Devito, Sam Gross, Matthew Johnson, Daniel J. Seita, Peter J. Battaglia, Alex Botev, Alexey Goldin, Siddharth Goel, George Hotz, Sherjil Ozair, James S. Smith, Balaji Sreenivasan, Austin Stone, Matthew Tucker, Charles R. Harris, Kevin J. Beaumont, Bart van Merriënboer, Joshua V. Dillon, David So, Ryan P. Adams, and Sander Dieleman, 2019 2019 USENIX Conference on Systems for Machine Learning (SysML) (IEEE) DOI: 10.1109/SysML.2019.00003 - 描述了 JAX 的基础设计，详细说明了其追踪机制和与 XLA 的集成以实现加速计算。
- JAX Documentation: JIT compilation, JAX core contributors, 2023 - 解释了 JAX 的 jit 转换如何工作，包括其追踪行为以及对 Python 值（例如闭包中捕获的值）的影响。