jax.extend module

Modules for JAX extensions.

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

This module is not the only means by which JAX aims to be extensible. For example, the main JAX API offers mechanisms for customizing derivatives, registering custom pytree definitions, and more.

API policy

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

Modules

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