用户指南

用户指南更深入地介绍了 JAX 中的特定主题，这些主题在您的 JAX 项目成熟为更大的代码库或部署的代码库时变得相关。

调试和性能

• 如何在 JAX 中思考
• 分析计算
• 分析设备内存
• 调试运行时值
• GPU 性能提示
• 持久编译缓存

接口

• Pytrees
• 错误
• 提前降低和编译
• 导出和序列化
• 类型提升语义
• 传输保护

自定义操作

• 外部回调
• Pallas: JAX 内核语言
• 外部函数接口 (FFI)

示例应用程序

• 训练简单的神经网络，使用 tensorflow/datasets 数据加载
• 训练简单的神经网络，使用 PyTorch 数据加载
• 贝叶斯推断的自动批处理