JAX PathwaysHomeCourse catalog
Your route through JAX.
Open a lesson. Build something you can explain.
22 phases · 103 lessons available
Open a step to see its lessons. You can start anywhere.
Start Here · CPU Setup & Optional TPU Track
Start with the local CPU environment. Add the optional TPU specialization when you need to qualify accelerator execution.
00Onboarding · Local CPUSetup & first steps
Start here4 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
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.
Keep a first-run log that another learner can reproduce, then pair each phase's CPU baseline with the parallel TPU launch and precision guides.
Phase study guide: Can someone else reproduce your first result?
Start after: No previous JAX experience needed. Begin with the guided setup.
Hardware and future work
CPU; the TPU bridge is optional accelerator practice
- 00.01Set up your learning workspaceCPU practice · 40 min
- 00.02Meet your arrays and devicesCPU practice · 45 min
- 00.03Practice with four virtual CPU devicesLogical CPU devices · 60 min
- 00.04Move your experiment to a TPUCPU practice · 45 min
Bring the phase together · Integration Project
Check the CPU environment and a saved matrix arithmetic identity
Checkpoint: Explain which device ran your code and reproduce the result in a fresh environment.
What to do & where to put your code: Write your implementation in projects/welcome-audit/starter/model.py (copied from the starter scaffold in projects/welcome-audit/starter/) and run projects/welcome-audit/tests/check.py from the top-level folder to verify each stage.
Open the staged project · stages 1, 2, 3, 4
Work in projects/welcome-audit/starter/model.py. Complete stages 1 and 2 for independent calculations and changed-condition diagnosis, then retain the stage 3 evidence with --evidence-dir runs/welcome-audit. Stage 4 integrates the checks. Learner commands explicitly select the starter; --implementation solution is instructor reference execution. This bounded practice does not certify the entire phase or target hardware.
Additional project: Connect arrays, transformations and explicit state
19Cloud TPU Mini-Course · Step 1 of 3TPU Setup: Provisioning & Runtime Verification
Start here2 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Step 1 of the 3-part Cloud TPU mini-course. Choose the right TPU slice, provision a Cloud TPU VM (or Colab TPU), bootstrap ~/jax-tpu-lab/.venv with jax[tpu], verify a synchronized pmap runtime receipt, and practice clean teardown.
Complete this short 2-lesson setup course right after 00: Setup & first steps, then continue directly to 20: TPU Workflows and 21: TPU Systems.
Start after: 00 Setup & first steps
Hardware and future work
CPU topology/contract verifier; Cloud TPU VM (v5litepod-4 / v6e-4) or Colab TPU for target runs
- 19.01Select a TPU slice and provision your Cloud TPU VMCPU practice · 25 min
- 19.02Bootstrap .venv and verify a TPU runtime receiptCPU practice · 20 min
Bring the phase together · Integration Project
Inspect CPU device inventory, BF16 storage, and collective participant arithmetic
Checkpoint: Select a single-host TPU topology, provision a Cloud TPU VM (or Colab TPU), bootstrap jax[tpu] in an isolated .venv, verify pmap execution, and confirm clean VM deletion.
What to do & where to put your code: Write your implementation in projects/tpu-runtime-audit/starter/model.py (copied from the starter scaffold in projects/tpu-runtime-audit/starter/) and run projects/tpu-runtime-audit/tests/check.py from the top-level folder to verify each stage.
Open the staged project · stages 1, 2, 3, 4
Work in projects/tpu-runtime-audit/starter/model.py. Complete stages 1 and 2 for independent calculations and changed-condition diagnosis, then retain the stage 3 evidence with --evidence-dir runs/tpu-runtime-audit. Stage 4 integrates the checks. Learner commands explicitly select the starter; --implementation solution is instructor reference execution. This bounded practice does not certify the entire phase or target hardware.
Additional project: Operate and recover a local JAX workload
20Cloud TPU Mini-Course · Step 2 of 3TPU Workflows: Launching Jobs & Checkpointed Experiments
Start here2 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Step 2 of the 3-part Cloud TPU mini-course. Stage scripts with gcloud compute tpus tpu-vm scp, separate XLA compilation warmup from steady-state step timing, recover interrupted runs from atomic checkpoints, and launch single-host or --worker=all Pod slice experiments.
Use the staging pattern from this course to run any lesson script in public/exercises/ on your Cloud TPU VM alongside your local CPU checks.
Start after: 19 TPU Setup: Provisioning & Runtime Verification
Hardware and future work
CPU launcher rehearsal; Cloud TPU VM (single-host or --worker=all Pod slice) for target execution
- 20.01Stage and launch checkpointed TPU jobsCPU practice · 25 min
- 20.02Run course scripts and multi-host TPU Pod slicesCPU practice · 20 min
Bring the phase together · Integration Project
Gate simulated multi-worker checkpoint publication on complete receipts
Checkpoint: Stage scripts with gcloud compute tpus tpu-vm scp, separate XLA warmup from timed step loops, verify zero-drift checkpoint resume, and build single-host and --worker=all Pod launch commands.
What to do & where to put your code: Write your implementation in projects/tpu-workflow-audit/starter/model.py (copied from the starter scaffold in projects/tpu-workflow-audit/starter/) and run projects/tpu-workflow-audit/tests/check.py from the top-level folder to verify each stage.
Open the staged project · stages 1, 2, 3, 4
Work in projects/tpu-workflow-audit/starter/model.py. Complete stages 1 and 2 for independent calculations and changed-condition diagnosis, then retain the stage 3 evidence with --evidence-dir runs/tpu-workflow-audit. Stage 4 integrates the checks. Learner commands explicitly select the starter; --implementation solution is instructor reference execution. This bounded practice does not certify the entire phase or target hardware.
Additional project: Operate and recover a local JAX workload
21Cloud TPU Mini-Course · Step 3 of 3TPU Systems: Generations, Memory, Precision & XProf
Start here2 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Step 3 of the 3-part Cloud TPU mini-course. Compare TPU v5e, v5p, and v6e specifications and ridge points, budget training and KV-cache HBM before provisioning, measure BF16 and INT8 numerical error, and capture warmed XProf traces.
After completing this third TPU mini-course, you are ready to run any accelerator phase (Transformers, Performance, Distributed, Kernels, Deployment, Operations, Pretraining, Post-training) on a Cloud TPU VM.
Start after: 20 TPU Workflows: Launching Jobs & Checkpointed Experiments
Hardware and future work
CPU numerical & profiler harness; Cloud TPU VM for target XProf traces
- 21.01TPU generations and HBM memory budgetingCPU practice · 25 min
- 21.02bfloat16, int8 precision, and XProf profiling on TPUCPU practice · 25 min
Bring the phase together · Integration Project
Estimate optimizer-state bytes and calculate supplied-rate scaling efficiency
Checkpoint: Compare TPU v5e, v5p, and v6e ridge points, budget training and KV-cache HBM, measure BF16/INT8 matmul error with FP32 accumulation, and capture a warmed XProf trace.
What to do & where to put your code: Write your implementation in projects/tpu-systems-audit/starter/model.py (copied from the starter scaffold in projects/tpu-systems-audit/starter/) and run projects/tpu-systems-audit/tests/check.py from the top-level folder to verify each stage.
Open the staged project · stages 1, 2, 3, 4
Work in projects/tpu-systems-audit/starter/model.py. Complete stages 1 and 2 for independent calculations and changed-condition diagnosis, then retain the stage 3 evidence with --evidence-dir runs/tpu-systems-audit. Stage 4 integrates the checks. Learner commands explicitly select the starter; --implementation solution is instructor reference execution. This bounded practice does not certify the entire phase or target hardware.
Additional project: Operate and recover a local JAX workload
Part I · Shared CPU Foundations
Master arrays, pure functions, jit/grad/vmap transforms, explicit state, and numerical optimization on CPU.
01Core Foundation · Local CPUArrays & pure functions
Shared foundations4 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Treat every axis as part of a contract. Work through normalization, broadcasting and immutable updates, then diagnose mistakes that still return plausible numbers.
Draw the input and output shapes before running each change.
Phase study guide: Can you spot a wrong axis even when the program runs?
Start after: 00 Setup & first steps
Hardware and future work
CPU
- 01.01From NumPy to jax.numpyCPU practice · 45 min
- 01.02Shapes, broadcasting, and dtypesCPU practice · 50 min
- 01.03Pure functions and explicit inputsCPU practice · 45 min
- 01.04Immutable updates and indexingCPU practice · 45 min
Bring the phase together · Integration Project
Audit masked batch moments and padding-sensitive denominators
Checkpoint: Predict output shapes and explain why a pure function is easier to transform.
What to do & where to put your code: Write your implementation in projects/arrays-audit/starter/model.py (copied from the starter scaffold in projects/arrays-audit/starter/) and run projects/arrays-audit/tests/check.py from the top-level folder to verify each stage.
Open the staged project · stages 1, 2, 3, 4
Work in projects/arrays-audit/starter/model.py. Complete stages 1 and 2 for independent calculations and changed-condition diagnosis, then retain the stage 3 evidence with --evidence-dir runs/arrays-audit. Stage 4 integrates the checks. Learner commands explicitly select the starter; --implementation solution is instructor reference execution. This bounded practice does not certify the entire phase or target hardware.
Additional project: Connect arrays, transformations and explicit state
02Core Foundation · Local CPUJAX transformations
Shared foundations5 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
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.
Keep eager and transformed outputs plus a deliberately broken assumption.
Phase study guide: What changes when you differentiate, batch or compile a function?
Start after: 01 Arrays & pure functions
Hardware and future work
CPU
- 02.01Your first gradientCPU practice · 55 min
- 02.02Losses and value_and_gradCPU practice · 55 min
- 02.03Batch a function with vmapCPU practice · 60 min
- 02.04Compile a function with jitCPU practice · 55 min
- 02.05Tracing, static arguments, and recompilationCPU practice · 60 min
Bring the phase together · Integration Project
Compare scalar, vmap, autodiff, and changed-shape JIT behavior
Checkpoint: Compose grad, vmap, and jit; explain what each transformation changes.
What to do & where to put your code: Write your implementation in projects/transforms-audit/starter/model.py (copied from the starter scaffold in projects/transforms-audit/starter/) and run projects/transforms-audit/tests/check.py from the top-level folder to verify each stage.
Open the staged project · stages 1, 2, 3, 4
Work in projects/transforms-audit/starter/model.py. Complete stages 1 and 2 for independent calculations and changed-condition diagnosis, then retain the stage 3 evidence with --evidence-dir runs/transforms-audit. Stage 4 integrates the checks. Learner commands explicitly select the starter; --implementation solution is instructor reference execution. This bounded practice does not certify the entire phase or target hardware.
Additional project: Connect arrays, transformations and explicit state
03Core Foundation · Local CPUState, randomness & control flow
Shared foundations4 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Make randomness, parameters and loop state explicit. Learn how to repeat an experiment without accidentally reusing randomness or mutating a caller’s state.
Record the state before and after a transition and explain every changed leaf.
Phase study guide: Have you saved everything needed to determine the next step?
Start after: 02 JAX transformations
Hardware and future work
CPU
- 03.01Random keys without surprisesCPU practice · 60 min
- 03.02Pytrees and structured parametersCPU practice · 60 min
- 03.03Compiled loops with lax.scanCPU practice · 65 min
- 03.04Branches with lax.condCPU practice · 60 min
Bring the phase together · Integration Project
A deterministic batched simulation
Checkpoint: Replay the same experiment with saved state and keys; explain the update order.
What to do & where to put your code: Write your implementation in projects/foundation-toolkit/my_toolkit.py (copied from the starter scaffold in projects/foundation-toolkit/starter/) and run projects/foundation-toolkit/tests/check.py from the top-level folder to verify each stage.
Open the staged project · stages 4
Copy projects/foundation-toolkit/starter/toolkit.py to projects/foundation-toolkit/my_toolkit.py and write your code in projects/foundation-toolkit/my_toolkit.py. Run python3 projects/foundation-toolkit/tests/check.py --implementation projects/foundation-toolkit/my_toolkit.py --stage 4 from the top-level folder to verify stage 4.
04Core Foundation · Local CPUMath & optimization
Shared foundations12 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
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.
For the initial regression project, focus on Linear algebra and loss intuition, Gradient checking and numerical accuracy, Write gradient descent yourself, and Optimize with Optax. The other eight lessons deepen the mathematical explanation; return to them as needed.
Phase study guide: Can you explain an optimizer’s behavior from the objective and data?
Direct links for the first project: Linear algebra and loss intuitionGradient checking and numerical accuracyWrite gradient descent yourselfOptimize with Optax
Start after: 03 State, randomness & control flow
Hardware and future work
CPU
- 04.01Vectors, norms, and projectionsCPU practice · 60 min
- 04.02Linear algebra and loss intuitionCPU practice · 80 min
- 04.03Least squares, rank, and conditioningCPU practice · 75 min
- 04.04The chain rule, Jacobians, and directionsCPU practice · 80 min
- 04.05Gradient checking and numerical accuracyCPU practice · 80 min
- 04.06Curvature, Hessians, and learning ratesCPU practice · 75 min
- 04.07Write gradient descent yourselfCPU practice · 80 min
- 04.08Probability, likelihood, and stable lossesCPU practice · 80 min
- 04.09Regularization and honest validationCPU practice · 75 min
- 04.10Minibatches, expectation, and gradient noiseCPU practice · 75 min
- 04.11Optimize with OptaxCPU practice · 80 min
- 04.12Adam, clipping, and learning-rate schedulesCPU practice · 90 min
Bring the phase together · Integration Project
A fitted regression model
Checkpoint: Explain prediction geometry and loss assumptions; verify derivatives; diagnose conditioning, curvature and gradient noise; justify an optimizer using held-out evidence.
What to do & where to put your code: Write your implementation in projects/optimizer-audit/my_optimization.py (copied from the starter scaffold in projects/optimizer-audit/starter/) and run projects/optimizer-audit/tests/check.py from the top-level folder to verify each stage.
Open the staged project · stages 1, 2, 3, 4, 5
Complete the regression audit after the four core lessons. This deeper project follows all twelve optimization lessons and connects curvature, conditioning, gradient noise, optimizer state and regularization. Copy projects/optimizer-audit/starter/optimization.py to projects/optimizer-audit/my_optimization.py and write your code in projects/optimizer-audit/my_optimization.py. Run python3 projects/optimizer-audit/tests/check.py --implementation projects/optimizer-audit/my_optimization.py --stage 1 from the top-level folder to verify stages 1, 2, 3, 4, 5.
Part II · Complete Training Systems
Build neural networks, checkpoint recovery, Transformers, pretraining, post-training, and performance profiling.
05Training Systems · Local CPUNeural network training
Training systems5 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Build the same ideas into small neural networks. Understand hidden layers, parameter state, stable objectives and evaluation before moving to a small image classifier.
Explain a loss curve together with held-out errors; a lower training loss alone is not sufficient.
Phase study guide: Did the network learn the task, or exploit an evaluation mistake?
Start after: 04 Math & optimization
Hardware and future work
CPU first; accelerator for larger experiments
- 05.01Build a tiny multilayer perceptronCPU practice · 85 min
- 05.02Model and state with Flax NNXCPU practice · 80 min
- 05.03A compiled train and evaluation stepCPU practice · 95 min
- 05.04Train a CNN on a small datasetCPU practice · 110 min
- 05.05Debug unstable learningCPU practice · 85 min
Bring the phase together · Integration Project
Build and audit a neural classifier
Checkpoint: Keep train and evaluation behavior separate and reproduce metrics from a saved configuration.
What to do & where to put your code: Write your implementation in projects/mlp-classifier/my_model.py (copied from the starter scaffold in projects/mlp-classifier/starter/) and run projects/mlp-classifier/tests/check.py from the top-level folder to verify each stage.
Copy projects/mlp-classifier/starter/model.py to projects/mlp-classifier/my_model.py and write your code in projects/mlp-classifier/my_model.py. Run python3 projects/mlp-classifier/tests/check.py --implementation projects/mlp-classifier/my_model.py --stage 1 from the top-level folder to verify all 3 stages.
Additional project: Reproducible JAX/Flax experiments workshop
06Training Systems · Local CPUData & checkpoint recovery
Training systems6 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
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.
Compare sample IDs, optimizer memory, random state and losses after restore.
Phase study guide: Can an interrupted run make exactly the same next update?
Start after: 05 Neural network training
Hardware and future work
CPU first; TPU for system measurements
- 06.01Design an input pipelineCPU practice · 70 min
- 06.02Load, batch, and checkpoint order with GrainCPU practice · 75 min
- 06.03Save model, optimizer, and random stateCPU practice · 100 min
- 06.04Recover data position and resumeCPU practice · 110 min
- 06.05Asynchronous checkpoints and failure boundariesCPU practice · 90 min
- 06.06Track experiments and lineage with MLflowCPU practice · 70 min
Bring the phase together · Integration Project
Resume optimizer and PRNG state across an interrupted training update
Checkpoint: Compare resumed and uninterrupted runs under a stated determinism tolerance.
What to do & where to put your code: Write your implementation in projects/recovery-audit/starter/model.py (copied from the starter scaffold in projects/recovery-audit/starter/) and run projects/recovery-audit/tests/check.py from the top-level folder to verify each stage.
Open the staged project · stages 1, 2, 3, 4
Work in projects/recovery-audit/starter/model.py. Complete stages 1 and 2 for independent calculations and changed-condition diagnosis, then retain the stage 3 evidence with --evidence-dir runs/recovery-audit. Stage 4 integrates the checks. Learner commands explicitly select the starter; --implementation solution is instructor reference execution. This bounded practice does not certify the entire phase or target hardware.
Additional project: Ship a tracked and containerized model release
Additional project: Resume and profile a sharded training run
Additional project: Reproducible JAX/Flax experiments workshop
07Training Systems · Local CPUTransformers & language models
Training systems4 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Build attention and causal masking from small arrays, then connect tokenization, mixed precision, training and generation.
Inspect what each token can attend to and count only valid targets in the loss.
Phase study guide: Can you prove that a predicted token cannot see its future?
Start after: 06 Data & checkpoint recovery
Hardware and future work
CPU for small blocks; TPU or GPU for training experiments
- 07.01Attention from small piecesCPU practice · 90 min
- 07.02Masking, tokenization, and sequence packingCPU practice · 90 min
- 07.03A Transformer block and mixed precisionCPU practice · 90 min
- 07.04Train, checkpoint, and generateCPU practice · 120 min
Bring the phase together · Integration Project
A checkpointed causal language model
Checkpoint: Verify masking, sequence handling, restoration, and generation before comparing speed.
What to do & where to put your code: Write your implementation in projects/text-harness/my_model.py (copied from the starter scaffold in projects/text-harness/starter/) and run projects/text-harness/tests/check.py from the top-level folder to verify each stage.
Open the staged project · stages 1, 2, 3, 4
Begin with token contracts and training. Complete the linked performance and deployment prerequisites before cache, export and profiling checks. Copy projects/text-harness/starter/model.py to projects/text-harness/my_model.py and write your code in projects/text-harness/my_model.py. Run python3 projects/text-harness/tests/check.py --implementation projects/text-harness/my_model.py --stage 1 from the top-level folder to verify stages 1, 2, 3, 4.
Additional project: Trace a tiny bidirectional encoder from pretraining to retrieval
08Training Systems · Local CPUPerformance diagnosis
Training systems5 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
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.
Keep workload shapes, warm-up policy and synchronized timings with your device report.
Phase study guide: What measured evidence identifies the bottleneck?
Start after: 05 Neural network training
Hardware and future work
CPU for timing concepts; TPU or GPU for profiling
- 08.01Benchmark asynchronous work correctlyCPU practice · 80 min
- 08.02Diagnose recompilation and host synchronizationCPU practice · 80 min
- 08.03Read an XProf traceCPU practice · 100 min
- 08.04Reason with HLO and the roofline modelCPU practice · 95 min
- 08.05Memory, rematerialization, and optimization tradeoffsCPU practice · 100 min
Bring the phase together · Integration Project
Measure CPU JIT latency and changed-shape traces with retained JAXPR evidence
Checkpoint: Separate compilation, execution, input stalls, and synchronization before claiming a speedup.
What to do & where to put your code: Write your implementation in projects/performance-audit/starter/model.py (copied from the starter scaffold in projects/performance-audit/starter/) and run projects/performance-audit/tests/check.py from the top-level folder to verify each stage.
Open the staged project · stages 1, 2, 3, 4
Work in projects/performance-audit/starter/model.py. Complete stages 1 and 2 for independent calculations and changed-condition diagnosis, then retain the stage 3 evidence with --evidence-dir runs/performance-audit. Stage 4 integrates the checks. Learner commands explicitly select the starter; --implementation solution is instructor reference execution. This bounded practice does not certify the entire phase or target hardware.
Additional project: Resume and profile a sharded training run
17Training Systems · Local CPUSelf-supervised pretraining: masked and contrastive learning
Training systems3 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Train small masked-token, hidden-patch and paired-view models. Audit information leakage, reduction counts and pair identity before interpreting a falling objective.
Predict the mask, denominator or preference direction before training. Keep independent objective checks and changed-condition evidence.
Phase study guide: What learning signal can unlabelled data provide?
Start after: 07 Transformers & language models
Hardware and future work
Executed CPU synthetic objectives; large-scale and human-feedback evaluation remain separate
- 17.01Masked language modeling: predict hidden tokensCPU practice · 105 min
- 17.02Masked image modeling: reconstruct missing patchesCPU practice · 105 min
- 17.03Contrastive learning: views, positives and negativesCPU practice · 105 min
Bring the phase together · Integration Project
An audited pretraining and adaptation objective toolkit
Checkpoint: Complete stages 1–3 of training-methods with independent loss references, real training curves, changed masks or pairs and explicit synthetic-data limits.
What to do & where to put your code: Write your implementation in projects/training-methods/my_methods.py (copied from the starter scaffold in projects/training-methods/starter/) and run projects/training-methods/tests/check.py from the top-level folder to verify each stage.
Open the staged project · stages 1, 2, 3
Copy projects/training-methods/starter/methods.py to projects/training-methods/my_methods.py and write your code in projects/training-methods/my_methods.py. Run python3 projects/training-methods/tests/check.py --implementation projects/training-methods/my_methods.py --stage 1 from the top-level folder to verify stages 1, 2, 3.
Additional project: Trace a tiny bidirectional encoder from pretraining to retrieval
18Training Systems · Local CPUPost-training: SFT, LoRA, reward models and RLHF
Training systems5 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Connect supervised demonstrations, frozen-base LoRA, reward modeling, sampled PPO and direct preference optimization. Keep the reference, behavior policy, reward model and dataset identities separate.
Predict the mask, denominator or preference direction before training. Keep independent objective checks and changed-condition evidence.
Phase study guide: Which parameters and signals change during adaptation?
Start after: 17 Self-supervised pretraining: masked and contrastive learning
Hardware and future work
Executed CPU synthetic objectives; large-scale and human-feedback evaluation remain separate
- 18.01Supervised fine-tuning with response-only token lossCPU practice · 105 min
- 18.02LoRA: adapt, save and merge low-rank updatesCPU practice · 105 min
- 18.03Learn a reward model from pairwise preferencesCPU practice · 105 min
- 18.04RLHF mechanics: a frozen reward and PPO policy updatesCPU practice · 105 min
- 18.05Direct preference optimization and reference-corrected marginsCPU practice · 105 min
Bring the phase together · Integration Project
Audit response-token SFT masking and preference-pair loss
Checkpoint: Complete stages 4–8 of training-methods. Retain shifted token-mask checks, frozen-base/merge evidence, preference oracles, signed PPO clipping and fixed-reference DPO comparisons.
What to do & where to put your code: Write your implementation in projects/posttraining-audit/starter/model.py (copied from the starter scaffold in projects/posttraining-audit/starter/) and run projects/posttraining-audit/tests/check.py from the top-level folder to verify each stage.
Open the staged project · stages 1, 2, 3, 4
Work in projects/posttraining-audit/starter/model.py. Complete stages 1 and 2 for independent calculations and changed-condition diagnosis, then retain the stage 3 evidence with --evidence-dir runs/posttraining-audit. Stage 4 integrates the checks. Learner commands explicitly select the starter; --implementation solution is instructor reference execution. This bounded practice does not certify the entire phase or target hardware.
Additional project: Audit pretraining and post-training objectives
Part III · Accelerator Scale, Kernels, Deployment & Operations
Partition across device meshes, author Pallas TPU kernels, ship inference services, and operate production workloads.
09Systems Specialization · CPU ReferenceDistributed training
Specializations4 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Partition arrays and training updates, compare all-reduce with reduce-scatter, and restore a complete sharded state across an epoch boundary. Verify CPU placement and arithmetic before separate cluster validation.
Compare against a single-device reference; measure real accelerator performance separately.
Phase study guide: Is the distributed update still optimizing the same objective?
Start after: 06 Data & checkpoint recovery08 Performance diagnosis
Hardware and future work
Logical CPU devices for sharding practice; multi-host work needs a real cluster
- 09.01Arrays, meshes, and shardingLogical CPU devices · 80 min
- 09.02A sharded training stepLogical CPU devices · 90 min
- 09.03Communication-efficient algorithmsLogical CPU devices · 100 min
- 09.04Resilient distributed trainingLogical CPU devices · 105 min
Bring the phase together · Integration Project
A resilient sharded training experiment
Checkpoint: Explain partitioning, communication costs, and recovery behavior with measured evidence.
What to do & where to put your code: Write your implementation in projects/sharded-training/my_model.py (copied from the starter scaffold in projects/sharded-training/starter/) and run projects/sharded-training/tests/check.py from the top-level folder to verify each stage.
Open the staged project · stages 1, 2, 3, 4
This connected project also uses performance and distributed lessons. Work through its prerequisites before the full integration check; return here with the completed evidence. Copy projects/sharded-training/starter/model.py to projects/sharded-training/my_model.py and write your code in projects/sharded-training/my_model.py. Run python3 projects/sharded-training/tests/check.py --implementation projects/sharded-training/my_model.py --stage 1 from the top-level folder to verify stages 1, 2, 3, 4.
13Systems Specialization · CPU ReferencePallas kernels
Specializations4 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Write actual Pallas grids and buffered pipelines, verify boundaries and precision on CPU, and prepare guarded target measurements before claiming an optimization win.
Begin with a profiled operation and independent oracle. Complete CPU semantics first; actual TPU compilation, correctness and performance require the separate target runner on supported hardware.
Phase study guide: Does a custom kernel earn its complexity?
Start after: 09 Distributed training
Hardware and future work
CPU interpretation/simulation reference; TPU target lab requires separately validated hardware
- 13.01Pallas grids and BlockSpecsCPU practice · 95 min
- 13.02A first TPU kernelCPU practice · 110 min
- 13.03Tiling, memory, and pipeliningCPU practice · 120 min
- 13.04Iterate with correctness and performance evidenceCPU practice · 120 min
Bring the phase together · Integration Project
A checked Pallas kernel with an explicit target qualification plan
Checkpoint: Compare against a reference across shapes and justify each optimization with evidence.
What to do & where to put your code: Write your implementation in projects/kernel-audit/my_model.py (copied from the starter scaffold in projects/kernel-audit/starter/) and run projects/kernel-audit/tests/check.py from the top-level folder to verify each stage.
Copy projects/kernel-audit/starter/model.py to projects/kernel-audit/my_model.py and write your code in projects/kernel-audit/my_model.py. Run python3 projects/kernel-audit/tests/check.py --implementation projects/kernel-audit/my_model.py --stage 1 from the top-level folder to verify all 3 stages.
15Systems Specialization · CPU ReferenceDeployment, interoperability & edge AI
Specializations8 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
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.
Preserve checkpoint, data and calibration provenance. Compare held-out inputs before and after conversion; separate local request timing from queue simulation and target-device evidence.
Phase study guide: Does the exported service preserve the model’s complete contract?
Start after: 06 Data & checkpoint recovery08 Performance diagnosis
Hardware and future work
CPU for core labs; optional framework environments and target devices for conversion and edge validation
- 15.01Adapt a pretrained model and choose a post-training objectiveCPU practice · 130 min
- 15.02Move models between JAX, Keras, TensorFlow and PyTorchCPU practice · 110 min
- 15.03Convert PyTorch weights to Flax and locate numerical errorsCPU practice · 130 min
- 15.04Export a computation and verify its serving contractCPU practice · 110 min
- 15.05Choose weight, activation and accumulation precisionCPU practice · 110 min
- 15.06Containerize a model service and verify its boundaryCPU practice · 110 min
- 15.07Inference capacity, batching, and autoscalingCPU practice · 130 min
- 15.08Deploy at the edge: conversion, budgets and device checksCPU practice · 110 min
Bring the phase together · Integration Project
Verify an inference artifact and capacity hypothesis
Checkpoint: Verify an exported prediction contract, report weight/activation precision error, and distinguish CPU request measurements from target-device evidence.
What to do & where to put your code: Write your implementation in projects/deployment-audit/my_model.py (copied from the starter scaffold in projects/deployment-audit/starter/) and run projects/deployment-audit/tests/check.py from the top-level folder to verify each stage.
Copy projects/deployment-audit/starter/model.py to projects/deployment-audit/my_model.py and write your code in projects/deployment-audit/my_model.py. Run python3 projects/deployment-audit/tests/check.py --implementation projects/deployment-audit/my_model.py --stage 1 from the top-level folder to verify all 3 stages.
Additional project: Ship a tracked and containerized model release
Additional project: Convert PyTorch weights and audit Flax architecture parity
16Systems Specialization · CPU ReferenceWorkload operations
Specializations8 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Launch and observe real local jobs, detect failures and stalls, restore complete training state, verify artifact quality and rehearse rollback before planning capacity. Connect MLOps data/release gates, LLMOps evaluation/traces and ModelOps ownership/approval to the workload runbook.
Use a bounded failure drill and retain its raw event log. Check state and data identity before resuming; separate measured execution from hypothetical demand and pricing.
Phase study guide: Can you detect, explain and safely recover a failed model release?
Start after: 09 Distributed training
Hardware and future work
Executed local CPU processes; accelerator and cluster transfer require separate validation
- 16.01Accelerator jobs and runtime lifecycleCPU practice · 90 min
- 16.02Health, logs, and workload observabilityCPU practice · 85 min
- 16.03Failure recovery and operational runbooksCPU practice · 110 min
- 16.04Changes, rollback, and artifact provenanceCPU practice · 95 min
- 16.05Capacity, utilization, and operating costCPU practice · 90 min
- 16.06MLOps: data contracts, CI gates and monitoringCPU practice · 70 min
- 16.07LLMOps: version prompts, evaluate behavior and trace requestsCPU practice · 75 min
- 16.08ModelOps: own, approve, roll back and retire a modelCPU practice · 65 min
Bring the phase together · Integration Project
An observable, restartable local workload
Checkpoint: Diagnose a failed or stalled job, restore it correctly, and justify its resource plan from recorded measurements.
What to do & where to put your code: Write your implementation in projects/workload-operations/my_model.py (copied from the starter scaffold in projects/workload-operations/starter/) and run projects/workload-operations/tests/check.py from the top-level folder to verify each stage.
Copy projects/workload-operations/starter/model.py to projects/workload-operations/my_model.py and write your code in projects/workload-operations/my_model.py. Run python3 projects/workload-operations/tests/check.py --implementation projects/workload-operations/my_model.py --stage 1 from the top-level folder to verify all 5 stages.
Additional project: Ship a tracked and containerized model release
Part IV · Research & Domain Specializations
Apply differentiable programming to scientific simulations, Bayesian inference, reinforcement learning, and JAX internals.
10Domain Specialization · Local CPUScientific computing
Specializations4 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Build and differentiate a physical simulator, recover a positive rate from observations, and separate numerical accuracy, identifiability and measured execution cost.
Complete gradients, explicit state and scan first. Start with analytic references, check discretization and sensitivity, then fit parameters and defend held-out predictions.
Phase study guide: Are you fitting the physical model or the solver’s error?
Start after: 04 Math & optimization
Hardware and future work
CPU; target accelerator execution requires separate evidence
- 10.01A functional physical simulationCPU practice · 85 min
- 10.02Vectorize trajectories and scan timeCPU practice · 100 min
- 10.03Differentiate through a solverCPU practice · 105 min
- 10.04Fit parameters and scale the simulationCPU practice · 115 min
Bring the phase together · Integration Project
A differentiable inverse problem
Checkpoint: Recover a known physical parameter and report error, stability, and gradient checks.
What to do & where to put your code: Write your implementation in projects/scientific-inverse/my_model.py (copied from the starter scaffold in projects/scientific-inverse/starter/) and run projects/scientific-inverse/tests/check.py from the top-level folder to verify each stage.
Copy projects/scientific-inverse/starter/model.py to projects/scientific-inverse/my_model.py and write your code in projects/scientific-inverse/my_model.py. Run python3 projects/scientific-inverse/tests/check.py --implementation projects/scientific-inverse/my_model.py --stage 1 from the top-level folder to verify all 3 stages.
11Domain Specialization · Local CPUProbabilistic modeling
Specializations4 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Specify probability models, derive exact Bayesian regression, implement and diagnose HMC, and compare variational approximations with predictive evidence.
Start with likelihoods, stable losses and explicit random keys. Check a known posterior before trusting a sampler, then separate inference error from model mismatch.
Phase study guide: Which uncertainty does your model actually represent?
Start after: 04 Math & optimization
Hardware and future work
CPU; transparent JAX samplers with an optional NumPyro ecosystem bridge
- 11.01Distributions and samplingCPU practice · 90 min
- 11.02Build a Bayesian regressionCPU practice · 110 min
- 11.03Hamiltonian Monte Carlo and sampler diagnosticsCPU practice · 140 min
- 11.04Variational inference and model checkingCPU practice · 125 min
Bring the phase together · Integration Project
A Bayesian model with diagnostic evidence
Checkpoint: Report sampler diagnostics and distinguish uncertainty from predictive error.
What to do & where to put your code: Write your implementation in projects/bayesian-regression/my_model.py (copied from the starter scaffold in projects/bayesian-regression/starter/) and run projects/bayesian-regression/tests/check.py from the top-level folder to verify each stage.
Copy projects/bayesian-regression/starter/model.py to projects/bayesian-regression/my_model.py and write your code in projects/bayesian-regression/my_model.py. Run python3 projects/bayesian-regression/tests/check.py --implementation projects/bayesian-regression/my_model.py --stage 1 from the top-level folder to verify all 3 stages.
12Domain Specialization · Local CPUReinforcement learning
Specializations4 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Build a functional environment, collect randomized trajectories, train policy-gradient and clipped-policy agents, and evaluate independent seeds.
Keep terminal and timeout behavior explicit. Verify a small enumerated expectation before trusting a sampled policy gradient, then separate training from held-out evaluation.
Phase study guide: Is the policy improving under a fixed evaluation protocol?
Start after: 05 Neural network training
Hardware and future work
CPU; target accelerator execution requires separate evidence
- 12.01Write a functional environmentCPU practice · 75 min
- 12.02Batch environments and scan rolloutsCPU practice · 85 min
- 12.03Policy gradients and PPOCPU practice · 110 min
- 12.04Evaluate agents across seedsCPU practice · 90 min
Bring the phase together · Integration Project
An evaluated policy in a functional environment
Checkpoint: Compare a trained policy to a baseline across several seeds and report variability.
What to do & where to put your code: Write your implementation in projects/policy-evaluation/my_model.py (copied from the starter scaffold in projects/policy-evaluation/starter/) and run projects/policy-evaluation/tests/check.py from the top-level folder to verify each stage.
Copy projects/policy-evaluation/starter/model.py to projects/policy-evaluation/my_model.py and write your code in projects/policy-evaluation/my_model.py. Run python3 projects/policy-evaluation/tests/check.py --implementation projects/policy-evaluation/my_model.py --stage 1 from the top-level folder to verify all 4 stages.
14Domain Specialization · Local CPUAutodiff & JAX internals
Specializations4 available
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Trace functions into jaxpr, check forward and reverse sensitivities, write exact custom derivatives and compile a small transformation with explicit supported primitives.
Keep an independent derivative reference and test changed values, higher derivatives and unsupported cases before trusting a transformation.
Phase study guide: Which mathematical operation does the transformed program represent?
Start after: 04 Math & optimization
Hardware and future work
CPU
- 14.01Read your first jaxprCPU practice · 90 min
- 14.02JVPs, VJPs, and higher-order derivativesCPU practice · 105 min
- 14.03Custom derivative rulesCPU practice · 120 min
- 14.04Lowering, compilation, and a tiny transformationCPU practice · 130 min
Bring the phase together · Integration Project
Audit derivatives and a tiny transformation
Checkpoint: Explain tracing and verify the derivative rule numerically.
What to do & where to put your code: Write your implementation in projects/derivative-audit/my_model.py (copied from the starter scaffold in projects/derivative-audit/starter/) and run projects/derivative-audit/tests/check.py from the top-level folder to verify each stage.
Copy projects/derivative-audit/starter/model.py to projects/derivative-audit/my_model.py and write your code in projects/derivative-audit/my_model.py. Run python3 projects/derivative-audit/tests/check.py --implementation projects/derivative-audit/my_model.py --stage 1 from the top-level folder to verify all 3 stages.
Jump to a course step
103 lesson drafts with CPU companions.
No available lessons in this view.