教程

• 快速入门
• 核心概念
• 即时编译
• 自动向量化
• 自动微分
• 调试入门
• 伪随机数
• 使用 PyTree
• 并行编程入门
• 有状态计算
• 使用 JIT 的控制流和逻辑运算符

高级教程

• 高级自动微分
• 外部回调
• 使用 jax.checkpoint (jax.remat) 进行梯度检查点
• JAX 内部原理：原语
• JAX 内部原理：jaxpr 语言