原英文教程地址: ReadTheDocs
施工中... Work in Progress...
- 启蒙
- 快速开始JAX
- 以JAX的方式思考
- JAX锋芒毕露
- 教程：JAX 101
- 加速版Numpy——JAX
- JAX的即时编译
- JAX中的自动向量化
- JAX中的高级自动微分
- JAX中的伪随机数
- 使用 pytrees
- 样例：ML模型参数
- 自定义 pytree节点
- 常见 pytree 陷阱和模式
- JAX中的并行评估
- 有状态的计算
- 参考文档
- JAX常见问题 （FAQ）
- 变换
- 异步调度
- 理解Jaxprs
- JAX中的卷积
- Pytrees
- 类型提升语义
- JAX错误
- JAX术语表
- 变更记录
- 高级JAX教程
Autodiff指导手册- 自动批处理日志密度示例
- 使用Tensorflow数据加载来训练简单的神经网络
- JAX可转换的Python函数自定义派生规则
- JAX原语如何工作
- 用JAX编写自定义Jaxpr解释器
- 使用PyTorch数据加载来训练简单的神经网络
- Python中的XLA
- 带有JAX的MAML教程
- 通过估计JAX中数据分布梯度的生成建模