JAX 教程草案

注意

以下教程仍在开发中；目前，请参阅旧教程内容，包括 JAX 入门、用户指南 和现已删除的 JAX 101 教程。

JAX 101

大部分已完成，位于 教程！

• 快速入门
• 关键概念
• 即时编译
• 自动矢量化
• 自动微分
• 调试入门
• 伪随机数
• 使用 pytrees
• 并行编程入门
• 有状态计算
• 示例：编写简单的神经网络

JAX 201

• 并行计算
• 高级自动微分
• 使用 jax.checkpoint (jax.remat) 的梯度检查点
• 高级调试
• 外部回调
• 分析和性能

JAX 301

• JAX 内部：原语
• JAX 内部：jaxpr 语言
• 高级编译