Skip to content

Course in preparation · Coming soon

JAX for Research: Transforms, jit and Sharding

Use JAX the way research teams on TPUs do. Write pure functions, transform them with grad, vmap and jit, and shard a model across devices you simulate on one machine.

Who it is for. Researchers and engineers aiming at labs that train on TPUs, who know NumPy and want the functional model behind JAX.

The planned syllabus

Every module will meet the book's bar before it opens: the math derived on the page, a Colab notebook on Python 3.12, and measurements from real models.

  1. 01

    Pure functions and PRNG keys

    Why JAX has no hidden state, and how explicit randomness makes a run reproducible.

  2. 02

    grad, vmap and jit

    Compose the transforms, read what XLA compiles, and see what a recompile costs.

  3. 03

    A model in Flax

    Parameters as data, a training step as a function, and a small language model trained on Colab.

  4. 04

    Sharding

    Meshes and partition specs on simulated devices: data and tensor parallelism without a TPU pod.

  5. 05

    Debugging compiled code

    Tracers, shape errors and the tools for finding what jit changed.

Start now with what it builds on

All courses