Skip to main content
Ctrl+K
JAX documentation - Home

Getting Started

  • Installing JAX
  • JAX Quickstart
  • How to Think in JAX
  • 🔪 JAX - The Sharp Bits 🔪
  • JAX Frequently Asked Questions (FAQ)
  • 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

Further Resources

  • User Guides
    • Profiling JAX programs
    • Device Memory Profiling
    • Runtime value debugging in JAX
      • jax.debug.print and jax.debug.breakpoint
      • The checkify transformation
      • JAX debugging flags
    • GPU peformance tips
    • Understanding Jaxprs
    • External Callbacks in JAX
    • Type promotion semantics
    • Pytrees
    • Ahead-of-time lowering and compilation
    • JAX Errors
    • Transfer guard
    • Pallas: a JAX kernel language
      • Pallas Design
      • Pallas Quickstart
      • Writing TPU kernels with Pallas
  • Advanced Tutorials
    • Training a Simple Neural Network, with tensorflow/datasets Data Loading
    • Training a Simple Neural Network, with PyTorch Data Loading
    • Autobatching for Bayesian Inference
    • Using JAX in multi-host and multi-process environments
    • Distributed arrays and automatic parallelization
    • Named axes and easy-to-revise parallelism with xmap
    • The Autodiff Cookbook
    • Custom derivative rules for JAX-transformable Python functions
    • Control autodiff’s saved values with jax.checkpoint (aka jax.remat)
    • How JAX primitives work
    • Writing custom Jaxpr interpreters in JAX
    • Custom operations for GPUs with C++ and CUDA
    • Generalized Convolutions in JAX
  • 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
      • 9263: Typed keys & pluggable RNGs
      • 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
      • 14273: `shard_map` (`shmap`) for simple per-device code
      • 15856: `jax.extend`, an extensions module
      • 17111: Efficient transposition of `shard_map` (and other maps)
      • 18137: Scope of JAX NumPy & SciPy Wrappers
    • Investigating a regression
  • Building on JAX
  • Notes
    • API compatibility
    • Python and NumPy version support policy
    • jax.Array migration
    • Asynchronous dispatch
    • Concurrency
    • GPU memory allocation
    • Rank promotion warning
  • Public API: jax package
    • jax.numpy module
    • jax.scipy module
    • jax.lax module
    • jax.random module
    • jax.sharding module
    • jax.debug module
    • jax.dlpack module
    • jax.distributed module
    • jax.dtypes module
    • jax.flatten_util module
    • jax.image module
    • jax.nn module
      • jax.nn.initializers module
    • jax.ops module
    • jax.profiler module
    • jax.stages module
    • jax.tree_util module
    • jax.typing module
    • jax.extend module
    • jax.example_libraries module
      • jax.example_libraries.optimizers module
      • jax.example_libraries.stax module
    • jax.experimental module
      • jax.experimental.array_api module
      • jax.experimental.checkify module
      • jax.experimental.host_callback module
      • jax.experimental.maps module
      • jax.experimental.pjit module
      • jax.experimental.sparse module
      • jax.experimental.jet module
      • jax.experimental.custom_partitioning module
      • jax.experimental.multihost_utils module
      • jax.experimental.compilation_cache module
    • jax.lib module
  • Change log
  • JAX Glossary of Terms
  • .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

previous

JAX Frequently Asked Questions (FAQ)

next

JAX As Accelerated NumPy

By The JAX authors

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