所有课程

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: Custom Operation Cookbook, The JAX Developers, 2024 (The JAX Project) - JAX 自定义操作的官方指南，详细介绍了抽象评估、下沉和微分的实现步骤。
JAX: Autograd and XLA, Together, The JAX team, 2018 (Google Research) - JAX 架构的官方概览，解释了其设计以及与 XLA 和自动微分的集成。
XLA: Optimizing Compilers for Machine Learning, Google Developers, 2024 (Google Developers) - XLA 的文档，这是 JAX 使用的编译器，有助于理解下沉细节和指令集。

© 2025 ApX Machine Learning用心打造