⚡ JAX 科学计算
高效应用
📘 30章 · 从入门到生产
🔥 自动微分
⚡ JIT / GPU
🧪 实战驱动
📁 完整 30 讲
🧑🏫
01
JAX初探
JAX是什么？核心优势：自动微分、JIT编译、GPU/TPU加速概览。
02
环境准备
安装JAX (CPU/GPU)、验证、Colab配置、与NumPy第一次接触。
03
核心数据结构
DeviceArray、与ndarray异同、设备间显式传输 (.device_buffer)。
04
函数式纯变换
jax.numpy基础、随机数机制、显式PRNGKey。
05
自动微分基石
jax.grad入门、标量梯度、与Python控制流兼容。
06
高阶自动微分
jacfwd/jacrev、hessian、神经网络损失梯度实战。
07
向量化映射
jax.vmap魔力、自动批处理、与手动循环性能对比。
08
并行计算
jax.pmap入门、多GPU/TPU、数据并行与模型并行。
09
JIT编译加速
jit装饰器、静态/动态参数、约束与陷阱。
10
JIT高级技巧
缓存、三值运算、lax.cond/scan/while_loop正确用法。
11
纯函数与副作用
函数式哲学、无全局状态、in_place替代方案。
12
状态管理
lax.scan有状态计算、全局状态容器、训练循环实战。
13
自定义梯度
custom_jvp/custom_vjp、前向/反向模式、数值稳定softmax。
14
优化器实现
从零SGD、Adam、jax.tree_util参数树。
15
线性代数
jnp.linalg、LU/QR/SVD分解、求解线性系统。
16
稀疏计算
experimental.sparse、存储格式、稀疏自动微分。
17
概率编程
jax.scipy.stats、随机采样、与NumPyro对比。
18
常微分方程求解
odeint、刚性问题、洛伦兹系统模拟。
19
傅里叶变换
jnp.fft模块、信号处理、GPU加速FFT。
20
自定义Pytree
register_pytree_node、自定义类支持JAX变换、神经网络参数管理。
21
性能剖析
jax.profiler、NVIDIA Nsight、识别JIT瓶颈。
22
内存优化
梯度检查点 (checkpoint)、内存权衡、大型Transformer实战。
23
分布式计算
distributed初始化、多主机配置、all_gather/all_reduce。
24
混合精度训练
convert_element_type、bfloat16/float16、损失缩放。
25
XLA编译深入
HLO/StableHLO、优化Pass、自定义XLA操作。
26
JAX与深度学习框架
Flax/Haiku基础、与PyTorch互操作 (jax2tf, torch2jax)。
27
科学计算实战1
物理模拟：二维热传导方程、温度场可视化。
28
科学计算实战2
贝叶斯推断：MCMC采样。
29
科学计算实战3
量子计算模拟：量子门操作与态演化。
30
生产化部署
模型序列化 (safetensors)、Serving框架、监控与日志。