用户指南

用户指南深入探讨了 JAX 中随着您的 JAX 项目发展为更大或已部署的代码库而变得相关的特定主题。

调试与性能

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

接口

• Pytrees
• 错误
• 提前降低和编译
• 导出和序列化
• 类型提升语义
• 转移守卫

自定义操作

• 外部回调
• Pallas: 一个 JAX 内核语言
• Foreign function interface (FFI)

示例应用

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