jax.extend 模块#

用于 JAX 扩展的模块。

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

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

API 策略#

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

模块#