jax.extend 模块#

JAX 扩展的模块。

jax.extend 模块提供了用于访问 JAX 内部机制的模块。请参阅 JEP #15856。

此模块并非 JAX 实现可扩展性的唯一途径。例如，JAX 的主 API 提供了用于 自定义导数、注册自定义 Pytree 定义 等的机制。

API 策略#

与 公共 API 不同，此模块在版本之间不提供兼容性保证。重大更改将通过 JAX 项目更新日志 进行公告。

模块#