JAX 入门

欢迎来到 JAX！JAX 文档包含许多有用的入门资源。 快速入门：如何在 JAX 中思考 是最容易上手并了解 JAX 项目、其执行模型以及与 NumPy 差异的地方。

如果您刚开始探索 JAX，您可能还会发现以下资源很有帮助：

• 关键概念 介绍了 JAX 的关键概念，例如转换、跟踪、jaxprs 和 pytrees。

• 🔪 JAX - 锋利的方面 🔪 列出了一些 JAX 的“锋利角落”。

• 常见问题解答 (FAQ) 回答了一些常见的 JAX 问题。

JAX 101

如果您已准备好更深入地探索 JAX，JAX 101 教程将提供更详细的介绍。

• JAX 101
  • 即时编译
  • 自动向量化
  • 自动微分
  • 使用 Pytree
  • 伪随机数
  • 并行编程简介
  • JIT 下的控制流和逻辑运算符
  • 追踪
  • 有状态计算

如果您更喜欢视频介绍，这里有一个由 JAX 贡献者 Jake VanderPlas 制作的视频。

基于 JAX 的构建

JAX 为社区开发的许多工具提供了核心数值计算基元。例如，如果您有兴趣使用 JAX 训练神经网络，两个支持良好的选项是 Flax 和 Haiku。

有关社区在广泛领域内精选的 JAX 相关项目列表，请查看 Awesome JAX。

寻求帮助

如果您有关于 JAX 的问题，我们很乐意为您解答！有两个很好的地方可以得到您问题的答案：

• JAX GitHub 讨论

• StackOverflow 上的 JAX