jax.extend module
Contents
jax.extend module#
Modules for JAX extensions.
Thejax.extend module provides modules for access to JAXinternal machinery. SeeJEP #15856.
This module is not the only means by which JAX aims to beextensible. For example, the main JAX API offers mechanisms forcustomizing derivatives,registering custom pytree definitions,and more.
API policy#
Unlike thepublic API,this module offersno compatibility guarantee across releases.Breaking changes will be announced via theJAX project changelog.
Modules#
Contents