趋近智
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 与训练框架结合使用
梯度累积
梯度检查点（再物化）
混合精度训练
模型并行策略
大规模优化算法
实践：实现梯度检查点
雅可比向量积 (JVPs) 与 jax.jvp
这部分内容有帮助吗？
- Automatic Differentiation of Algorithms: From Simulation to Optimization, Andreas Griewank, Andrea Walther, 2008 (Springer) DOI: 10.1007/978-3-540-68937-5 - 一本关于自动微分理论与实践的综合性教科书，详细介绍了前向模式自动微分和雅可比向量积。
jax.jvp, JAX core developers, 2024 -jax.jvp函数的官方文档，提供了API细节、使用示例以及它在JAX框架中的作用。- JAX: Composable transformations of Python+NumPy programs, James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Skye Lew, George H. Nash, Nicolas Nieto, Sam Oden, Mark Omernick, Jonathan Renjie So, Matt Von-Rose and Sasha Rush, 2018 - 介绍JAX库的原始参考文献，概述了其自动微分和其他程序转换的设计原则。
- The mechanics of automatic differentiation, Atilim Baydin, Barak A. Pearlmutter and Alexey Andreyevich Radul, 2018 Journal of Machine Learning Research, Vol. 18 (Journal of Machine Learning Research) DOI: 10.5555/3278500.3278502 - 一篇关于自动微分的综述和教程，涵盖了其数学基础、各种模式（包括前向模式）和实际实现。