JAX API 兼容性与弃用政策全解析三个月弃用周期、公开 API 边界与数值稳定性承诺【免费下载链接】jaxComposable transformations of PythonNumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/gh_mirrors/jax/jaxJAX 作为一套持续快速演进的 PythonNumPy 可组合变换框架在推进 API 改进的同时也需要尽量降低对用户社区的破坏性影响。本文基于 docs/api_compatibility.md 全面解读 JAX 的 API 兼容性承诺、三个月弃用Deprecation政策、受保护与不受保护的 API 边界以及数值与随机数方面的稳定性限制并结合仓库源码如 jax/_src/deprecations.py、jax/init.py、CHANGELOG.md揭示其底层实现机制。读完本文你将能够准确判断一个 JAX 接口是否受兼容性保护、理解弃用警告的生命周期并为自己的项目制定稳妥的升级与迁移策略。JAX 的兼容性承诺与总体原则JAX 处于持续演进之中团队希望保留对 API 进行改进的空间。与此同时为了最大限度减少对用户社区的冲击JAX 会尽量减少破坏性变更breaking changes并且只在必要时才做出不兼容改动。其总体原则可以概括为三点允许演进JAX 不会冻结 API新的设计如jax.tree取代jax.tree_util的部分接口会持续推进降低冲击任何不兼容的改动都必须经过正式的弃用流程给用户留出迁移时间明确定界兼容性承诺只覆盖公开 API内部实现细节不受保护。需要注意的是这份政策包括弃用周期长度随时可能被 JAX 团队调整文档中明确保留了这一权利。因此在依赖较新的 JAX 功能时建议关注 CHANGELOG.md 和官方发布说明。三个月弃用Deprecation政策详解这是整个兼容性文档的核心。当一个 API 发生不兼容变更时JAX 会尽力遵循以下流程弃用宣布阶段变更会在 CHANGELOG.md 中宣布被弃用 API 的docstring中也会标注弃用说明旧 API 会发出DeprecationWarningPython 标准库warnings机制中的类别。从源码看这一机制在 jax/_src/deprecations.py 中有完整实现deprecation_getattr()会在模块属性被访问时检查_deprecations字典命中则调用warnings.warn(message, DeprecationWarning, stacklevel2)并返回被弃用的对象见该文件 第 47-57 行。以 jax/init.py 中实际注册的弃用条目为例_deprecations { tree_flatten: ( jax.tree_flatten is deprecated: use jax.tree.flatten (jax v0.4.25 or newer) or jax.tree_util.tree_flatten (any JAX version)., _deprecated_tree_flatten ), tree_map: ( jax.tree_map is deprecated: use jax.tree.map (jax v0.4.25 or newer) or jax.tree_util.tree_map (any JAX version)., _deprecated_tree_map ), clear_backends: (jax.clear_backends is deprecated., _deprecated_clear_backends), xla_computation: ( jax.xla_computation is deprecated. Please use the AOT APIs; see https://jax.readthedocs.io/en/latest/aot.html. For example, replace xla_computation(f)(*xs) with jit(f).lower(*xs).compiler_ir(hlo). See CHANGELOG.md for 0.4.30 for more examples., _deprecated_xla_computation ), }这些条目会通过deprecation_getattr注册为模块级__getattr__见 jax/init.py#L251-L254从而在用户访问旧名字时自动发出警告。可以看到真实的弃用消息会给出明确的替代 API和JAX 版本要求这正是迁移友好设计的体现。移除阶段从弃用该 API 的 JAX 版本发布之日起三个月后JAX 随时可能移除该 API三个月是下限lower bound而非承诺的移除时间点该周期有意设计得比许多更成熟的项目更短例如 NumPy/SciPy 通常采用数年甚至无限期的弃用窗口在实践中移除往往比三个月晚得多尤其是当某个功能有大量用户时JAX 会倾向于延长弃用时间如果三个月的弃用期对你的项目造成了困难文档明确建议向 JAX 团队反馈。从 CHANGELOG.md 可以观察到真实的弃用与移除节奏例如jax.xla_computation于 0.4.30 被弃用并建议迁移到 AOT API而后续版本中previously-deprecated的内部 API 才被批量移除。对于长期维护的项目建议在 JAX 发布新版本时留意 changelog 中标记为deprecated的条目。受兼容性政策保护的公开 API 范围JAX 的兼容性承诺仅覆盖公开 API具体包括以下模块模块说明jax顶层核心 APIjit、grad、vmap、pmap、tree等jax.dlpackDLPack 张量互操作接口jax.image图像处理辅助函数jax.lax底层操作原语控制流、随机、卷积等jax.nn神经网络辅助函数激活、softmax、one-hot 等jax.numpyNumPy 风格 APIjax.ops操作辅助函数jax.profiler性能分析接口jax.random伪随机数生成见下文数值与随机数一节jax.scipySciPy 风格 APIjax.tree树工具jax.tree.map、jax.tree.leaves等新接口jax.tree_util树工具传统接口jax.tree_util.tree_flatten等jax.test_util测试工具需要特别指出两点边界并非这些模块中的所有内容都算公开 API。JAX 正在逐步将公开与私有 API 分离公开 API 以 JAX 官方文档中有记录为准JAX 的目标是让所有非公开 API 都以下划线开头命名但目前尚未完全做到这一点。因此即使一个名字出现在上述公开模块中只要它未被文档化就可能不受兼容性保护。一个典型的例子是jax.core虽然它作为顶层模块可见但官方将其视为内部模块其中pp_*系列调试打印工具、DimSize、Shape、canonicalize_shape等在 jax/core.py 中均以_deprecations字典形式注册了弃用并明确标注Finalized ... remove after ...的移除时间节点——这充分说明这些名字不属于受保护的公开 API。不受保护的范围内部 API 与实验性模块以下内容明确不在兼容性承诺的覆盖范围内任何以下划线_开头的名字jax._src真正的源码实现目录jax.corejax.libjax.interpretersjax.experimental实验性模块如jax.experimental.maps、jax.experimental.host_callback等jax.example_libraries示例库如optimizers、staxjax.extend扩展 API如jax.extend.ffi、jax.extend.mlir。文档同时强调这份清单并非穷尽性的This list is not exhaustive。也就是说未列入上表但同样属于内部实现的内容也不受保护。这一边界在源码中有直接体现。例如 jax/experimental/maps.py 中if deprecations.is_accelerated(jax-experimental-maps-module): raise ImportError(...)实验性模块可以随时变更或移除而不经过完整弃用周期。jax.experimental与jax.extend的划分本身也在演进中例如jax.experimental.export已被弃用并迁移到jax.export见 CHANGELOG.md相关弃用条目同样通过deprecation_getattr机制注册在 jax/experimental/export/init.py 中。实践建议如果你的代码依赖jax.experimental或jax.core中的接口应当将其视为随时可能变动并尽量通过薄封装隔离便于未来迁移。对于jax.extend其本身就是为开发者提供稳定扩展点而设计但仍需参照 jax.extend 文档 确认哪些接口是受支持的。数值结果与随机数的兼容性边界除了 API 命名层面的兼容性JAX 还明确了数值层面的承诺边界数值结果不做精确稳定承诺数值运算的精确输出值不保证跨 JAX 版本稳定即使在同一个 JAX 版本内精确数值也不一定稳定可能因以下因素而异不同的加速器平台CPU / GPU / TPU是否在jax.jit内执行其他执行环境差异。这是所有编译器类框架的共性XLA 的算子融合、重排与底层指令集都可能改变浮点运算的舍入顺序。因此将 JAX 用于需要逐位bit-exact复现数值结果的场景如回归测试中的精确断言是不可靠的建议改用容差断言如jax.test_util提供的check_close类工具。随机数只保证分布不保证具体输出对于固定的 PRNG key 输入jax.random中伪随机函数的输出值可能跨 JAX 版本变化。兼容性政策仅适用于输出的分布。例如文档给出的示例jax.random.gumbel(jax.random.key(72))这个表达式在不同 JAX 版本间可能返回不同的具体数值但jax.random.gumbel会始终保持为Gumbel 分布的伪随机生成器。关于随机数变更JAX 采取以下实践尽量不频繁修改伪随机输出值一旦发生修改会在 changelog 中宣布但不遵循弃用周期因为分布本身没有改变只是抽样序列变了在某些情况下JAX 可能提供一个**临时配置开关transient configuration flag**来恢复旧行为帮助用户诊断和迁移受影响代码这类开关的存续时间与一个弃用窗口的时长相当即约三个月到期后会被移除。从源码看这类加速/回退机制在 jax/_src/deprecations.py 中以DeprecationState和register/unregister/accelerate/is_accelerated/warn一组函数实现第 77-113 行默认弃用只发警告一旦被加速accelerated则直接抛错jax/_src/core.py中的 tracer 哈希弃用is_accelerated(tracer-hash)和 jax/_src/numpy/lax_numpy.py 中的数组替代 dtype 弃用is_accelerated(jax-numpy-array-none)都是这套机制的实例。仓库源码中的弃用机制实现JAX 的弃用机制在 jax/_src/deprecations.py 中集中实现核心包括两套互补机制机制一模块属性级弃用import 型适用于被移除的模块顶层名字。通过deprecation_getattr(module, deprecations)生成模块的__getattr__当访问已弃用名字时发出DeprecationWarning并返回替代对象若某条目被加速fn is None则直接抛出AttributeError。该机制被 jax/init.py、jax/core.py、jax/lax/init.py、jax/numpy/init.py、jax/nn/init.py 等多个公开模块复用。机制二注册式弃用非 import 型适用于函数参数、行为变化等无法用属性访问触发的弃用。通过全局字典_registered_deprecations以字符串 ID 注册DeprecationStatewarn(id, message, stacklevel)在未加速时发出警告、加速时抛ValueError。jax.experimental.maps模块的弃用即属此类。弃用加速Accelerationaccelerate_getattr_deprecation与accelerate用于把警告升级为硬错误主要用于 Google 内部代码库推进快速弃用jax/_src/deprecations.py#L60-L67。测试工具 jax/_src/test_util.py 中还有对应的unaccelerate_getattr_deprecation用于测试环境中解除加速。对于普通用户而言DeprecationWarning默认通常不会显示可通过python -W error::DeprecationWarning或warnings.simplefilter(always)显式开启以便及早发现迁移点。给开发者的迁移与升级实践建议结合以上政策在使用 JAX 时建议遵循以下实践升级前必读 changelog每次升级 JAX 前查看 CHANGELOG.md 中标记deprecated的条目确认你使用的接口是否受影响把DeprecationWarning当错误处理在 CI 中开启-W error::DeprecationWarning建议配合-W ignore::DeprecationWarning过滤无关第三方警告让弃用提示在升级早期暴露只依赖受保护的公开 API优先使用官方文档有记录的接口对jax.experimental、jax.core、jax.interpreters等未保护模块中的依赖用封装层隔离数值断言用容差不要依赖跨版本、跨平台的逐位数值结果随机数测试不要断言具体序列值只断言统计分布性质关注替代 API 的版本要求如弃用消息中usejax.tree.flatten(jax v0.4.25 or newer)这类信息迁移时需确认目标替代接口在你支持的最低 JAX 版本中可用及时反馈如果三个月的弃用窗口对你的项目不够可按文档建议向 JAX 团队提出实践中大量弃用尤其用户量大的功能会持续更久。延伸阅读本政策文档原始出处docs/api_compatibility.mdJAX 的 Python/NumPy/SciPy 版本支持政策遵循科学计算社区 SPEC 0 并延长 Python 支持窗口至 45 个月docs/deprecation.md历次版本的弃用与移除记录CHANGELOG.md弃用机制的源码实现jax/_src/deprecations.py顶层模块弃用条目示例jax/init.py发布说明与开发者文档索引docs/notes.rst【免费下载链接】jaxComposable transformations of PythonNumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/gh_mirrors/jax/jax创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
