jax.extend 模块#

用于 JAX 扩展的模块。

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

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

API 策略#

与公共 API 不同，此模块不保证跨版本兼容性。重大变更将通过JAX 项目更新日志公布。

模块#