Start with JAX

Who this path is for

Python learners who can write functions and use lists; NumPy experience helps but is not required.

Preparation

Python basics and a little NumPy

Your first useful artifact

A regression implementation with an independent gradient check and held-out predictions.

How to study

Begin with setup, arrays, transformations and explicit state, completing one foundation-toolkit stage per phase. In Math & optimization, first complete loss intuition, gradient checking, handwritten descent and Optax; then build the regression project. Return for the deeper math sequence and optimizer-audit to explain conditioning, curvature, noisy updates and regularization. Predict each output before running it, then test your explanation with a changed input.

Course phases and lessons

1. Setup & first steps

   Run a saved program and explain its environment, array output and device placement. Start on one CPU, practice sharding on four logical CPU devices, and learn how to get a TPU running, launch a TPU job, and run course experiments on TPU in parallel from Phase 00 onward.

   • Set up your learning workspace
   • Meet your arrays and devices
   • Practice with four virtual CPU devices
   • Move your experiment to a TPU
2. Arrays & pure functions

   Treat every axis as part of a contract. Work through normalization, broadcasting and immutable updates, then diagnose mistakes that still return plausible numbers.

   • From NumPy to jax.numpy
   • Shapes, broadcasting, and dtypes
   • Pure functions and explicit inputs
   • Immutable updates and indexing
3. JAX transformations

   Begin with a scalar function, then differentiate it, batch it and compile it. Compare each transformed result with a simple reference so you can explain what changed.

   • Your first gradient
   • Losses and value_and_grad
   • Batch a function with vmap
   • Compile a function with jit
   • Tracing, static arguments, and recompilation
4. State, randomness & control flow

   Make randomness, parameters and loop state explicit. Learn how to repeat an experiment without accidentally reusing randomness or mutating a caller’s state.

   • Random keys without surprises
   • Pytrees and structured parameters
   • Compiled loops with lax.scan
   • Branches with lax.cond
5. Math & optimization

   Connect geometry and probability to the updates in a training loop. Begin with linear algebra, gradient checks and gradient descent, then investigate conditioning, curvature, likelihoods, regularization and optimizer behavior.

   • Vectors, norms, and projections
   • Linear algebra and loss intuition
   • Least squares, rank, and conditioning
   • The chain rule, Jacobians, and directions
   • Gradient checking and numerical accuracy
   • Curvature, Hessians, and learning rates
   • Write gradient descent yourself
   • Probability, likelihood, and stable losses
   • Regularization and honest validation
   • Minibatches, expectation, and gradient noise
   • Optimize with Optax
   • Adam, clipping, and learning-rate schedules

A fitted regression model

A guided project is available for this route.

Show a tested loss, numerically checked gradients, and a reproducible fit.

• A shape contract and independent loss and gradient calculations
• A reproducible fit, held-out predictions and a changed-data experiment
• A diagnosis of a deliberately wrong gradient or broadcasting rule

Read the portfolio project guide

Work through the synthesis assessment

Math & optimization synthesis

Use this path in your learning workspace