📘 JAX·核心技巧
⚡ 30章 从入门到部署
1
JAX初探
JAX是什么？为什么比NumPy快？安装与Hello World
2
核心数据结构
DeviceArray详解，与NumPy互操作，CPU/GPU流转
3
即时编译 (JIT)
@jit装饰器，静态/动态参数，编译缓存
4
自动微分 (Grad)
grad函数，高阶导数，分支/循环处理
5
向量化映射 (vmap)
自动批处理，标量函数向量化，性能对比
6
并行计算 (pmap)
多设备并行，数据/模型并行，跨设备通信
7
随机数生成
显式状态，PRNGKey与split，与NumPy区别
8
控制流
lax.cond, while_loop, fori_loop，JIT下使用
9
纯函数与副作用
为什么要求纯函数？处理打印、累加等
10
性能分析
block_until_ready, profiler, 常见瓶颈
11
自定义梯度
@custom_vjp，自定义反向传播，不可微操作
12
JAX与NumPy互操作
numpy()/device_put()，拷贝与共享，开销
13
高级JIT
动态形状，partial固定参数，编译时/运行时
14
分布式计算
多主机多设备，pjit分区，SPMD编程
15
JAX生态概览
Flax, Haiku, Optax, RLax 等库简介
16
实战：线性回归
从零实现，grad求梯度，jit加速训练
17
实战：逻辑回归与分类
二分类，损失函数，准确率评估
18
实战：多层感知机 (MLP)
手动前向/反向，vmap批处理
19
实战：卷积神经网络 (CNN)
卷积实现，池化层，图像分类
20
实战：循环神经网络 (RNN)
简单RNN，序列处理，梯度裁剪
21
实战：变分自编码器 (VAE)
重参数化，ELBO损失，生成样本
22
实战：生成对抗网络 (GAN)
生成器/判别器，交替训练，模式崩溃
23
实战：物理模拟 (微分方程)
求解ODE，刚性问题处理
24
实战：贝叶斯推断
MCMC，随机梯度Langevin动力学
25
实战：强化学习基础
策略梯度，环境交互，vmap加速采样
26
实战：自然语言处理 (NLP)
词嵌入，注意力机制，Transformer块
27
实战：图神经网络 (GNN)
消息传递，图卷积，节点分类
28
实战：科学计算 (FFT/线性代数)
快速傅里叶变换，矩阵分解，特征值
29
实战：自定义算子与C++扩展
jax.ops新算子，调用CUDA内核
30
部署与生产化
模型序列化，TF Serving集成，调优清单