JAX 自动微分机制深度拆解 · 课程目录

JAX 初探哲学

什么是 JAX？与 NumPy 的关系 · 函数式、不可变性、转换

安装与第一行代码上手

CPU/GPU/TPU 安装 · API 对比 · 第一个 JAX 程序

核心数据结构 DeviceArray

DeviceArray 详解 · 设备内存管理 · CPU↔GPU 流转

函数式纯变换 jit

jit 编译入门 · 为什么需要 jit？Tracing 与 XLA

自动微分基础核心

什么是自动微分？前向模式 vs 反向模式

grad 函数详解梯度

grad 基本用法 · 高阶函数支持 · 与 jit 组合

jacfwd 与 jacrev 雅可比

雅可比矩阵计算 · 前向/反向选择 · 实战案例

vmap 向量化映射向量化

vmap 原理 · 与手动批处理对比 · 与 jit 协同

pmap 并行计算多GPU

pmap 数据并行 · 多 GPU 通信 · 局限性

随机数生成 PRNG

PRNGKey 设计 · 分叉与状态管理 · 与 NumPy 区别

控制流与条件 lax

lax.cond/switch/while_loop/fori_loop 陷阱

自定义 VJP 规则反向

为什么需要自定义 VJP？custom_vjp · 数值稳定 softmax

自定义 JVP 规则前向

custom_jvp 注册 · 前向模式自定义 · 与 VJP 对比

高阶梯度 & Hessian 二阶

grad 嵌套 · Hessian 计算 · 优化器中的二阶信息

调试技巧 debug

print 调试 · 断点 · tree_map 结构调试

JAX 与神经网络 Flax

Flax/Haiku 简介 · 手写线性层 · 参数管理

优化器实现 SGD/Adam

SGD、Adam 的 JAX 实现 · optax 库简介

数据加载与预处理 TFDS

TensorFlow Datasets + JAX · 数据流水线设计

训练循环实战 MNIST

完整 MNIST 分类器 · 损失函数 · 评估指标

序列化 checkpoint

参数保存/加载 · orbax checkpoint 管理

分布式计算多机多卡

多机多卡训练 · all-reduce · 分布式策略

稀疏计算稀疏

稀疏张量 · 稀疏自动微分 · 图神经网络

量化与混合精度 bfloat16

bfloat16 支持 · 混合精度训练 · 损失缩放

科学计算 PINN

微分方程求解 · 物理信息神经网络 · JAX 优势

自定义算子 Pallas

Pallas 编写自定义 CUDA 内核 · 与 JAX 集成

编译优化 XLA

XLA 深入 · HLO/MHLO · 性能分析工具

函数变换组合组合拳

grad+jit+vmap+pmap 组合 · 大规模模型训练

JAX 生态概览生态

Flax, Haiku, Optax, Orbax, Diffrax, Equinox 选型

性能调优 profiler

内存优化 · 计算图优化 · jax.profiler 分析

实战项目：GPT-2 从零实现

从零实现 GPT-2 小模型 · 训练推理 · 性能评估