用户指南

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

调试和性能

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

接口

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

自定义操作

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

示例应用

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