用户指南# 用户指南深入探讨了 JAX 中的特定主题,这些主题在您的 JAX 项目成熟并发展为更大规模或已部署的代码库时变得相关。 调试与性能 计算性能分析 设备内存性能分析 调试运行时值 GPU 性能提示 持久化编译缓存 接口 Pytrees 错误 提前降低和编译 导出和序列化 传输守卫 自定义操作 Pallas:一种 JAX 内核语言 外部函数接口 (FFI) 示例应用 使用 tensorflow/datasets 数据加载训练简单神经网络 使用 PyTorch 数据加载训练简单神经网络 贝叶斯推理的自动批处理