JAX 入门

欢迎使用 JAX！JAX 文档包含许多有用的入门资源。快速入门 是快速上手并了解 JAX 项目概况的最简单方法。

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

• JAX 的思考方式 是对 JAX 执行模型的概念性讲解。

• 🔪 JAX - 棘手之处 🔪 列出了一些 JAX 的潜在陷阱。

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

教程

如果你准备更深入地探索 JAX，JAX 教程会提供更详细的内容

• 教程
  • 快速入门
  • 核心概念
  • 即时编译
  • 自动向量化
  • 自动微分
  • 调试入门
  • 伪随机数
  • 使用 pytrees
  • 并行编程入门
  • 有状态计算
  • 使用 JIT 的控制流和逻辑运算符
  • 高级自动微分
  • 外部回调
  • 使用 jax.checkpoint (jax.remat) 进行梯度检查点
  • JAX 内部：原语
  • JAX 内部：jaxpr 语言

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

基于 JAX 构建

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

有关跨多个领域的 JAX 相关项目的社区整理列表，请查看 Awesome JAX。

寻求帮助

如果你对 JAX 有疑问，我们很乐意解答！以下是两个可以获得解答的好去处：

• JAX GitHub 讨论区

• StackOverflow 上的 JAX