所有课程

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 高级编程：性能、优化与规模化

熟练运用 JAX 的进阶方法，以达成高性能机器学习。本课程内容包含 JAX 核心机制、针对 GPU 和 TPU 的性能优化策略、使用 pmap 进行分布式计算、进阶自动微分、自定义操作，以及大型模型训练技术。运用 JAX 的函数式编程和编译能力，构建复杂且高效的机器学习系统。

先修课程 Python, 机器学习入门, JAX 入门

级别:

高级

性能优化
对 JAX 代码进行性能分析和优化，以在 GPU 和 TPU 等加速器上获得最佳性能。
分布式计算
实现数据并行和模型并行，运用 pmap 及其他分布式基本操作以进行多设备执行。
进阶变换
使用复杂的控制流基本操作（如 scan, cond, while_loop），并理解它们与 JAX 变换的相互关系。
自定义自动微分规则
定义自定义的向量-雅可比乘积（VJP）和雅可比-向量乘积（JVP），用于非标准操作。
JAX 核心机制
透彻理解 JAX 的编译过程 (XLA) 及其内部表示 (jaxprs)。
大型训练
应用 JAX 进阶模式和库（如 Flax 或 Haiku），以高效训练大型神经网络。

本课程没有先修课程。

目前没有推荐的后续课程。

登录以撰写评论

分享您的反馈以帮助其他学习者。

© 2026 ApX Machine Learning用心打造