jax.extend 模块#

用于 JAX 扩展的模块。

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

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

API 策略#

与公共 API不同，此模块不提供跨版本的兼容性保证。破坏性更改将通过JAX 项目更新日志进行公告。

模块#