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