JAX 入门

欢迎来到 JAX！JAX 文档包含许多有用的入门资源。快速入门 是最容易上手并了解 JAX 项目概况的地方。

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

• 如何用 JAX 思考 是对 JAX 执行模型的概念性演练。

• 🔪 JAX - 锋利之处 🔪 列出了一些 JAX 的尖锐角落。

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

教程

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

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

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

构建于 JAX 之上

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

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

查找帮助

如果您对 JAX 有疑问，我们很乐意回答！以下是两个可以获得问题解答的好地方

• JAX GitHub 讨论区

• StackOverflow 上的 JAX