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.

Getting Started

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

• Tutorial: JAX 101
  • JAX As Accelerated NumPy
  • Just In Time Compilation with JAX
  • Automatic Vectorization in JAX
  • Advanced Automatic Differentiation in JAX
  • Pseudo Random Numbers in JAX
  • Working with Pytrees
  • Parallel Evaluation in JAX
  • Stateful Computations in JAX
  • Introduction to pjit

Reference Documentation

• JAX Frequently Asked Questions (FAQ)
• Asynchronous dispatch
• Understanding Jaxprs
• Convolutions in JAX
• Pytrees
• Type promotion semantics
• JAX Errors
• JAX Glossary of Terms
• Change log

Advanced JAX Tutorials

• The Autodiff Cookbook
• 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
• Using JAX in multi-host and multi-process environments

Notes

• API compatibility
• Python and NumPy version support policy
• Concurrency
• GPU memory allocation
• Profiling JAX programs
• Device Memory Profiling
• Rank promotion warning
• custom_vjp and nondiff_argnums update guide
• Transfer guard

Developer documentation

• Contributing to JAX
  • Ways to contribute
  • Contributing code using pull requests
  • JAX pull request checklist
• Building from source
  • Building or installing jaxlib
  • Installing jax
• Running the tests
  • Doctests
• Type checking
• Linting
• Update documentation
  • Update notebooks
  • Documentation building on readthedocs.io
• Internal APIs
  • core
• Autodidax: JAX core from scratch
  • Part 1: Transformations as interpreters: standard evaluation, jvp, and vmap
  • Pytrees and flattening user functions’ inputs and outputs
  • Part 2: Jaxprs
  • Part 3: jit, simplified
  • Part 4: linearize and vjp (and grad!)
  • Part 5: the control flow primitives cond
• Design Notes
  • Custom JVP/VJP rules for JAX-transformable functions
  • Jax and Jaxlib versioning
  • Omnistaging
  • JAX PRNG Design
  • Design of Type Promotion Semantics for JAX
  • Sequencing side-effects in JAX

API documentation

• Public API: jax package
  • Subpackages
    • jax.numpy package
    • jax.scipy package
    • JAX configuration
    • jax.dlpack module
    • jax.distributed module
    • jax.example_libraries package
    • jax.experimental package
    • jax.flatten_util package
    • jax.image package
    • jax.lax package
    • jax.nn package
    • jax.ops package
    • jax.profiler module
    • 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