jax.extend module

Modules for JAX extensions.

The jax.extend package provides modules for access to JAX internal machinery. See JEP #15856.

API policy

Unlike the public API, this package offers no compatibility guarantee across releases. Breaking changes will be announced via the JAX project changelog.

Modules

• jax.extend.linear_util module
• jax.extend.mlir module
• jax.extend.random module