# JAX Pathways

Hands-on JAX, Flax NNX, Optax, and Cloud TPU engineering course.

- [Catalog](api/v1/catalog.json)
- [Pathways](api/v1/pathways.json)
- [Roles](api/v1/roles.json)
- [API contract](openapi.json)

## Available lessons

- [Set up your learning workspace](lesson.html?lesson=welcome-01) · [Source](api/v1/lessons/welcome-01.json) · phases/00-welcome/01-set-up-your-learning-workspace
- [Meet your arrays and devices](lesson.html?lesson=welcome-02) · [Source](api/v1/lessons/welcome-02.json) · phases/00-welcome/02-meet-your-arrays-and-devices
- [Practice with four virtual CPU devices](lesson.html?lesson=welcome-cpu) · [Source](api/v1/lessons/welcome-cpu.json) · phases/00-welcome/03-virtual-cpu-devices
- [Move your experiment to a TPU](lesson.html?lesson=welcome-03) · [Source](api/v1/lessons/welcome-03.json) · phases/00-welcome/03-move-your-experiment-to-a-tpu
- [From NumPy to jax.numpy](lesson.html?lesson=arrays-01) · [Source](api/v1/lessons/arrays-01.json) · phases/01-arrays/01-from-numpy-to-jax-numpy
- [Shapes, broadcasting, and dtypes](lesson.html?lesson=arrays-02) · [Source](api/v1/lessons/arrays-02.json) · phases/01-arrays/02-shapes-broadcasting-and-dtypes
- [Pure functions and explicit inputs](lesson.html?lesson=arrays-03) · [Source](api/v1/lessons/arrays-03.json) · phases/01-arrays/03-pure-functions-and-explicit-inputs
- [Immutable updates and indexing](lesson.html?lesson=arrays-04) · [Source](api/v1/lessons/arrays-04.json) · phases/01-arrays/04-immutable-updates-and-indexing
- [Your first gradient](lesson.html?lesson=first-gradient) · [Source](api/v1/lessons/first-gradient.json) · phases/02-transforms/01-your-first-gradient
- [Losses and value_and_grad](lesson.html?lesson=transforms-02) · [Source](api/v1/lessons/transforms-02.json) · phases/02-transforms/02-losses-and-value-and-grad
- [Batch a function with vmap](lesson.html?lesson=transforms-03) · [Source](api/v1/lessons/transforms-03.json) · phases/02-transforms/03-batch-a-function-with-vmap
- [Compile a function with jit](lesson.html?lesson=transforms-04) · [Source](api/v1/lessons/transforms-04.json) · phases/02-transforms/04-compile-a-function-with-jit
- [Tracing, static arguments, and recompilation](lesson.html?lesson=transforms-05) · [Source](api/v1/lessons/transforms-05.json) · phases/02-transforms/05-tracing-static-arguments-and-recompilation
- [Random keys without surprises](lesson.html?lesson=state-01) · [Source](api/v1/lessons/state-01.json) · phases/03-state/01-random-keys-without-surprises
- [Pytrees and structured parameters](lesson.html?lesson=state-02) · [Source](api/v1/lessons/state-02.json) · phases/03-state/02-pytrees-and-structured-parameters
- [Compiled loops with lax.scan](lesson.html?lesson=state-03) · [Source](api/v1/lessons/state-03.json) · phases/03-state/03-compiled-loops-with-lax-scan
- [Branches with lax.cond](lesson.html?lesson=state-04) · [Source](api/v1/lessons/state-04.json) · phases/03-state/04-branches-with-lax-cond
- [Vectors, norms, and projections](lesson.html?lesson=optimization-05) · [Source](api/v1/lessons/optimization-05.json) · phases/04-optimization/05-vectors-norms-and-projections
- [Linear algebra and loss intuition](lesson.html?lesson=optimization-01) · [Source](api/v1/lessons/optimization-01.json) · phases/04-optimization/01-linear-algebra-and-loss-intuition
- [Least squares, rank, and conditioning](lesson.html?lesson=optimization-06) · [Source](api/v1/lessons/optimization-06.json) · phases/04-optimization/06-least-squares-rank-and-conditioning
- [The chain rule, Jacobians, and directions](lesson.html?lesson=optimization-07) · [Source](api/v1/lessons/optimization-07.json) · phases/04-optimization/07-chain-rule-jacobians-and-directional-derivatives
- [Gradient checking and numerical accuracy](lesson.html?lesson=optimization-02) · [Source](api/v1/lessons/optimization-02.json) · phases/04-optimization/02-gradient-checking-and-numerical-accuracy
- [Curvature, Hessians, and learning rates](lesson.html?lesson=optimization-08) · [Source](api/v1/lessons/optimization-08.json) · phases/04-optimization/08-curvature-hessians-and-learning-rates
- [Write gradient descent yourself](lesson.html?lesson=optimization-03) · [Source](api/v1/lessons/optimization-03.json) · phases/04-optimization/03-write-gradient-descent-yourself
- [Probability, likelihood, and stable losses](lesson.html?lesson=optimization-09) · [Source](api/v1/lessons/optimization-09.json) · phases/04-optimization/09-probability-likelihood-and-stable-losses
- [Regularization and honest validation](lesson.html?lesson=optimization-10) · [Source](api/v1/lessons/optimization-10.json) · phases/04-optimization/10-regularization-and-validation
- [Minibatches, expectation, and gradient noise](lesson.html?lesson=optimization-11) · [Source](api/v1/lessons/optimization-11.json) · phases/04-optimization/11-minibatches-expectation-and-gradient-noise
- [Optimize with Optax](lesson.html?lesson=optimization-04) · [Source](api/v1/lessons/optimization-04.json) · phases/04-optimization/04-optimize-with-optax
- [Adam, clipping, and learning-rate schedules](lesson.html?lesson=optimization-12) · [Source](api/v1/lessons/optimization-12.json) · phases/04-optimization/12-adam-clipping-and-learning-rate-schedules
- [Build a tiny multilayer perceptron](lesson.html?lesson=networks-01) · [Source](api/v1/lessons/networks-01.json) · phases/05-networks/01-build-a-tiny-multilayer-perceptron
- [Model and state with Flax NNX](lesson.html?lesson=networks-02) · [Source](api/v1/lessons/networks-02.json) · phases/05-networks/02-model-and-state-with-flax-nnx
- [A compiled train and evaluation step](lesson.html?lesson=networks-03) · [Source](api/v1/lessons/networks-03.json) · phases/05-networks/03-a-compiled-train-and-evaluation-step
- [Train a CNN on a small dataset](lesson.html?lesson=networks-04) · [Source](api/v1/lessons/networks-04.json) · phases/05-networks/04-train-a-cnn-on-a-small-dataset
- [Debug unstable learning](lesson.html?lesson=networks-05) · [Source](api/v1/lessons/networks-05.json) · phases/05-networks/05-debug-unstable-learning
- [Design an input pipeline](lesson.html?lesson=recovery-01) · [Source](api/v1/lessons/recovery-01.json) · phases/06-recovery/01-design-an-input-pipeline
- [Load, batch, and checkpoint order with Grain](lesson.html?lesson=recovery-02) · [Source](api/v1/lessons/recovery-02.json) · phases/06-recovery/02-load-batch-and-prefetch-with-grain
- [Save model, optimizer, and random state](lesson.html?lesson=recovery-03) · [Source](api/v1/lessons/recovery-03.json) · phases/06-recovery/03-save-model-optimizer-and-random-state
- [Recover data position and resume](lesson.html?lesson=recovery-04) · [Source](api/v1/lessons/recovery-04.json) · phases/06-recovery/04-recover-data-position-and-resume
- [Asynchronous checkpoints and failure boundaries](lesson.html?lesson=recovery-05) · [Source](api/v1/lessons/recovery-05.json) · phases/06-recovery/05-asynchronous-checkpoints-and-failure-boundaries
- [Track experiments and lineage with MLflow](lesson.html?lesson=recovery-06) · [Source](api/v1/lessons/recovery-06.json) · phases/06-recovery/06-track-experiments-and-lineage-with-mlflow
- [Attention from small pieces](lesson.html?lesson=transformers-01) · [Source](api/v1/lessons/transformers-01.json) · phases/07-transformers/01-attention-from-small-pieces
- [Masking, tokenization, and sequence packing](lesson.html?lesson=transformers-02) · [Source](api/v1/lessons/transformers-02.json) · phases/07-transformers/02-masking-tokenization-and-sequence-packing
- [A Transformer block and mixed precision](lesson.html?lesson=transformers-03) · [Source](api/v1/lessons/transformers-03.json) · phases/07-transformers/03-a-transformer-block-and-mixed-precision
- [Train, checkpoint, and generate](lesson.html?lesson=transformers-04) · [Source](api/v1/lessons/transformers-04.json) · phases/07-transformers/04-train-checkpoint-and-generate
- [Benchmark asynchronous work correctly](lesson.html?lesson=performance-01) · [Source](api/v1/lessons/performance-01.json) · phases/08-performance/01-benchmark-asynchronous-work-correctly
- [Diagnose recompilation and host synchronization](lesson.html?lesson=performance-02) · [Source](api/v1/lessons/performance-02.json) · phases/08-performance/02-diagnose-recompilation-and-host-synchronization
- [Read an XProf trace](lesson.html?lesson=performance-03) · [Source](api/v1/lessons/performance-03.json) · phases/08-performance/03-read-an-xprof-trace
- [Reason with HLO and the roofline model](lesson.html?lesson=performance-04) · [Source](api/v1/lessons/performance-04.json) · phases/08-performance/04-reason-with-hlo-and-the-roofline-model
- [Memory, rematerialization, and optimization tradeoffs](lesson.html?lesson=performance-05) · [Source](api/v1/lessons/performance-05.json) · phases/08-performance/05-memory-rematerialization-and-optimization-tradeoffs
- [Arrays, meshes, and sharding](lesson.html?lesson=distributed-01) · [Source](api/v1/lessons/distributed-01.json) · phases/09-distributed/01-arrays-meshes-and-sharding
- [A sharded training step](lesson.html?lesson=distributed-02) · [Source](api/v1/lessons/distributed-02.json) · phases/09-distributed/02-a-sharded-training-step
- [Communication-efficient algorithms](lesson.html?lesson=distributed-03) · [Source](api/v1/lessons/distributed-03.json) · phases/09-distributed/03-communication-efficient-algorithms
- [Resilient distributed training](lesson.html?lesson=distributed-04) · [Source](api/v1/lessons/distributed-04.json) · phases/09-distributed/04-resilient-distributed-training
- [A functional physical simulation](lesson.html?lesson=science-01) · [Source](api/v1/lessons/science-01.json) · phases/10-science/01-a-functional-physical-simulation
- [Vectorize trajectories and scan time](lesson.html?lesson=science-02) · [Source](api/v1/lessons/science-02.json) · phases/10-science/02-vectorize-trajectories-and-scan-time
- [Differentiate through a solver](lesson.html?lesson=science-03) · [Source](api/v1/lessons/science-03.json) · phases/10-science/03-differentiate-through-a-solver
- [Fit parameters and scale the simulation](lesson.html?lesson=science-04) · [Source](api/v1/lessons/science-04.json) · phases/10-science/04-fit-parameters-and-scale-the-simulation
- [Distributions and sampling](lesson.html?lesson=probability-01) · [Source](api/v1/lessons/probability-01.json) · phases/11-probability/01-distributions-and-sampling
- [Build a Bayesian regression](lesson.html?lesson=probability-02) · [Source](api/v1/lessons/probability-02.json) · phases/11-probability/02-build-a-bayesian-regression
- [Hamiltonian Monte Carlo and sampler diagnostics](lesson.html?lesson=probability-03) · [Source](api/v1/lessons/probability-03.json) · phases/11-probability/03-hamiltonian-monte-carlo-and-sampler-diagnostics
- [Variational inference and model checking](lesson.html?lesson=probability-04) · [Source](api/v1/lessons/probability-04.json) · phases/11-probability/04-variational-inference-and-model-checking
- [Write a functional environment](lesson.html?lesson=rl-01) · [Source](api/v1/lessons/rl-01.json) · phases/12-rl/01-write-a-functional-environment
- [Batch environments and scan rollouts](lesson.html?lesson=rl-02) · [Source](api/v1/lessons/rl-02.json) · phases/12-rl/02-batch-environments-and-scan-rollouts
- [Policy gradients and PPO](lesson.html?lesson=rl-03) · [Source](api/v1/lessons/rl-03.json) · phases/12-rl/03-policy-gradients-and-ppo
- [Evaluate agents across seeds](lesson.html?lesson=rl-04) · [Source](api/v1/lessons/rl-04.json) · phases/12-rl/04-evaluate-agents-across-seeds
- [Pallas grids and BlockSpecs](lesson.html?lesson=kernels-01) · [Source](api/v1/lessons/kernels-01.json) · phases/13-kernels/01-pallas-grids-and-blockspecs
- [A first TPU kernel](lesson.html?lesson=kernels-02) · [Source](api/v1/lessons/kernels-02.json) · phases/13-kernels/02-a-first-tpu-kernel
- [Tiling, memory, and pipelining](lesson.html?lesson=kernels-03) · [Source](api/v1/lessons/kernels-03.json) · phases/13-kernels/03-tiling-memory-and-pipelining
- [Iterate with correctness and performance evidence](lesson.html?lesson=kernels-04) · [Source](api/v1/lessons/kernels-04.json) · phases/13-kernels/04-iterate-with-correctness-and-performance-evidence
- [Read your first jaxpr](lesson.html?lesson=internals-01) · [Source](api/v1/lessons/internals-01.json) · phases/14-internals/01-read-your-first-jaxpr
- [JVPs, VJPs, and higher-order derivatives](lesson.html?lesson=internals-02) · [Source](api/v1/lessons/internals-02.json) · phases/14-internals/02-jvps-vjps-and-higher-order-derivatives
- [Custom derivative rules](lesson.html?lesson=internals-03) · [Source](api/v1/lessons/internals-03.json) · phases/14-internals/03-custom-derivative-rules
- [Lowering, compilation, and a tiny transformation](lesson.html?lesson=internals-04) · [Source](api/v1/lessons/internals-04.json) · phases/14-internals/04-lowering-compilation-and-a-tiny-transformation
- [Adapt a pretrained model and choose a post-training objective](lesson.html?lesson=deployment-02) · [Source](api/v1/lessons/deployment-02.json) · phases/15-deployment/02-adapt-a-pretrained-model-and-choose-a-post-training-objective
- [Move models between JAX, Keras, TensorFlow and PyTorch](lesson.html?lesson=deployment-01) · [Source](api/v1/lessons/deployment-01.json) · phases/15-deployment/01-keras-and-pytorch-bridges-to-explicit-jax
- [Convert PyTorch weights to Flax and locate numerical errors](lesson.html?lesson=deployment-08) · [Source](api/v1/lessons/deployment-08.json) · phases/15-deployment/08-pytorch-to-flax-numerical-parity
- [Export a computation and verify its serving contract](lesson.html?lesson=deployment-03) · [Source](api/v1/lessons/deployment-03.json) · phases/15-deployment/03-export-and-serve-a-trained-computation
- [Choose weight, activation and accumulation precision](lesson.html?lesson=deployment-05) · [Source](api/v1/lessons/deployment-05.json) · phases/15-deployment/05-weight-activation-and-accumulation-precision
- [Containerize a model service and verify its boundary](lesson.html?lesson=deployment-07) · [Source](api/v1/lessons/deployment-07.json) · phases/15-deployment/07-containerize-a-model-service
- [Inference capacity, batching, and autoscaling](lesson.html?lesson=deployment-04) · [Source](api/v1/lessons/deployment-04.json) · phases/15-deployment/04-inference-capacity-batching-and-autoscaling
- [Deploy at the edge: conversion, budgets and device checks](lesson.html?lesson=deployment-06) · [Source](api/v1/lessons/deployment-06.json) · phases/15-deployment/06-edge-ai-conversion-and-device-validation
- [Accelerator jobs and runtime lifecycle](lesson.html?lesson=operations-01) · [Source](api/v1/lessons/operations-01.json) · phases/16-operations/01-accelerator-jobs-and-runtime-lifecycle
- [Health, logs, and workload observability](lesson.html?lesson=operations-02) · [Source](api/v1/lessons/operations-02.json) · phases/16-operations/02-health-logs-and-workload-observability
- [Failure recovery and operational runbooks](lesson.html?lesson=operations-03) · [Source](api/v1/lessons/operations-03.json) · phases/16-operations/03-failure-recovery-and-operational-runbooks
- [Changes, rollback, and artifact provenance](lesson.html?lesson=operations-04) · [Source](api/v1/lessons/operations-04.json) · phases/16-operations/04-changes-rollback-and-artifact-provenance
- [Capacity, utilization, and operating cost](lesson.html?lesson=operations-05) · [Source](api/v1/lessons/operations-05.json) · phases/16-operations/05-capacity-utilization-and-operating-cost
- [MLOps: data contracts, CI gates and monitoring](lesson.html?lesson=operations-06) · [Source](api/v1/lessons/operations-06.json) · phases/16-operations/06-mlops-data-contracts-and-release-pipelines
- [LLMOps: version prompts, evaluate behavior and trace requests](lesson.html?lesson=operations-07) · [Source](api/v1/lessons/operations-07.json) · phases/16-operations/07-llmops-prompts-evaluation-and-traces
- [ModelOps: own, approve, roll back and retire a model](lesson.html?lesson=operations-08) · [Source](api/v1/lessons/operations-08.json) · phases/16-operations/08-modelops-ownership-approval-and-retirement
- [Masked language modeling: predict hidden tokens](lesson.html?lesson=pretraining-01) · [Source](api/v1/lessons/pretraining-01.json) · phases/17-pretraining/01-mlm
- [Masked image modeling: reconstruct missing patches](lesson.html?lesson=pretraining-02) · [Source](api/v1/lessons/pretraining-02.json) · phases/17-pretraining/02-mim
- [Contrastive learning: views, positives and negatives](lesson.html?lesson=pretraining-03) · [Source](api/v1/lessons/pretraining-03.json) · phases/17-pretraining/03-contrastive
- [Supervised fine-tuning with response-only token loss](lesson.html?lesson=posttraining-01) · [Source](api/v1/lessons/posttraining-01.json) · phases/18-posttraining/01-sft
- [LoRA: adapt, save and merge low-rank updates](lesson.html?lesson=posttraining-02) · [Source](api/v1/lessons/posttraining-02.json) · phases/18-posttraining/02-lora
- [Learn a reward model from pairwise preferences](lesson.html?lesson=posttraining-03) · [Source](api/v1/lessons/posttraining-03.json) · phases/18-posttraining/03-reward
- [RLHF mechanics: a frozen reward and PPO policy updates](lesson.html?lesson=posttraining-04) · [Source](api/v1/lessons/posttraining-04.json) · phases/18-posttraining/04-rlhf
- [Direct preference optimization and reference-corrected margins](lesson.html?lesson=posttraining-05) · [Source](api/v1/lessons/posttraining-05.json) · phases/18-posttraining/05-dpo
- [Select a TPU slice and provision your Cloud TPU VM](lesson.html?lesson=tpu-01) · [Source](api/v1/lessons/tpu-01.json) · phases/19-tpu/01-get-a-tpu-running-and-verify-the-runtime
- [Bootstrap .venv and verify a TPU runtime receipt](lesson.html?lesson=tpu-04) · [Source](api/v1/lessons/tpu-04.json) · phases/19-tpu/02-bootstrap-venv-and-verify-tpu-runtime-receipt
- [Stage and launch checkpointed TPU jobs](lesson.html?lesson=tpu-02) · [Source](api/v1/lessons/tpu-02.json) · phases/20-tpu-workflows/01-stage-and-launch-checkpointed-tpu-jobs
- [Run course scripts and multi-host TPU Pod slices](lesson.html?lesson=tpu-05) · [Source](api/v1/lessons/tpu-05.json) · phases/20-tpu-workflows/02-run-course-scripts-and-multihost-tpu-slices
- [TPU generations and HBM memory budgeting](lesson.html?lesson=tpu-03) · [Source](api/v1/lessons/tpu-03.json) · phases/21-tpu-systems/01-tpu-generations-and-hbm-memory-budgeting
- [bfloat16, int8 precision, and XProf profiling on TPU](lesson.html?lesson=tpu-06) · [Source](api/v1/lessons/tpu-06.json) · phases/21-tpu-systems/02-bfloat16-int8-precision-and-xprof-profiling

## Open Engineering TODOs

- Add runnable starter/solution harnesses for the 10 planned large-model and retrieval capstones in projects/.
- Add automated Cloud TPU VM CI execution receipts alongside the local CPU verification receipts.