JAX 入门

欢迎来到 JAX！JAX 文档包含许多有用的入门资源。快速入门：JAX 思维模式是了解 JAX 项目、其执行模型以及与 NumPy 差异的最简单切入点。

如果您刚开始探索 JAX，以下资源可能对您有所帮助

• 核心概念介绍了 JAX 的核心概念，例如变换、跟踪、jaxprs 和 pytrees。

• 🔪 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