趋近智
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 Documentation, JAX core developers, 2024 - JAX 自动微分功能的官方综合文档，包含计算雅可比矩阵和海森矩阵的函数。
- Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, Andreas Griewank, Andrea Walther, 2008 (SIAM) DOI: 10.1137/1.9780898717765 - 一本涵盖自动微分理论和算法的学术著作，包括前向模式、反向模式、雅可比矩阵和海森矩阵。
- Numerical Optimization, Jorge Nocedal and Stephen J. Wright, 2006 (Springer) DOI: 10.1007/978-0-387-40065-5 - 一本数值优化标准教材，解释了各种优化算法中一阶和二阶导数（雅可比矩阵、海森矩阵）的应用。