{ "schemaVersion": 1, "pathways": [ { "id": "foundations", "title": "Start with JAX", "description": "Learn how arrays, derivatives and explicit state fit together by building and checking a regression experiment. Use the deeper math lessons when you need to explain why an update works or fails.", "prerequisites": "Python basics and a little NumPy", "outcome": "A fitted regression model", "libraries": "jax · jax.numpy · optax", "source": "https://docs.jax.dev/en/latest/beginner_guide.html", "phaseIds": [ "welcome", "arrays", "transforms", "state", "optimization" ], "capstone": { "title": "A fitted regression model", "status": "authored", "assessment": "Show a tested loss, numerically checked gradients, and a reproducible fit.", "evidence": [ "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" ], "projectId": "regression-audit", "assessmentDraft": { "id": "foundations", "status": "active", "source": "assessments/foundations.md", "url": "assessments/foundations.html" } }, "audience": "Python learners who can write functions and use lists; NumPy experience helps but is not required.", "firstArtifact": "A regression implementation with an independent gradient check and held-out predictions.", "focusPhaseIds": [ "arrays", "transforms", "state", "optimization" ], "modalityTrackIds": [ "image", "text", "audio", "cross-modal" ], "studyAdvice": "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.", "assessmentExtensions": [ { "id": "math", "title": "Math & optimization synthesis", "url": "assessments/math.html", "status": "active" } ], "engineeringLessonIds": [ "recovery-06" ] }, { "id": "models", "title": "Train neural networks", "description": "Train a CPU image classifier, preserve its complete state through interruption, inspect held-out errors and verify its exported inference. Language modeling and TPU work are optional branches.", "prerequisites": "Shared foundations", "outcome": "A recoverable image classifier with verified exported inference", "libraries": "Flax · Optax · Orbax", "source": "https://flax.readthedocs.io/", "phaseIds": [ "welcome", "arrays", "transforms", "state", "optimization", "networks", "recovery", "performance", "deployment" ], "capstone": { "title": "A recoverable image classifier with verified exported inference", "status": "authored", "assessment": "Complete the eight image lifecycle stages and changed-condition synthesis: independent CNN checks, full-state recovery, frozen error analysis, precision, export and measured inference. Transformer-specific evidence belongs to the complementary text harness.", "evidence": [ "Training and evaluation splits, seeded configuration and learning curves", "A saved full state whose next update matches an uninterrupted run", "Predictions from the restored model and examples explaining its failures" ], "assessmentDraft": { "id": "models", "projectId": "image-harness", "title": "Classifier and image lifecycle synthesis", "status": "active", "scope": "project-synthesis", "source": "assessments/models.md", "url": "assessments/models.html" }, "projectId": "image-harness", "assessmentId": "models" }, "audience": "Learners who can trace shapes, differentiate a scalar loss and carry state through an update.", "firstArtifact": "A CPU classifier with stable loss, replay checks and count-weighted evaluation.", "focusPhaseIds": [ "networks", "recovery", "deployment" ], "modalityTrackIds": [ "image", "text" ], "studyAdvice": "Build the classifier first, freeze its evaluation split and add optimizer, random and data-position state before rehearsing interruption. Continue through CPU performance and deployment to the image harness. Use the reproducible-experiments workshop to compare changes. The optional language branch starts at Transformers after recovery; add pretraining and supervised adaptation next. Complete the RL phase before learned-reward PPO. These branches are separate from the classifier evidence.", "engineeringLessonIds": [ "recovery-06", "deployment-07", "operations-06" ], "additionalGuides": [ { "title": "Reproducible JAX/Flax experiments workshop", "url": "project-guides/reproducible-experiments.html", "scope": "CPU-tested experiments, logging, paired comparisons and a Markdown research notebook. Optional trackers run locally; accelerator qualification is separate." }, { "title": "Optional TPU and Google Cloud specialization", "url": "pathways.html?path=tpu", "scope": "The required path runs on CPU. This optional route adds cloud provisioning, TPU execution and hardware-specific evidence; CPU examples do not qualify TPU performance." }, { "title": "Optional language-model branch", "url": "course.html#phase-transformers", "scope": "After recovery: Transformers → self-supervised pretraining → supervised adaptation. The learned-reward PPO lesson additionally requires the four RL lessons. Keep text-harness evidence separate from the image capstone." }, { "title": "Optional staged embedding-model pilot", "url": "project-guides/embedding-lineage.html", "scope": "A 70,765-parameter CPU encoder carries checkpoints from MLM through contrastive continuation and query-positive adaptation to held-out retrieval and fresh-process export. Its measured poor result is retained. Follow its language and recovery prerequisites; this is separate from the image capstone and from planned 300M training." } ] }, { "id": "science", "title": "Simulate and discover", "description": "Simulate physical dynamics, check numerical and derivative accuracy, and recover unknown parameters with held-out evidence. Build a reusable CPU inverse-problem audit before extending to adaptive solvers or accelerators.", "prerequisites": "Shared foundations and calculus", "outcome": "A scientific inverse problem", "libraries": "JAX · NumPy; optional Diffrax ecosystem extension", "source": "https://docs.kidger.site/diffrax/", "phaseIds": [ "welcome", "arrays", "transforms", "state", "optimization", "science" ], "capstone": { "title": "Audit a differentiable scientific inverse problem", "status": "authored", "assessment": "Complete the staged project contracts and the changed-condition synthesis assessment with independent calculations, observed outputs, and failure diagnosis.", "evidence": [ "Convergence against the analytic solution at multiple rates and time grids", "Solver sensitivities checked by independent finite differences", "Parameter recovery, held-out errors and a diagnosis of numerical bias or non-identifiability" ], "projectId": "scientific-inverse", "assessmentId": "science", "assessmentDraft": { "id": "science", "projectId": "scientific-inverse", "title": "Scientific inverse-problem synthesis", "status": "active", "scope": "project-synthesis", "source": "assessments/science.md", "url": "assessments/science.html" } }, "audience": "Scientists and engineers with a domain model to test; calculus is introduced in the shared math lessons.", "firstArtifact": "A cooling simulator with a convergence table and an independent sensitivity check.", "focusPhaseIds": [ "science" ], "modalityTrackIds": [], "studyAdvice": "Work from the cooling model through batched trajectories, sensitivities and parameter fitting. At each step compare against the analytic solution before changing the physical rate or time grid. Use the deeper linear algebra and curvature lessons when parameters are hard to identify. Finish the inverse-problem project, then attempt the two-compartment assessment with its independent conservation check.", "engineeringLessonIds": [ "recovery-06" ] }, { "id": "probability", "title": "Model uncertainty", "description": "Express assumptions as probability models, verify a known posterior, then implement HMC and variational inference with explicit diagnostics. Use predictive checks to distinguish a sampler problem from a model that does not explain the data.", "prerequisites": "Shared foundations and probability", "outcome": "A checked Bayesian regression", "libraries": "JAX · NumPy; optional NumPyro ecosystem extension", "source": "https://num.pyro.ai/en/stable/", "phaseIds": [ "welcome", "arrays", "transforms", "state", "optimization", "probability" ], "capstone": { "title": "Audit a Bayesian regression model", "status": "authored", "assessment": "Defend the prior and likelihood, verify posterior and sampler calculations independently, distinguish latent and observation uncertainty, and diagnose changed-condition predictive failures.", "evidence": [ "Prior, likelihood and posterior calculations checked against an independent reference", "Multiple seeded chains with trace, acceptance and between-chain diagnostics", "Predictive intervals separating latent and observation variance, including a model-mismatch example" ], "projectId": "bayesian-regression", "assessmentId": "probability", "assessmentDraft": { "id": "probability", "projectId": "bayesian-regression", "title": "Bayesian regression synthesis", "status": "active", "scope": "project-synthesis", "source": "assessments/probability.md", "url": "assessments/probability.html" } }, "audience": "Learners who want to move beyond a point prediction and explain assumptions and uncertainty.", "firstArtifact": "An exact Bayesian regression posterior checked against an independent host calculation.", "focusPhaseIds": [ "probability" ], "modalityTrackIds": [], "studyAdvice": "Start with distributions and the exact regression posterior so you have a reference for approximate inference. Run multiple HMC chains and deliberately inspect an unstable sampler before studying variational inference. Compare latent uncertainty with uncertainty in a future observation. Complete the Bayesian project with a changed prior and held-out data before moving to an optional NumPyro model.", "engineeringLessonIds": [ "recovery-06" ] }, { "id": "rl", "title": "Learn by interacting", "description": "Train a policy from actual environment interaction, verify its objective against an independent small-state reference, and compare multiple seeded agents on fixed held-out evaluation streams.", "prerequisites": "Shared foundations and basic ML", "outcome": "An evaluated policy", "libraries": "JAX · Optax", "source": "https://github.com/FLAIROx/Jumanji", "phaseIds": [ "welcome", "arrays", "transforms", "state", "optimization", "networks", "rl" ], "capstone": { "title": "Train and audit a policy from interaction", "status": "authored", "assessment": "Complete the staged project contracts and the changed-condition synthesis assessment with independent calculations, observed outputs, and failure diagnosis.", "evidence": [ "Verified transitions, terminal masks and reconstructed return targets", "A policy-gradient comparison against independent state enumeration", "Learning curves and fixed held-out evaluation for several trained seeds, with failed trajectories explained" ], "projectId": "policy-evaluation", "assessmentId": "rl", "assessmentDraft": { "id": "rl", "projectId": "policy-evaluation", "title": "Reinforcement-learning synthesis", "status": "active", "scope": "project-synthesis", "source": "assessments/rl.md", "url": "assessments/rl.html" } }, "audience": "Learners interested in sequential decisions who can already implement and evaluate a neural network.", "firstArtifact": "A tested environment transition table and reproducible batched trajectories.", "focusPhaseIds": [ "rl" ], "modalityTrackIds": [], "studyAdvice": "Test one environment transition before batching it. Reconstruct a trajectory by hand, then compare the policy gradient with the small-state enumeration reference. Train multiple agents with different seeds and freeze the evaluation policy and random streams. Use the policy project to separate an implementation bug, Monte Carlo variation and an algorithm that learns poorly.", "engineeringLessonIds": [ "recovery-06", "operations-06" ], "additionalGuides": [ { "title": "Optional TPU and Google Cloud specialization", "url": "pathways.html?path=tpu", "scope": "The required path runs on CPU. This optional route adds cloud provisioning, TPU execution and hardware-specific evidence; CPU examples do not qualify TPU performance." } ] }, { "id": "scale", "title": "Make it fast. Make it scale.", "description": "Profile a correct workload, preserve its global objective while partitioning arrays, and audit actual Pallas kernels before requesting target measurements. CPU semantics are available; accelerator speed and multi-host claims require their own receipts.", "prerequisites": "Shared foundations and a working training loop", "outcome": "A checked kernel and a measurement-backed optimization argument", "libraries": "JAX sharding · profiling · Pallas", "source": "https://docs.jax.dev/en/latest/", "phaseIds": [ "welcome", "arrays", "transforms", "state", "optimization", "networks", "recovery", "performance", "distributed", "kernels" ], "capstone": { "title": "A checked kernel and a measurement-backed optimization argument", "status": "authored", "assessment": "Defend the profile, global objective, partitioning and kernel correctness; quantify a whole-workload opportunity. CPU checks establish semantics. Actual accelerator execution and performance remain a separate required qualification for device claims.", "evidence": [ "Reference agreement across more than one shape and parameter setting", "Synchronized compile and execution measurements with a trace of the bottleneck", "Array partitioning, collective communication and a before/after comparison on the stated hardware" ], "projectId": "kernel-audit", "assessmentId": "scale", "assessmentDraft": { "id": "scale", "projectId": "kernel-audit", "title": "Profiling, partitioning and kernel synthesis", "status": "active", "scope": "project-synthesis", "source": "assessments/scale.md", "url": "assessments/scale.html" } }, "audience": "Engineers comfortable with training loops who want to reason about memory, communication and throughput.", "firstArtifact": "A synchronized timing report and a checked array placement on logical CPU devices.", "focusPhaseIds": [ "performance", "distributed", "kernels" ], "modalityTrackIds": [ "image", "text" ], "studyAdvice": "Keep one correct reference workload throughout this path. Measure completed work, inspect its trace and relate the bottleneck to memory traffic or computation. Verify partitioning and collectives on logical devices, then audit a Pallas kernel against independent references. Run the separate guarded TPU measurement only on supported target hardware before claiming a device or system speedup.", "engineeringLessonIds": [ "deployment-07", "operations-06" ], "additionalGuides": [ { "title": "Optional TPU and Google Cloud specialization", "url": "pathways.html?path=tpu", "scope": "The required path runs on CPU. This optional route adds cloud provisioning, TPU execution and hardware-specific evidence; CPU examples do not qualify TPU performance." } ] }, { "id": "internals", "title": "Understand JAX itself", "description": "Trace a function into jaxpr, verify forward and reverse sensitivities, write exact custom derivative rules, and compile a small transformation whose limits you can explain.", "prerequisites": "Shared foundations and advanced autodiff", "outcome": "Audit derivatives and a tiny transformation", "libraries": "jaxpr · autodiff · XLA", "source": "https://docs.jax.dev/en/latest/", "phaseIds": [ "welcome", "arrays", "transforms", "state", "optimization", "internals" ], "capstone": { "title": "Audit derivatives and a tiny transformation", "status": "authored", "assessment": "Connect a jaxpr to your function and numerically verify the derivative rule.", "evidence": [ "An annotated jaxpr linking primitives and shapes to the original function", "Directional and adjoint derivative checks on changed inputs", "A custom derivative or transformation with a deliberate unsupported case and explicit rejection" ], "projectId": "derivative-audit", "assessmentId": "internals", "assessmentDraft": { "id": "internals", "projectId": "derivative-audit", "title": "Derivative synthesis", "status": "active", "scope": "project-synthesis", "source": "assessments/internals.md", "url": "assessments/internals.html" } }, "audience": "Curious JAX users who can already use grad and jit and want to diagnose transformation behavior.", "firstArtifact": "An annotated jaxpr that explains a computation and a changed input shape.", "focusPhaseIds": [ "internals" ], "modalityTrackIds": [ "text", "image" ], "studyAdvice": "Start by connecting each jaxpr primitive to the function that produced it. Revisit directional derivatives and the chain rule before studying JVPs, VJPs and custom rules. Keep an independent derivative reference beside every transformation. Finish the derivative audit, then extend the tiny interpreter in the synthesis assessment and explain its unsupported cases.", "engineeringLessonIds": [ "deployment-07" ] }, { "id": "tpu", "title": "Learn TPUs and Google Cloud", "description": "Build and recover an experiment on CPU, then learn GCP provisioning, TPU generations, training and inference memory, precision policies and XProf profiling. Connect measured error and real target traces to model engineering.", "prerequisites": "Shared foundations and basic machine learning", "outcome": "A reproducible run and a hardware, memory, precision and profiling decision report", "libraries": "JAX · Flax NNX · Optax · Orbax · gcloud · XProf", "source": "https://docs.jax.dev/en/latest/", "phaseIds": [ "welcome", "tpu", "tpu-workflows", "tpu-systems", "arrays", "transforms", "state", "optimization", "networks", "recovery", "transformers", "pretraining", "rl", "posttraining", "performance", "distributed", "deployment", "operations" ], "capstone": { "title": "Causal Transformer training-system capstone", "status": "authored", "assessment": "Train and recover a byte Transformer, verify cache precision and exported inference, and benchmark on CPU and Cloud TPU VM.", "evidence": [ "Independent causal-attention and valid-token loss checks on disjoint data", "Fresh-process recovery of parameters, both Adam moments, dropout keys and document position", "FP32/BF16/INT8 cache parity, stored bytes and numerical-error comparisons", "Six serialized prefill/decode endpoints with fresh-process reload and artifact rejection", "A genuine profile and synchronized CPU measurements, plus separate actual-target qualification before any TPU claim", "Cloud resource identity, resolved environment, actual backend events, retrieved checkpoints and verified resource cleanup for any claimed TPU run", "Generation and topology choice, explicit weight/activation/accumulator/cache policy, memory budget and one tested profile hypothesis with retained quality evidence" ], "projectId": "text-harness", "assessmentId": "tpu", "assessmentDraft": { "id": "tpu", "projectId": "text-harness", "title": "Training-system synthesis", "status": "active", "scope": "project-synthesis", "source": "assessments/tpu.md", "url": "assessments/tpu.html", "pathwayId": "tpu", "phaseId": "transformers" } }, "audience": "Model builders who want to provision TPUs, launch single-host and multi-host experiments, and connect sequence models to accelerator engineering.", "firstArtifact": "A verified TPU runtime receipt, a checkpoint-resumed TPU experiment log, and a memory/precision/XProf decision dossier.", "focusPhaseIds": [ "tpu", "tpu-workflows", "tpu-systems", "transformers", "performance", "distributed", "pretraining", "posttraining", "operations" ], "modalityTrackIds": [ "text", "cross-modal" ], "studyAdvice": "Complete the 3-part Cloud TPU mini-course (19: TPU Setup, 20: TPU Workflows, 21: TPU Systems) immediately after 00: Setup & first steps so you know how to get a TPU running, stage scripts with gcloud compute tpus tpu-vm scp, launch single-host and --worker=all jobs, and capture XProf traces before starting any TPU-dependent course. Then build the tiny Transformer, verify causal masking and held-out token loss, restore the complete state in a fresh process, and run the distributed, pretraining, post-training, and operations experiments on your TPU VM in parallel with the course.", "engineeringLessonIds": [ "tpu-01", "tpu-04", "tpu-02", "tpu-05", "tpu-03", "tpu-06", "recovery-06", "deployment-07", "operations-07", "performance-03", "performance-04" ], "practicalGuide": { "title": "From a laptop experiment to a TPU training run", "url": "tpu-gcp.html", "scope": "CPU-tested launcher; cloud instructions require target qualification." }, "additionalGuides": [ { "title": "TPU generations, precision and profiling", "url": "tpu-performance.html", "scope": "CPU-tested numerical lab and recorded figures; TPU execution and performance require target qualification." } ] }, { "id": "ship", "title": "Ship and adapt JAX workloads", "description": "Move trained computations into a measured inference service. Check framework/precision parity, package a non-root Docker image, verify readiness and connect evaluation, model identity and rollback evidence.", "prerequisites": "Shared foundations and a working model", "outcome": "Verify an inference artifact and capacity hypothesis", "libraries": "JAX export · Keras · TensorFlow · PyTorch · LiteRT · Docker · MLflow", "source": "https://docs.jax.dev/en/latest/", "phaseIds": [ "welcome", "arrays", "transforms", "state", "optimization", "networks", "recovery", "performance", "deployment" ], "capstone": { "title": "Verify an inference artifact and capacity hypothesis", "status": "authored", "assessment": "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.", "evidence": [ "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" ], "projectId": "deployment-audit", "assessmentId": "ship", "assessmentDraft": { "id": "ship", "projectId": "deployment-audit", "title": "Shipping synthesis", "status": "active", "scope": "project-synthesis", "source": "assessments/ship.md", "url": "assessments/ship.html" } }, "audience": "Model builders who want to package reliable inference on servers or edge devices.", "firstArtifact": "An exported computation checked against its source on the same inputs.", "focusPhaseIds": [ "deployment" ], "modalityTrackIds": [ "image", "text", "audio", "cross-modal" ], "studyAdvice": "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.", "engineeringLessonIds": [ "recovery-06", "deployment-07", "operations-06", "operations-07", "operations-08" ], "additionalGuides": [ { "title": "Optional TPU and Google Cloud specialization", "url": "pathways.html?path=tpu", "scope": "The required path runs on CPU. This optional route adds cloud provisioning, TPU execution and hardware-specific evidence; CPU examples do not qualify TPU performance." } ] }, { "id": "ops", "title": "Operate JAX workloads", "description": "Operate reproducible model workloads from tracked experiments to containerized inference. Practice failure recovery, MLOps release gates, LLMOps evaluation and tracing, and ModelOps ownership, approval and rollback.", "prerequisites": "Shared foundations and a runnable training system", "outcome": "An observable, restartable local workload", "libraries": "JAX distributed · Grain · Orbax · XProf · MLflow · Docker", "source": "https://docs.jax.dev/en/latest/", "phaseIds": [ "welcome", "arrays", "transforms", "state", "optimization", "networks", "recovery", "performance", "distributed", "deployment", "operations" ], "capstone": { "title": "An observable, restartable local workload", "status": "authored", "assessment": "Diagnose real local subprocess failures and timeouts, restore the accepted complete state, reject a degraded artifact and defend rollback and capacity decisions. Accelerator and cluster operation require additional target evidence.", "evidence": [ "Versioned configuration, structured progress signals and an incident timeline", "A controlled failure followed by a verified restore of the next update", "An accepted artifact, rollback criteria and capacity calculations separating measurements from assumptions" ], "projectId": "workload-operations", "assessmentId": "ops", "assessmentDraft": { "id": "ops", "projectId": "workload-operations", "title": "Workload operations synthesis", "status": "active", "scope": "project-synthesis", "source": "assessments/ops.md", "url": "assessments/ops.html" } }, "audience": "Engineers interested in reproducibility, failure recovery and operating accelerator workloads.", "firstArtifact": "An interrupted CPU run whose next data batches, losses and full state match uninterrupted execution.", "focusPhaseIds": [ "recovery", "performance", "distributed", "operations", "deployment" ], "modalityTrackIds": [ "image", "text", "audio", "cross-modal" ], "studyAdvice": "First reproduce the next update after a local interruption. Then inspect workload traces and placement so an incident has interpretable signals. Work through the operations project’s process, event, restore, release and capacity stages. Keep a runbook naming the accepted checkpoint and rollback decision; validate an actual cluster separately before claiming accelerator operations experience. After the deployment bridge, complete engineering-release to connect MLflow, a verified container and application/release policy. Keep local Docker evidence separate from cloud and accelerator qualification.", "engineeringLessonIds": [ "recovery-06", "deployment-07", "operations-06", "operations-07", "operations-08" ], "additionalGuides": [ { "title": "Reproducible JAX/Flax experiments workshop", "url": "project-guides/reproducible-experiments.html", "scope": "CPU-tested experiments, logging, paired comparisons and a Markdown research notebook. Optional trackers run locally; accelerator qualification is separate." }, { "title": "Optional TPU and Google Cloud specialization", "url": "pathways.html?path=tpu", "scope": "The required path runs on CPU. This optional route adds cloud provisioning, TPU execution and hardware-specific evidence; CPU examples do not qualify TPU performance." } ] } ] }