趋近智
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 Autodiff Cookbook, JAX core team, 2024 - JAX官方文档，解释其自动微分系统，包含处理不可微分操作的
stop_gradient和custom_vjp。 - Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, Andreas Griewank, Andrea Walther, 2008 (Society for Industrial and Applied Mathematics (SIAM)) DOI: 10.1137/1.9781611975571 - 关于算法微分的学术著作，覆盖JVP/VJP的理论基础以及与不可微分性相关的问题。
- Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation, Yoshua Bengio, Nicholas Léonard, and Aaron Courville, 2013 arXiv preprint arXiv:1308.3432 DOI: 10.48550/arXiv.1308.3432 - 提出了直通估计器等方法，用于通过离散或随机单元反向传播梯度，与自定义VJP示例相关。
- Deep Learning, Ian Goodfellow, Yoshua Bengio, and Aaron Courville, 2016 (MIT Press) - 提供了深度学习背景下自动微分的介绍，讨论了梯度计算以及不可微分激活函数或操作可能产生的问题。