用户指南

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

调试与性能

• 计算性能分析
• 设备内存性能分析
• 调试运行时值
• GPU 性能提示
• 持久化编译缓存

接口

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

自定义操作

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

示例应用

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