关于项目#
JAX 项目由 JAX 核心团队领导。我们以开放方式进行开发，并欢迎来自社区的开源贡献。我们经常收到来自 Google DeepMind、更广泛的 Alphabet、NVIDIA 以及其他地方的贡献。
项目的核心是 JAX 核心库，它专注于大规模机器学习和数值计算的基础。
在开发核心时，我们希望保持敏捷性和专注的范围，因此我们严重依赖周围的模块化技术堆栈。首先，我们将 jax 模块设计为可组合的和可扩展的，以便各种领域特定的库可以以去中心化的方式在其外部蓬勃发展。其次，我们严重依赖模块化后端堆栈（编译器和运行时）来针对不同的加速器。无论您是使用 JAX 构建新的领域特定库，还是希望支持新硬件，您通常都可以对 JAX 核心代码库进行极少或不进行修改地贡献这些内容。
JAX 的许多核心贡献者都源于开源软件和研究领域，涵盖计算机科学和自然科学。我们致力于持续推动机器学习和数值计算的前沿发展——跨越所有计算平台和加速器——并探索大规模数组编程的真谛。
开放开发#
JAX 的日常开发在 GitHub 上公开进行，使用拉取请求、问题跟踪器、讨论和 JAX 增强提案 (JEPs)。阅读并参与这些是参与其中的好方法。我们还维护涵盖 JAX 内部设计的开发者笔记。
JAX 核心团队决定是否接受更改和增强。目前，保持简单的决策结构有助于我们以研究前沿的速度进行开发。开放开发是我们的核心价值观，如果/当它变得有用时，我们可能会随着时间的推移适应更复杂的决策结构（例如，有指定的区域所有者）。
更多信息请参阅为 JAX 贡献。
模块化堆栈#
为了实现 (a) 跨数值领域的不断增长的用户社区，以及 (b) 不断发展的硬件环境，我们严重依赖模块化。
基于 JAX 构建的库#
虽然 JAX 核心库专注于基础功能，但我们希望鼓励在 JAX 之上构建领域特定的库和工具。事实上，许多库已经围绕 JAX 出现，以提供更高级的功能和扩展。
我们如何鼓励这种去中心化的开发？我们通过几种技术选择来指导它。首先，JAX 的主要 API 侧重于基本构建块（例如数值原语、NumPy 操作、数组和变换），鼓励辅助库根据其领域需要开发实用程序。此外，JAX 还公开了一些更高级的 API，用于定制和可扩展性。库可以依赖这些 API，以便将 JAX 用作内部实现方式，更深入地与其变换（如自动微分）集成等等。
JAX 生态系统中的项目以分布式且通常开放的方式开发。它们不受 JAX 核心团队管理，尽管有时团队成员会向其贡献或与开发者保持联系。
可插拔的后端#
我们希望 JAX 能够在 CPU、GPU、TPU 以及其他新出现的硬件平台上运行。为了鼓励 JAX 在新平台上得到不受阻碍的支持，JAX 核心也在其后端强调了模块化。
为了管理硬件设备和内存，以及向这些设备进行编译，JAX 调用了开放的 XLA 编译器和 PJRT 运行时。这两个项目都独立于 JAX 核心，由 OpenXLA 管理和维护（同样，JAX 核心开发者也频繁贡献并参与讨论）。
XLA 旨在实现加速器之间的互操作性（例如，通过接收 StableHLO 作为输入），而 PJRT 则通过插件设备 API 提供可扩展性。添加对新设备的支持是通过为 XLA 实现后端降低，并实现 PJRT 定义的插件设备 API 来完成的。如果您希望为编译或支持新硬件做出贡献，我们鼓励您在 XLA 和 PJRT 层进行贡献。
这些开放系统组件允许第三方在新的加速器平台上支持 JAX，无需修改 JAX 核心。目前有几个插件正在开发中。例如，Apple 的一个团队正在开发一个 PJRT 插件，以使 JAX 在 Apple Metal 上运行。