Topolearn

Learn JAX

NumPy-style arrays with composable grad, jit and vmap on CPU, GPU and TPU

17 lessons · 4 hours · 4 levels · covers JAX 0.11

JAX is a Python library for accelerated array computing and program transformation. It offers a NumPy-like API (jax.numpy) whose functions can be differentiated with grad, compiled with XLA via jit, vectorized with vmap and sharded across many devices. It powers large-scale ML research and training stacks such as Flax and Optax, and scientific computing that needs fast gradients.

Sign in free to track progressOfficial docs ↗

Before you start

JAX builds on these; the tree links to the exact lessons you need.

PythonPractical levelNumPyEssentials level

How deep do you want to go?

1. EssentialsRead and write simple code, do the everyday basics5 lessons · 52 min in total
2. PracticalUse it productively at work, with the usual idioms10 lessons · 2 hours in total
3. AdvancedThe less common features, performance and trade-offs14 lessons · 3 hours in total
4. ExpertInternals and the edge cases experts know17 lessons · 4 hours in total

The JAX tree

JAXjax.numpy arrays and devicesjax.numpy arrays anddevicesImmutable arrays and .at updatesImmutable arrays and.at updatesAutomatic differentiation with gradAutomaticdifferentiation with…Compiling with jitCompiling with jitAutomatic batching with vmapAutomatic batchingwith vmapExplicit random keysExplicit random keysPytrees: nested containers of arraysPytrees: nestedcontainers of arraysTracing, static arguments and recompilesTracing, staticarguments and rec…Structured control flow with laxStructured controlflow with laxA functional training loopA functional trainingloopThe ndarrayThe ndarrayData types (dtypes)Data types (dtypes)Indexing and slicingIndexing and slicingDefining functionsDefining functionsShape and reshapeShape and reshapeDictionariesDictionaries
Essentials and Practical shown (10 of 17 lessons); the full list is below. Click a lesson to read it.

All 17 lessons, in order

Arrays and the JAX model

  1. jax.numpy arrays and devices · Essentials · 10 min

    Use jax.numpy as a drop-in NumPy, and see where JAX arrays live and what dtypes they default to.

  2. Immutable arrays and .at updates · Essentials · 8 min

    Replace in-place assignment with the x.at[idx].set/add/mul syntax that returns a new array.

Core transformations

  1. Automatic differentiation with grad · Essentials · 12 min

    Turn a scalar-valued Python function into a function that returns its gradient.

  2. Compiling with jit · Essentials · 12 min

    Compile a function with XLA so that sequences of array operations run as one fused, fast kernel.

  3. Automatic batching with vmap · Essentials · 10 min

    Write a function for one example and let vmap turn it into a batched function without manual axis bookkeeping.

Functional programming in JAX

  1. Explicit random keys · Practical · 10 min

    Generate random numbers with explicit keys that you split, instead of a hidden global seed.

  2. Pytrees: nested containers of arrays · Practical · 12 min

    Treat dicts, lists, tuples and dataclasses of arrays as single values that every JAX transformation understands.

  3. A functional training loop · Practical · 15 min

    Put grad, jit, pytrees and keys together into a pure train step that returns new parameters.

Tracing and control flow

  1. Tracing, static arguments and recompiles · Practical · 14 min

    Understand tracers and abstract values, and when to mark arguments static to make jit work.

  2. Structured control flow with lax · Practical · 14 min

    Express branches and loops that depend on traced values with lax.cond, fori_loop, while_loop and scan.

  3. Debugging jitted code · Advanced · 12 min

    Print, inspect and catch NaNs inside compiled functions with jax.debug and debug flags.

Autodiff and compiler internals

  1. Jacobians, JVPs and VJPs · Advanced · 14 min

    Go beyond grad: forward and reverse mode, Jacobians, Hessians and stopping gradients.

  2. Jaxprs, lowering and XLA · Expert · 14 min

    Inspect what JAX actually compiles: jaxprs, StableHLO and the compiled executable.

  3. Custom JVPs and VJPs · Expert · 16 min

    Override how JAX differentiates a function for numerical stability, efficiency or non-differentiable code.

  4. Rematerialization with jax.checkpoint · Expert · 12 min

    Trade compute for memory by recomputing activations in the backward pass instead of storing them.

Performance and scaling

  1. Performance: dispatch, donation, caching · Advanced · 14 min

    Make JAX programs fast by avoiding host syncs, reusing buffers, controlling compilation and profiling.

  2. Sharding across devices · Advanced · 16 min

    Spread arrays and computation over many devices with meshes, NamedSharding and jit's automatic partitioning.