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 🔪
• JAX pytrees
• How JAX primitives work
• Writing custom Jaxpr interpreters in JAX
• Training a Simple Neural Network, with Tensorflow Datasets Data Loading
• Autobatching log-densities example

Notes

• Change Log
• Understanding JAXPR
• Asynchronous dispatch
• Concurrency
• GPU memory allocation
• Profiling JAX programs
• Rank promotion warning
• Type promotion semantics

Developer documentation

• Building from source
  • Building or installing jaxlib
  • Installing jax
• Running the tests
• Update documentation
  • Update notebooks
  • Documentation building on readthedocs.io
• Internal APIs
  • core

API documentation

• Public API: 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
    • jax.dlpack module
  • Just-in-time compilation (jit)
  • Automatic differentiation
  • Vectorization (vmap)
  • Parallelization (pmap)

Indices and tables

• Index
• Module Index
• Search Page