指南 4：使用 JAX 进行研究项目

本指南总结了使用 JAX 进行研究项目时有用的一些提示、技巧和实践。在我看来，与 PyTorch 相比，JAX 缺少的一个关键方面是像PyTorch Lightning这样的框架可以大大减少代码开销，同时仍然足够灵活以支持几乎任何模型/任务。尽管存在用于某些常见任务的此类库，例如 trax或风景（基于注意力的简历），到目前为止，我还没有遇到过对我的研究足够灵活的简历。因此，在本指南中，我们构建了一个更简单版本的 PyTorch Lightning 训练器，它总结了几乎所有模型所需的所有训练、日志记录等行为，并且允许使用比从头开始少得多的行数来训练各种模型。此外，我们还实现了一些简单的示例来展示可