logo

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
  • Runtime value debugging in JAX
    • jax.debug.print and jax.debug.breakpoint
    • The checkify transformation
    • JAX debugging flags

Reference Documentation

  • JAX Frequently Asked Questions (FAQ)
  • Asynchronous dispatch
  • Ahead-of-time lowering and compilation
  • Understanding Jaxprs
  • Convolutions in JAX
  • Pytrees
  • Type promotion semantics
  • JAX Errors
  • Transfer guard
  • 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

Developer documentation

  • Contributing to JAX
  • Building from source
  • Internal APIs
  • Autodidax: JAX core from scratch
  • JAX Enhancement Proposals (JEPs)
    • 263: JAX PRNG Design
    • 2026: Custom JVP/VJP rules for JAX-transformable functions
    • 4008: Custom VJP and `nondiff_argnums` update
    • 4410: Omnistaging
    • 9407: Design of Type Promotion Semantics for JAX
    • 9419: Jax and Jaxlib versioning
    • 10657: Sequencing side-effects in JAX
    • 11830: `jax.remat` / `jax.checkpoint` new implementation
    • 12049: Type Annotation Roadmap for JAX

API documentation

  • Public API: jax package
    • jax.numpy package
    • jax.scipy package
    • JAX configuration
    • jax.debug package
    • jax.dlpack module
    • jax.distributed module
    • jax.example_libraries package
      • jax.example_libraries.optimizers module
      • jax.example_libraries.stax module
    • jax.experimental package
      • jax.experimental.checkify module
      • jax.experimental.global_device_array module
      • jax.experimental.host_callback module
      • jax.experimental.maps module
      • jax.experimental.pjit module
      • jax.experimental.sparse module
      • jax.experimental.jet module
    • jax.flatten_util package
    • jax.image package
    • jax.lax package
    • jax.nn package
      • jax.nn.initializers package
    • jax.ops package
    • jax.profiler module
    • jax.random package
    • jax.stages package
    • jax.tree_util package
    • jax.lib package
Theme by the Executable Book Project
  • .rst

Tutorial: JAX 101

Tutorial: JAX 101#

This is a tutorial developed by engineers and researchers at DeepMind.

Tutorials

  • 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

previous

🔪 JAX - The Sharp Bits 🔪

next

JAX As Accelerated NumPy

By The JAX authors
© Copyright 2020, The JAX Authors. NumPy and SciPy documentation are copyright the respective authors..