API 兼容性#

JAX 在不断发展,我们希望能够改进其 API。也就是说,我们希望尽量减少 JAX 用户社区的变动,并尽量不进行重大更改。

JAX 版本控制#

JAX 使用 基于努力的版本控制(请参阅 JEP 25516:JAX 的基于努力的版本控制),目前处于零版本阶段。这意味着对于版本 0.X.Y,增加 Y 将引入小的重大更改,而增加 X 将引入大的重大更改。

对于任何重大更改,JAX 目前遵循 3 个月的弃用策略。当 API 发生不兼容的更改时,我们将尽最大努力遵守以下步骤

  • 此变更将在 CHANGELOG.md 中以及已弃用 API 的文档字符串中声明,并且旧的 API 将发出 DeprecationWarning

  • jax 版本弃用 API 三个月后,我们可能会随时移除已弃用的 API。请注意,三个月是一个下限,并且有意选择比许多更成熟的项目更快。实际上,弃用可能需要相当长的时间,特别是当某个功能有许多用户时。如果三个月的弃用期出现问题,请向我们提出。

我们保留随时更改此政策的权利。

涵盖哪些内容?#

仅涵盖公共 JAX API,包括以下模块

  • jax

  • jax.dlpack

  • jax.image

  • jax.lax

  • jax.nn

  • jax.numpy

  • jax.ops

  • jax.profiler

  • jax.random (请参阅以下详细信息

  • jax.scipy

  • jax.tree

  • jax.tree_util

  • jax.test_util

并非这些模块中的所有内容都打算公开,并且随着时间的推移,我们正在努力分离公共和私有 API。公共 API 已在 JAX 文档中记录。此外,我们的目标是所有非公共 API 都应以带有下划线的前缀命名,尽管我们尚未完全遵守此规则。

不涵盖哪些内容?#

显式私有 API#

任何以带有下划线的前缀的 API 或导入路径都是显式私有的,并且可能会在 JAX 版本之间更改而不发出警告。我们正在努力将所有私有 API 移动到 jax._src 中,以使这些预期更加明确。

遗留内部 API#

此外,还有几个遗留模块当前公开了一些没有下划线的私有 API,包括

  • jax.core

  • jax.interpreters

  • jax.lib

  • jax.util

我们正在积极努力弃用这些模块及其包含的 API。在大多数情况下,此类弃用将遵循 3 个月的弃用期,但这并非总是可行。如果您使用任何此类 API,请预计它们将很快被弃用,并寻求替代方案。

实验性和示例库#

以下模块包含用于实验或演示目的的代码,并且 API 可能会在版本之间更改而不会发出警告

  • jax.experimental

  • jax.example_libraries

我们知道某些用户依赖于 jax.experimental,因此在大多数情况下,我们对更改遵循 3 个月的弃用期,但这并非总是可行。

JAX 扩展#

jax.extend 模块包含半公共的 JAX 内部 API,这些 API 旨在供下游项目使用,但不具有与主 JAX 包相同的稳定性保证。如果您的代码使用了 jax.extend,我们强烈建议针对 JAX 的每日版本运行 CI 测试,以便在发布之前捕获潜在的更改。

有关 jax.extend 的详细信息,请参阅 jax.extend 模块文档,或设计文档 jax.extend:用于扩展的模块

数值和随机性#

不保证数值运算的确切值在 JAX 版本之间保持稳定。事实上,在给定的 JAX 版本中,跨加速器平台、在 jax.jit 内外等等,确切的数值不一定稳定。

对于固定的 PRNG 密钥输入,jax.random 中伪随机函数的输出可能会在 JAX 版本之间有所不同。兼容性策略仅适用于输出分布。例如,表达式 jax.random.gumbel(jax.random.key(72)) 可能会在不同的 JAX 版本中返回不同的值,但 jax.random.gumbel 仍将是耿贝尔分布的伪随机生成器。

我们尽量不频繁地更改此类伪随机值。当它们发生时,这些更改会在变更日志中声明,但不遵循弃用周期。在某些情况下,JAX 可能会公开一个临时配置标志,该标志会恢复新行为,以帮助用户诊断和更新受影响的代码。此类标志将持续一个弃用窗口的时间。