用户指南

用户指南是对 JAX 中特定主题的深入探讨，当您的 JAX 项目成熟为更大或已部署的代码库时，这些主题将变得相关。

调试和性能

• 如何以 JAX 的方式思考
• 性能分析计算
• 性能分析设备内存
• 调试运行时值
• GPU 性能技巧
• 持久编译缓存

接口

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

自定义操作

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

示例应用

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