jax.extend 模块

用于 JAX 扩展的模块。

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

API 策略

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

模块

• jax.extend.ffi 模块
• jax.extend.linear_util 模块
• jax.extend.mlir 模块
• jax.extend.random 模块