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.