Udemy
Turn what you know into an opportunity and reach millions around the world.
Learn More
Your cart is empty.
Keep shopping
Introduction to JAX: JIT, Autodiff and the Compiler Stack
New
Rating: 5.0 out of 5(2 ratings)
2 students

Introduction to JAX: JIT, Autodiff and the Compiler Stack

For machine learning and scientific computing. The compiler stack, not just the API: tracing, JIT, XLA, autodiff, vmap
Created byDr. Pietropaolo Frisoni
Last updated 10/2026
English

What you'll learn

  • Read a Jaxpr, the intermediate representation JAX builds from your Python code, and see how tracing produces it
  • Predict when jax.jit traces, compiles and recompiles, and avoid its classic pitfalls
  • Inspect StableHLO and the compiled XLA program
  • Ask the compiler what a function costs before running it, and export a model with jax.export
  • Define a new JAX primitive from scratch, with its evaluation, lowering, differentiation and batching rules
  • Use grad, value_and_grad, Jacobians, Hessians and jvp
  • Vectorize single-example code with vmap, compute per-sample gradients, and compose jit, vmap and grad into one training step
  • Understand data-parallel training with pmap and the modern sharding API, practiced on devices simulated on a laptop CPU
  • Port NumPy code to JAX safely
  • Train a neural network twice, in pure JAX with pytrees and then with Flax NNX and Optax

Course content

13 sections • 115 lectures • 11h 32m total length

  • Introduction to the Course9:07

Requirements

  • Python: functions, classes, decorators and comprehensions.
  • NumPy basics: arrays, shapes, broadcasting and indexing.
  • No compiler background is assumed.
  • No GPU is required. Everything in the course runs on an ordinary laptop CPU.
  • Helpful but not required: knowing what a loss function and a gradient step are.

Description

Welcome to this course on JAX!

We are going to learn one of the best and most powerful Python libraries ever created: we will see how JAX can speed up our Python code by orders of magnitude if we use it properly!

This course is not simply a tour of the JAX API. The official documentation already does that well.

Here we go one level deeper.

What we do together

We follow a Python function all the way down, and then we climb back up:

  • Tracing and the Jaxpr: the one mechanism every JAX feature is built on

  • Just-in-time compilation: what gets cached, what triggers a recompile, and how to time JAX code correctly

  • The XLA backend: StableHLO, operator fusion and layout assignment, observed on real code

  • Asking the compiler what a function costs before running it, and exporting a model

  • Defining a new primitive from scratch, with its evaluation, lowering, differentiation, and batching rules

  • Automatic differentiation: grad, Jacobians, Hessians, forward mode, and gradients through control flow

  • Vectorization with vmap, including per-sample gradients

  • Parallel computation with pmap and the modern sharding API

  • Numerical computing: immutable arrays, explicit random keys and the silent differences from NumPy

  • Neural networks, first in pure JAX and then with Flax NNX and Optax

By the end, you can follow a Python function down to a compiled kernel and say what happened at every step. That is a different skill from knowing the API, and it is the one that makes JAX predictable.

How each chapter is built

An introduction that frames the problem. Notebooks where every line of code is written on screen as it is explained. A conclusion that connects to the next chapter. A short quiz with the reasoning behind each answer. The notebooks are in a public GitHub repository, linked in the lecture resources.

What you should know before buying

  • You need Python and NumPy. We do not review either one.

  • No compiler background is assumed, and no GPU is needed. Everything runs on an ordinary laptop CPU, and the parallel computation chapter simulates multiple devices on that CPU.

  • Machine learning is not the goal of this course. One chapter builds a neural network, and it is easier if you already know what a loss function and a gradient step are.

  • The chapters build on each other. This is a course to follow in order, not a reference to dip into.

  • The lectures were recorded against pinned library versions (JAX 0.10.0). JAX changes quickly, so the notebooks in the repository are kept up to date, while the videos are not re-recorded at every release. The concepts do not change.

Who is teaching

I am Pietropaolo Frisoni, a computational scientist working as a Quantum Software Engineer, with a background in theoretical physics and scientific computing. I have used JAX for several years, applying tracing and just-in-time compilation to quantum computing. This is the course I wish I had when I started with JAX.

Take a look at the curriculum and the free preview lectures.

If going one level deeper sounds like your kind of fun, I will see you inside!

Who this course is for:

  • Researchers who use JAX every day and want to know what actually happens below the API
  • Engineers evaluating JAX against PyTorch or TensorFlow, who need more than benchmark charts to decide
  • Python developers curious about how a modern machine learning compiler stack works
  • People who use JAX outside machine learning, for example in physics or scientific computing. Most of the course has nothing to do with ML
  • Not the right course if you only want a quick tour of the API to get a model training. The official documentation covers that well

Report abuse