JAX reference documentation

JAX is Autograd and XLA, brought together for high-performance numerical computing and machine learning research. It provides composable transformations of Python+NumPy programs: differentiate, vectorize, parallelize, Just-In-Time compile to GPU/TPU, and more.

Note

JAX 0.4.0 introduces new parallelism APIs, including breaking changes to jax.experimental.pjit() and a new unified jax.Array type. Please see Parallelism with JAX tutorial and the jax.Array migration guide for more information.

Getting Started

• Installing JAX
• JAX Quickstart
• How to Think in JAX
• 🔪 JAX - The Sharp Bits 🔪

• Tutorial: JAX 101

• Runtime value debugging in JAX
  • Interactive inspection with jax.debug
  • Functional error checks with jax.experimental.checkify
  • Throwing Python errors with JAX’s debug flags

Reference Documentation

• JAX Frequently Asked Questions (FAQ)
• Asynchronous dispatch
• Ahead-of-time lowering and compilation
• Understanding Jaxprs
• Convolutions in JAX
• Pytrees
• jax.Array migration
• Type promotion semantics
• JAX Errors
• Transfer guard
• JAX Glossary of Terms
• Change log

Advanced JAX Tutorials

• The Autodiff Cookbook
• Using JAX in multi-host and multi-process environments
• Parallelism with JAX
• Autobatching log-densities example
• Training a Simple Neural Network, with tensorflow/datasets Data Loading
• Custom derivative rules for JAX-transformable Python functions
• How JAX primitives work
• Writing custom Jaxpr interpreters in JAX
• Training a Simple Neural Network, with PyTorch Data Loading
• Named axes and easy-to-revise parallelism

Developer documentation

• Contributing to JAX
• Building from source
• Running the tests
• Type checking
• Linting
• Update documentation
• Internal APIs
• Autodidax: JAX core from scratch
• JAX Enhancement Proposals (JEPs)

API documentation

• Public API: jax package

Notes

• API compatibility
• Python and NumPy version support policy
• Concurrency
• GPU memory allocation
• Profiling JAX programs
• Device Memory Profiling
• Rank promotion warning

Indices and tables

• Index

• Module Index

• Search Page