Ship and adapt JAX workloads

Who this path is for

Model builders who want to package reliable inference on servers or edge devices.

Preparation

Shared foundations and a working model

Your first useful artifact

An exported computation checked against its source on the same inputs.

How to study

Begin with one saved model and a fixed set of parity inputs. Trace preprocessing, parameter layout and output semantics across each framework boundary. Separate training precision, exported precision and the target runtime before benchmarking. Use the precision and edge lessons to build a validation checklist; measure a real target device before claiming device latency or memory savings.

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
6. Neural network training

   Build the same ideas into small neural networks. Understand hidden layers, parameter state, stable objectives and evaluation before moving to a small image classifier.

   • Build a tiny multilayer perceptron
   • Model and state with Flax NNX
   • A compiled train and evaluation step
   • Train a CNN on a small dataset
   • Debug unstable learning
7. Data & checkpoint recovery

   Follow examples from the dataset into a training update, then save enough information to reproduce the next update after interruption. Track experiment configuration, data identity and metric histories with MLflow alongside complete checkpoints.

   • Design an input pipeline
   • Load, batch, and checkpoint order with Grain
   • Save model, optimizer, and random state
   • Recover data position and resume
   • Asynchronous checkpoints and failure boundaries
   • Track experiments and lineage with MLflow
8. Performance diagnosis

   Measure completed work, inspect a real profiler trace and relate compiler output to memory traffic. Compare rematerialization choices while separating saved residuals, compiler memory estimates and observed latency.

   • Benchmark asynchronous work correctly
   • Diagnose recompilation and host synchronization
   • Read an XProf trace
   • Reason with HLO and the roofline model
   • Memory, rematerialization, and optimization tradeoffs
9. Deployment, interoperability & edge AI

   Adapt a trained model, verify framework and export contracts, compare precision policies and measure local inference before planning server or edge capacity. Package a model service in a reproducible container and verify readiness, artifact identity and request contracts.

   • Adapt a pretrained model and choose a post-training objective
   • Move models between JAX, Keras, TensorFlow and PyTorch
   • Convert PyTorch weights to Flax and locate numerical errors
   • Export a computation and verify its serving contract
   • Choose weight, activation and accumulation precision
   • Containerize a model service and verify its boundary
   • Inference capacity, batching, and autoscaling
   • Deploy at the edge: conversion, budgets and device checks

Verify an inference artifact and capacity hypothesis

A guided project is available for this route.

Audit adaptation, exported precision artifacts and a validated local request boundary. Derive a capacity hypothesis from measured execution and explicitly simulated arrivals; network serving and target-device measurements remain extensions.

• Saved-model and exported-runtime parity on a fixed input contract
• Weight, activation and accumulation policies with accuracy and range checks
• Completed inference timings and a workload-based capacity argument, identifying unmeasured target runtimes

Read the portfolio project guide

Work through the synthesis assessment

Use this path in your learning workspace