关于项目#
JAX 项目由 JAX 核心团队领导。我们在开放环境中开发，并欢迎社区的开源贡献。我们经常收到来自 Google DeepMind、更广泛的 Alphabet、NVIDIA 以及其他方面的贡献。
该项目的核心是 JAX 核心库，该库专注于机器学习和数值计算的基础，并支持大规模应用。
在 开发 核心时，我们希望保持敏捷性和专注的范围，因此我们大量依赖于周围的 模块化技术堆栈。首先，我们设计了 jax 模块，使其具有 组合性 和 可扩展性，以便各种特定领域的库能够以去中心化的方式在其外部蓬勃发展。其次，我们大量依赖于模块化后端堆栈（编译器和运行时）来针对不同的加速器。无论您是 编写基于 JAX 的新领域特定库，还是希望 支持新硬件，您通常都可以通过对 JAX 核心代码库进行最少甚至零修改来贡献这些。
JAX 的许多核心贡献者都源于开源软件和研究领域，涵盖计算机科学和自然科学。我们努力持续推动机器学习和数值计算的前沿——跨越所有计算平台和加速器——并发现大规模数组编程的真谛。
开放开发#
JAX 的日常开发在 GitHub 上公开进行，使用拉取请求、问题跟踪器、讨论和 JAX 增强提案 (JEPs)。阅读和参与这些是参与其中的好方法。我们还维护 开发者笔记，涵盖 JAX 的内部设计。
JAX 核心团队决定是否接受更改和增强。目前保持简单的决策结构有助于我们在研究前沿的速度下进行开发。开放开发是我们的核心价值观，如果/当它变得有用时，我们可能会随着时间的推移适应更复杂的决策结构（例如，指定区域负责人）。
有关更多信息，请参阅 为 JAX 做贡献。
模块化堆栈#
为了支持 (a) 跨数值领域的不断增长的用户社区，以及 (b) 不断发展的硬件格局，我们大量依赖于**模块化**。
基于 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 核心。目前有几个插件正在开发中。例如，苹果公司的一个团队正在开发一个 PJRT 插件，以使 JAX 在 Apple Metal 上运行。