jax.extend 模块#

JAX 扩展的模块。

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

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

API 政策#

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

模块#