用户指南

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

调试和性能

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

接口

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

自定义操作

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

示例应用

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