JAX 入门

欢迎来到 JAX！JAX 文档包含大量有用的资源，可以帮助您入门。 快速入门 是开始并了解 JAX 项目概述的最简单方法。

如果您习惯于编写 NumPy 代码，并且开始探索 JAX，您可能会发现以下资源很有用

• 如何在 JAX 中思考 是对 JAX 执行模型的概念性逐步介绍。

• 🔪 JAX - 尖锐的部分 🔪 列出了 JAX 的一些尖锐角落。

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

教程

如果您已准备好更深入地探索 JAX，JAX 教程将提供更多详细信息

• 教程
  • 快速入门
  • 关键概念
  • 即时编译
  • 自动矢量化
  • 自动微分
  • 调试简介
  • 伪随机数
  • 使用 pytrees
  • 并行编程简介
  • 有状态计算

如果您更喜欢视频介绍，这里有一个来自 JAX 贡献者 Jake VanderPlas 的视频

建立在 JAX 之上

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

要查看跨广泛领域与 JAX 相关的社区策划项目列表，请查看 Awesome JAX。

寻求帮助

如果您对 JAX 有任何疑问，我们很乐意为您解答！以下两个地方是您获得答案的好去处：

• JAX GitHub 讨论

• StackOverflow 上的 JAX