API 兼容性#
JAX 正在不断发展,我们希望能够对其 API 进行改进。话虽如此,我们希望最大限度地减少对 JAX 用户社区的影响,并尽量避免进行重大更改。
JAX 版本控制#
JAX 使用基于努力的版本控制(请参阅JEP 25516:JAX 的基于努力的版本控制),目前处于零版本阶段。这意味着对于版本 0.X.Y,增加 Y 将引入次要的重大更改,而增加 X 将引入主要的重大更改。
对于任何重大更改,JAX 目前遵循 3 个月的弃用策略。当对 API 进行不兼容更改时,我们将尽最大努力遵守以下程序:
更改将在
CHANGELOG.md中以及弃用 API 的文档字符串中进行 anúncio,旧 API 将发出DeprecationWarning。在弃用 API 的
jax版本发布三个月后,我们随时可能删除被弃用的 API。请注意,三个月是一个下限,其选择时间比许多更成熟的项目要短。实际上,弃用可能需要相当长的时间,特别是当有许多用户使用某个功能时。如果三个月的弃用期出现问题,请与我们联系。
我们保留随时更改此策略的权利。
涵盖范围?#
仅涵盖公开的 JAX API,其中包括以下模块:
jaxjax.dlpackjax.imagejax.laxjax.nnjax.numpyjax.opsjax.profilerjax.random(详情请参阅下文)jax.scipyjax.treejax.tree_utiljax.test_util
并非这些模块中的所有内容都旨在公开,并且随着时间的推移,我们正在努力区分公开 API 和私有 API。公共 API 在 JAX 文档中有记录。此外,我们的目标是所有非公共 API 都应以前导下划线命名,尽管我们尚未完全遵守这一点。
不涵盖范围?#
明确的私有 API#
任何以_开头的 API 或导入路径均明确为私有,并且可能在 JAX 版本之间无警告地更改。我们正在努力将所有私有 API 移至 jax._src,以使这些预期更加明确。
jaxlib#
jaxlib 包中的任何导入路径均被视为私有,并且可能在版本之间无警告地更改。jaxlib 中定义的某些 API 在 jax 包中有公开的别名。
遗留内部 API#
此外,还有几个遗留模块目前在没有下划线的情况下公开了一些私有 API,包括:
jax.corejax.interpretersjax.libjax.util
我们正在积极致力于弃用这些模块及其包含的 API。在大多数情况下,此类弃用将遵循 3 个月的弃用期,但这可能并不总是可行的。如果您使用任何此类 API,请预计它们很快将被弃用,并寻求替代方案。
实验性和示例库#
以下模块包含用于实验或演示目的的代码,并且 API 可能在版本之间无警告地更改:
jax.experimentaljax.example_libraries
我们理解一些用户依赖于 jax.experimental,因此在大多数情况下,我们遵循 3 个月的弃用期来处理更改,但这可能并非总是可行的。
JAX 扩展#
模块包含半公开的 JAX 内部 API,这些 API 供下游项目使用,但没有与主 JAX 包相同的稳定性保证。如果您有使用 jax.extendjax.extend 的代码,我们强烈建议针对 JAX 的 nightly 版本进行 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 仍将是 Gumbel 分布的伪随机生成器。
我们尽量不频繁地更改伪随机值。当发生更改时,会在变更日志中宣布,但不遵循弃用周期。在某些情况下,JAX 可能会公开一个临时的配置标志来恢复新行为,以帮助用户诊断和更新受影响的代码。此类标志的有效期将相当于弃用窗口期。