📘 JAX·函数式编程
从零到精通
🎯 30章 · 完整目录
1
JAX初探
JAX是什么？对比NumPy/TensorFlow/PyTorch
➜
2
环境搭建
CPU/GPU/TPU · jaxlib/optax/flax
➜
3
DeviceArray详解
与ndarray互操作 · 设备内存管理
➜
4
函数式编程基础
纯函数 · 副作用 · JAX设计哲学
➜
5
JIT编译入门
@jit · 即时编译原理 · 性能差异
➜
6
JIT进阶
static_argnums · 形状变化 · 缓存
➜
7
自动微分 (一)
grad · 标量梯度 · 高阶梯度
➜
8
自动微分 (二)
jacfwd/jacrev · Hessian矩阵
➜
9
向量化映射 vmap
自动批处理 · 性能对比
➜
10
并行计算 pmap
多设备数据并行 · 模型并行
➜
11
随机数生成
PRNGKey · split/fold · 显式状态
➜
12
控制流 (一)
lax.cond/while_loop/fori_loop
➜
13
控制流 (二)
scan · 累积计算 · RNN应用
➜
14
自定义梯度
@custom_vjp · 前向/反向 · 数值稳定
➜
15
Pytree数据结构
tree_map/tree_reduce · 参数管理
➜
16
状态管理
Stateful vs 函数式 · kwargs传递状态
➜
17
Flax入门
Module定义 · 参数初始化
➜
18
Flax进阶
optax · 训练循环 · checkpoint
➜
19
线性回归实战
梯度下降 · 损失函数可视化
➜
20
逻辑回归与分类
二分类 · 交叉熵 · 决策边界
➜
21
多层感知机 MLP
Flax构建 · 激活函数 · 正则化
➜
22
卷积神经网络 CNN
卷积/池化 · Flax Conv · 图像分类
➜
23
循环神经网络 RNN
scan实现 · LSTM/GRU · 序列预测
➜
24
Transformer入门
自注意力 · 多头注意力 · 位置编码
➜
25
生成对抗网络 GAN
生成器/判别器 · 对抗训练
➜
26
性能优化
GPU调优 · XLA · 内存复用
➜
27
调试与Profiling
jax.debug · 断点 · nvprof
➜
28
分布式训练
多GPU · 梯度累积 · 混合精度
➜
29
模型部署
SavedModel · TF Serving · ONNX
➜
30
前沿与生态
科学计算 · DeepMind/Google Research
➜
✨ 30章 · 函数式 JAX 从零到精通