Jax 和 Jaxlib 版本控制#

为什么 jax 和 jaxlib 是单独的软件包？#

我们将 JAX 发布为两个单独的 Python Wheel，即 jax（一个纯 Python Wheel）和 jaxlib（一个主要由 C++ 组成的 Wheel，其中包含如下库：

  • XLA,

  • XLA 使用的 LLVM 片段，

  • MLIR 基础设施，例如 StableHLO Python 绑定。

  • 用于快速 JIT 和 PyTree 操作的 JAX 特定 C++ 库。

我们发布单独的 jax 和 jaxlib 软件包，因为这样可以轻松处理 JAX 的 Python 部分，而无需构建 C++ 代码，甚至无需安装 C++ 工具链。jaxlib 是一个大型库，对许多用户来说不容易构建，但 JAX 的大多数更改只涉及 Python 代码。通过允许 Python 部分独立于 C++ 部分进行更新，我们提高了 Python 更改的开发速度。

此外，jaxlib 的构建成本很高，但我们希望能够在 CPU 资源不多的环境中迭代和运行 JAX 测试，例如在 Github Actions 或笔记本电脑上。我们的许多 CI 构建只是使用预构建的 jaxlib，而无需在每个 PR 上重建 JAX 的 C++ 部分。

正如我们将看到的，单独分发 jax 和 jaxlib 会带来成本，因为它要求对 jaxlib 的更改保持向后兼容的 API。但是，我们认为，总的来说，即使以使 C++ 更改稍微困难为代价，使 Python 更改变得容易也是更可取的。

jax 和 jaxlib 如何进行版本控制？#

摘要：jax 和 jaxlib 在 JAX 源代码树中共享相同的版本号，但作为单独的 Python 软件包发布。安装后，jax 软件包版本必须大于或等于 jaxlib 的版本，并且 jaxlib 的版本必须大于或等于 jax 指定的最低 jaxlib 版本。

jax 和 jaxlib 版本都编号为 x.y.z，其中 x 是主版本，y 是次版本，z 是可选的补丁版本。版本号必须遵循 PEP 440。版本号比较是对整数元组的字典序比较。

每个 jax 版本都有一个关联的最低 jaxlib 版本 mx.my.mz。jax 版本 x.y.z 的最低 jaxlib 版本必须不大于 x.y.z。

为了使 jax 版本 x.y.z 和 jaxlib 版本 lx.ly.lz 兼容，必须满足以下条件

  • jaxlib 版本 (lx.ly.lz) 必须大于或等于最低 jaxlib 版本 (mx.my.mz)。

  • jax 版本 (x.y.z) 必须大于或等于 jaxlib 版本 (lx.ly.lz)。

这些约束意味着以下发布规则

  • jax 可以随时单独发布，而无需更新 jaxlib。

  • 如果发布新的 jaxlib，则必须同时发布 jax 版本。

这些 版本约束 目前在导入时由 jax 检查，而不是表示为 Python 软件包版本约束。jax 在运行时检查 jaxlib 版本，而不是使用 pip 软件包版本约束，因为我们 为各种硬件和软件版本（例如，GPU、TPU 等）提供单独的 jaxlib Wheel。由于我们不知道哪个是对任何给定用户的正确选择，因此我们不希望 pip 为我们自动安装 jaxlib 软件包。

未来，我们希望将 jaxlib 的硬件特定部分分离到单独的插件中，届时最低版本可以表示为 Python 软件包依赖项。目前，我们确实提供了平台特定的额外要求，以安装兼容的 jaxlib 版本，例如，jax[cuda]。

如何安全地更改 jaxlib 的 API？#

  • 只要最低 jaxlib 版本增加到兼容版本，jax 可以随时放弃与旧版 jaxlib 版本的兼容性。但是，请注意，即使对于未发布的 jax 版本，最低 jaxlib 也必须是已发布的版本！这使我们能够在 CI 构建中使用已发布的 jaxlib Wheel，并允许 Python 开发人员在 HEAD 上处理 jax，而无需构建 jaxlib。

    例如，要删除 jax Python 代码中的旧向后兼容性路径，只需提高最低 jaxlib 版本，然后删除兼容性路径即可。

  • jaxlib 可以放弃与低于其自身发布版本号的旧版 jax 版本的兼容性。jax 强制执行的版本约束将禁止使用不兼容的 jaxlib。

    例如，为了使 jaxlib 放弃旧版 jax 版本使用的 Python 绑定 API，必须增加 jaxlib 的次版本号或主版本号。

  • 如果可能，对 jaxlib 的更改应以向后兼容的方式进行。

    通常，jaxlib 可以自由更改其 API，只要遵循关于 jax 与所有至少与最低版本一样新的 jaxlib 兼容的规则即可。这意味着 jax 必须始终与至少两个版本的 jaxlib 兼容，即上一个版本和树尖版本，实际上是下一个版本。如果保持兼容性，这将更容易做到，尽管可以使用来自 jax 的版本测试进行不兼容的更改；请参见下文。

    例如，通常可以安全地向 jaxlib 添加新函数，但不安全地删除现有函数或更改其签名（如果当前的 jax 仍在 使用它）。对于所有大于最低版本直到 HEAD 的 jaxlib 版本，对 jax 的更改必须有效或优雅地降级。

请注意，此处的兼容性规则仅适用于 jax 和 jaxlib 的已发布版本。它们不适用于未发布的版本；也就是说，如果 API 从未发布，或者没有已发布的 jax 版本使用该 API，则可以从 jaxlib 中引入然后删除 API。

jaxlib 的源代码是如何布局的？#

jaxlib 分布在两个主要存储库中，即主 JAX 存储库中的 jaxlib/ 子目录 和 XLA 源代码树（位于 XLA 存储库内）。XLA 内部的 JAX 特定部分主要位于 xla/python 子目录中。

JAX 的 C++ 部分（例如 Python 绑定和运行时组件）位于 XLA 树内的原因部分是历史原因，部分是技术原因。

历史原因是最初 xla/python 绑定被设想为通用 Python 绑定，可以与其他框架共享。实际上，这种情况越来越不真实，xla/python 包含许多 JAX 特定部分，并且可能会包含更多。因此，最好简单地将 xla/python 视为 JAX 的一部分。

技术原因是 XLA C++ API 不稳定。通过将 XLA:Python 绑定保留在 XLA 树中，它们的 C++ 实现可以与 XLA 的 C++ API 原子地更新。维护 Python API 的向后和向前兼容性比 C++ API 更容易，因此 xla/python 公开 Python API，并负责在 Python 级别维护向后兼容性。

jaxlib 是使用 Bazel 从 jax 存储库构建的。来自 XLA 存储库的 jaxlib 片段作为 Bazel 子模块合并到构建中。要更新构建期间使用的 XLA 版本，必须更新 Bazel WORKSPACE 中的固定版本。这是根据需要手动完成的，但可以在每次构建时覆盖。

在版本发布之间，我们如何在 jax 和 jaxlib 边界之间进行更改？#

jaxlib 版本是一个粗略的工具：它只允许我们推断版本发布。

但是，由于 jax 和 jaxlib 代码分布在无法在单个更改中原子更新的存储库中，因此我们需要以比我们的发布周期更精细的粒度管理兼容性。为了管理细粒度的兼容性，我们有独立于 jaxlib 发布版本号的额外版本控制。

我们在 XLA 存储库中的 xla_client.py 中维护一个额外的版本号 (_version)。想法是，此版本号在 xla/python 中与 JAX 的 C++ 部分一起定义，也可以作为 jax._src.lib.jaxlib_extension_version 供 JAX Python 访问，并且每次对 XLA/Python 代码进行更改（对 jax 具有向后兼容性影响）时都必须递增。然后，JAX Python 代码可以使用此版本号来维护向后兼容性，例如

from jax._src.lib import jaxlib_extension_version

# 123 is the new version number for _version in xla_client.py
if jaxlib_extension_version >= 123:
  # Use new code path
  ...
else:
  # Use old code path.

请注意，此版本号除了对已发布版本号的约束之外，也就是说，此版本号的存在是为了帮助管理未发布代码的开发期间的兼容性。版本发布还必须遵循上面给出的兼容性规则。