JAX reference documentation

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more.

For an introduction to JAX, start at the JAX GitHub page.

Tutorials

• JAX Quickstart
• The Autodiff Cookbook
• Training a Simple Neural Network, with PyTorch Data Loading

Advanced JAX Tutorials

• 🔪 JAX - The Sharp Bits 🔪
• XLA in Python
• JAX pytrees
• How JAX primitives work
• Writing custom Jaxpr interpreters in JAX
• Training a Simple Neural Network, with Tensorflow Datasets Data Loading
• MAML Tutorial with JAX
• Generative Modeling by Estimating Gradients of Data Distribution in JAX
• Autobatching log-densities example

Notes

• Asynchronous dispatch
• Concurrency
• GPU memory allocation
• Profiling JAX programs
• Rank promotion warning
• Type promotion semantics

Developer documentation

• Building from source
• Running the tests
• Update documentation
  • Update notebooks
  • Documentation building on readthedocs.io

API documentation

• jax package
  • Subpackages
    • jax.numpy package
    • jax.scipy package
    • jax.experimental package
    • jax.lax package
    • jax.nn package
    • jax.ops package
    • jax.random package
    • jax.tree_util package
  • Just-in-time compilation (jit)
  • Automatic differentiation
  • Vectorization (vmap)
  • Parallelization (pmap)

Indices and tables

• Index
• Module Index
• Search Page