Extension guides

Guides for extending JAX’s capabilities, and for building libraries that use or interface with JAX.

Extensible JAX internals

• Writing custom Jaxpr interpreters in JAX
• Custom operations for GPUs with C++ and CUDA
• Check for correctness
• jax.extend module

Libraries and extensions

• Building on JAX