{ "schemaVersion": 1, "title": "JAX from first principles to working systems", "phases": [ { "id": "welcome", "number": "00", "title": "Setup & first steps", "part": "Start here", "prerequisitePhaseIds": [], "hardware": "CPU; the TPU bridge is optional accelerator practice", "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.", "source": "https://docs.jax.dev/en/latest/installation.html", "description": "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.", "learningAdvice": "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.", "projectId": "welcome-audit", "projectStageIds": [ "1", "2", "3", "4" ], "studyGuide": { "question": "Can someone else reproduce your first result?", "why": "Start with one saved program and a known answer. The terminal, interpreter and working folder are part of the experiment: a notebook that imports a package is not yet proof that the downloaded script uses the same environment.", "readiness": { "prompt": "You saved a script in your course folder, but the terminal reports that the file does not exist. Which observation should you check first?", "answer": "Check the terminal’s current folder and the exact filename, including its extension. Then run the environment’s Python explicitly. Reinstalling JAX will not repair a wrong path.", "lessonId": "welcome-01" }, "milestones": [ { "lessonIds": [ "welcome-01", "welcome-02" ], "title": "Make one reproducible result", "task": "Run the row-sum example as a saved file. Record the Python executable, package versions, backend, input shape and computed sums. Explain why a package version alone cannot identify which interpreter ran the file." }, { "lessonIds": [ "welcome-cpu", "welcome-03" ], "title": "Separate local practice from parallel TPU execution", "task": "Start a fresh process for virtual CPU devices and compare its device receipt with the ordinary process. Then learn how to get a TPU running (hosted TPU or single-host Cloud TPU VM), launch your experiment over SSH with explicit backend checks, and run subsequent course phases on TPU in parallel." } ], "transfer": { "prompt": "A colleague gets the right sums but a different device count. Decide whether the arithmetic is wrong and how to compare the runs.", "approach": "Compare values and shapes first, then backend, process startup flags and environment. Device enumeration is a runtime property; equal arithmetic does not imply equal hardware or performance." }, "pitfall": { "symptom": "Import works in a notebook but fails in the terminal.", "check": "Compare interpreter paths before changing packages; launch both from the intended environment." }, "completion": "Keep a saved script, its command and actual output, plus an explanation of one repaired setup error. Use foundation-toolkit stage 1. A screenshot without a rerunnable file is incomplete.", "furtherWork": "Fresh-machine Windows/Linux/macOS walkthroughs and novice usability sessions remain unverified; accelerator setup needs a separately observed device run." }, "projectNote": "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.", "additionalProjectIds": [ "foundation-toolkit" ], "lessons": [ { "id": "welcome-01", "title": "Set up your learning workspace", "status": "authored", "objective": "Download a starter workspace, select its Python environment, run a saved program, and verify one deliberate change.", "prerequisites": [], "exercise": "Edit the starter as shown in Step 6: change the array length from 4 to 6 and the expected sum from 6 to 15. Save, rerun, and keep my-first-run.txt plus a note explaining why the result changed.", "evidence": "Keep first_experiment.py, my-first-run.txt with package/device output, the environment interpreter path, and a note deriving the change from a sum of 6 to 15. Record a setup error and its repair if you encountered one.", "check": "Which evidence lets another learner reproduce your environment?", "artifacts": { "source": "phases/00-welcome/01-set-up-your-learning-workspace/docs/en.md", "script": "exercises/welcome-01.py", "notebook": "exercises/welcome-01.ipynb", "markdown": "exercises/welcome-01.md", "scriptSource": "phases/00-welcome/01-set-up-your-learning-workspace/code/main.py", "notebookSource": "phases/00-welcome/01-set-up-your-learning-workspace/notebooks/exercise.ipynb", "quizSource": "phases/00-welcome/01-set-up-your-learning-workspace/quiz.json", "evidenceSource": "phases/00-welcome/01-set-up-your-learning-workspace/outputs/evidence.md", "evidence": "exercises/welcome-01-evidence.md" }, "path": "phases/00-welcome/01-set-up-your-learning-workspace", "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 40, "contentHash": "f15bc8e6b9ae9d88474253a7e4224a1c0d7f398f5a98b3f18f76fcbbdd5ef7e0", "visualArtifact": { "image": "figures/welcome-01.svg", "png": "figures/welcome-01.png", "data": "figures/welcome-01.json", "generatedAt": "2026-10-20T10:45:04.797711+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "f15bc8e6b9ae9d88474253a7e4224a1c0d7f398f5a98b3f18f76fcbbdd5ef7e0", "kind": "executed", "mechanismImage": "figures/welcome-01-mechanism.svg", "mechanismPng": "figures/welcome-01-mechanism.png" }, "execution": { "contentHash": "f15bc8e6b9ae9d88474253a7e4224a1c0d7f398f5a98b3f18f76fcbbdd5ef7e0", "executedAt": "2026-10-20T10:45:05.346444+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Python: 3.14.3\nJAX: 0.9.2\nNumPy: 2.4.4\nBackend: cpu\nDevices: [CpuDevice(id=0)]\nSum: 6.0\nChanged experiment sum: 15.0\nPASS: welcome-01\n", "url": "executions/welcome-01.json" }, "url": "lesson.html?lesson=welcome-01" }, { "id": "welcome-02", "title": "Meet your arrays and devices", "status": "authored", "objective": "Create an array, inspect its dtype and device, and explain the result.", "prerequisites": [ "welcome-01" ], "exercise": "Compute column sums instead. State the resulting shape before running the code.", "evidence": "Keep the labeled table, row/column calculations, centered-data check, array dtype/device report and the failed reshape with its element-count diagnosis.", "check": "For an array of shape (2, 3), what shape results from sum(axis=0)?", "artifacts": { "source": "phases/00-welcome/02-meet-your-arrays-and-devices/docs/en.md", "script": "exercises/welcome-02.py", "notebook": "exercises/welcome-02.ipynb", "markdown": "exercises/welcome-02.md", "scriptSource": "phases/00-welcome/02-meet-your-arrays-and-devices/code/main.py", "notebookSource": "phases/00-welcome/02-meet-your-arrays-and-devices/notebooks/exercise.ipynb", "quizSource": "phases/00-welcome/02-meet-your-arrays-and-devices/quiz.json", "evidenceSource": "phases/00-welcome/02-meet-your-arrays-and-devices/outputs/evidence.md", "evidence": "exercises/welcome-02-evidence.md" }, "path": "phases/00-welcome/02-meet-your-arrays-and-devices", "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 45, "contentHash": "9cebbe5b29a2efa8d335a07bc2113a61deefcc78edea629b5398a996222f7023", "visualArtifact": { "image": "figures/welcome-02.svg", "png": "figures/welcome-02.png", "data": "figures/welcome-02.json", "generatedAt": "2026-10-20T10:45:05.555675+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "9cebbe5b29a2efa8d335a07bc2113a61deefcc78edea629b5398a996222f7023", "kind": "executed" }, "execution": { "contentHash": "9cebbe5b29a2efa8d335a07bc2113a61deefcc78edea629b5398a996222f7023", "executedAt": "2026-10-20T10:45:06.782121+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Shape: (2, 3)\nDtype: float32\nDevices: {CpuDevice(id=0)}\nRow sums: [ 3. 12.]\nShape: (2, 3)\nDtype: float32\nDevices: {CpuDevice(id=0)}\nRow sums: [ 3. 12.]\nMeans: [[1.5 2.5 3.5]]\nCentered: [[-1.5 -1.5 -1.5]\n [ 1.5 1.5 1.5]]\nHost dtype: float32\nExpected element-count mismatch\nPASS: welcome-02\n", "url": "executions/welcome-02.json" }, "url": "lesson.html?lesson=welcome-02" }, { "id": "welcome-cpu", "title": "Practice with four virtual CPU devices", "status": "authored", "objective": "Start four logical CPU devices, inspect row-sharded and replicated arrays, and diagnose an indivisible batch.", "prerequisites": [ "welcome-02" ], "exercise": "Predict and verify the per-device shapes for twelve rows, then demonstrate how replication repairs an indivisible five-row placement without deleting data.", "evidence": "Keep package/device reports, global and local shapes, shard index/value tables, NumPy comparisons, changed-shape predictions, and the five-row failure/repair. State that the run uses one host and does not validate TPU performance.", "check": "What does a (12,2) global array sharded with P(\"data\",None) over four devices mean?", "path": "phases/00-welcome/03-virtual-cpu-devices", "runtime": { "kind": "virtual-cpu", "deviceCount": 4 }, "artifacts": { "source": "phases/00-welcome/03-virtual-cpu-devices/docs/en.md", "script": "exercises/welcome-cpu.py", "notebook": "exercises/welcome-cpu.ipynb", "markdown": "exercises/welcome-cpu.md", "scriptSource": "phases/00-welcome/03-virtual-cpu-devices/code/main.py", "notebookSource": "phases/00-welcome/03-virtual-cpu-devices/notebooks/exercise.ipynb", "quizSource": "phases/00-welcome/03-virtual-cpu-devices/quiz.json", "evidenceSource": "phases/00-welcome/03-virtual-cpu-devices/outputs/evidence.md", "evidence": "exercises/welcome-cpu-evidence.md" }, "hardware": "4 logical CPU devices on one host", "minutes": 60, "contentHash": "733a1d70e7c11d1340a2a478102aed48eef83dc5f8dc9679741d851a71426806", "visualArtifact": { "image": "figures/welcome-cpu.svg", "png": "figures/welcome-cpu.png", "data": "figures/welcome-cpu.json", "generatedAt": "2026-10-20T10:45:06.292248+00:00", "environment": { "backend": "cpu", "deviceCount": 4, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "733a1d70e7c11d1340a2a478102aed48eef83dc5f8dc9679741d851a71426806", "kind": "executed" }, "execution": { "contentHash": "733a1d70e7c11d1340a2a478102aed48eef83dc5f8dc9679741d851a71426806", "executedAt": "2026-10-20T10:45:07.979734+00:00", "environment": { "backend": "cpu", "deviceCount": 4, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "JAX version: 0.9.2\nCPU devices: 4\nDevice 0 index (slice(0, 1, None), slice(None, None, None)) values [[0. 1.]]\nDevice 1 index (slice(1, 2, None), slice(None, None, None)) values [[2. 3.]]\nDevice 2 index (slice(2, 3, None), slice(None, None, None)) values [[4. 5.]]\nDevice 3 index (slice(3, 4, None), slice(None, None, None)) values [[6. 7.]]\nGlobal shape: (4, 2) result: [[1.0, 3.0], [5.0, 7.0], [9.0, 11.0], [13.0, 15.0]]\nJAX version: 0.9.2\nCPU devices: 4\nDevice 0 index (slice(0, 1, None), slice(None, None, None)) values [[0. 1.]]\nDevice 1 index (slice(1, 2, None), slice(None, None, None)) values [[2. 3.]]\nDevice 2 index (slice(2, 3, None), slice(None, None, None)) values [[4. 5.]]\nDevice 3 index (slice(3, 4, None), slice(None, None, None)) values [[6. 7.]]\nGlobal shape: (4, 2) result: [[1.0, 3.0], [5.0, 7.0], [9.0, 11.0], [13.0, 15.0]]\nNegative input result: [[-3.0, -5.0], [-7.0, -9.0], [-11.0, -13.0], [-15.0, -17.0]]\nReplicated local shapes: [(4, 2), (4, 2), (4, 2), (4, 2)]\nExpected indivisible row dimension\nExpected indivisible row dimension\nPASS: welcome-cpu\n", "url": "executions/welcome-cpu.json" }, "url": "lesson.html?lesson=welcome-cpu" }, { "id": "welcome-03", "title": "Move your experiment to a TPU", "status": "authored", "objective": "Provision a hosted TPU runtime or single-host Cloud TPU VM, launch an explicit TPU target path over SSH that fails when TPU is unavailable, verify the device of a completed prediction against NumPy, and set up parallel TPU execution for the rest of the course.", "prerequisites": [ "welcome-cpu" ], "exercise": "Predict the effect of replacing the bias by minus one quarter. Run a second compiled prediction on the selected device. Compare against the original prediction plus the expected constant shift and an independent NumPy calculation.", "evidence": "Keep separate CPU and TPU reports with requested/observed platform, package versions, completed predictions, NumPy error and changed-bias and row-permutation checks. Until the target command succeeds, TPU execution remains unverified.", "check": "The values match NumPy, but the output array is on CPU when the requested platform was TPU. What has been demonstrated?", "path": "phases/00-welcome/03-move-your-experiment-to-a-tpu", "runtime": { "kind": "cpu" }, "hardware": "CPU reference; hosted TPU notebook or single-host Google Cloud TPU VM", "artifacts": { "source": "phases/00-welcome/03-move-your-experiment-to-a-tpu/docs/en.md", "script": "exercises/welcome-03.py", "notebook": "exercises/welcome-03.ipynb", "markdown": "exercises/welcome-03.md", "scriptSource": "phases/00-welcome/03-move-your-experiment-to-a-tpu/code/main.py", "notebookSource": "phases/00-welcome/03-move-your-experiment-to-a-tpu/notebooks/exercise.ipynb", "quizSource": "phases/00-welcome/03-move-your-experiment-to-a-tpu/quiz.json", "evidenceSource": "phases/00-welcome/03-move-your-experiment-to-a-tpu/outputs/evidence.md", "evidence": "exercises/welcome-03-evidence.md" }, "minutes": 45, "contentHash": "c635ba5492b7276fae2ed1af1cbef264ceb432ae40c49d9b40291f8eb84821ac", "visualArtifact": { "image": "figures/welcome-03.svg", "png": "figures/welcome-03.png", "data": "figures/welcome-03.json", "generatedAt": "2026-10-20T10:45:06.987466+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "c635ba5492b7276fae2ed1af1cbef264ceb432ae40c49d9b40291f8eb84821ac", "kind": "executed" }, "execution": { "contentHash": "c635ba5492b7276fae2ed1af1cbef264ceb432ae40c49d9b40291f8eb84821ac", "executedAt": "2026-10-20T10:45:09.080862+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "{\n \"expected\": \"cpu\",\n \"output_devices\": [\n \"TFRT_CPU_0\"\n ],\n \"platform\": \"cpu\",\n \"jax\": \"0.9.2\",\n \"python\": \"3.14.3\",\n \"shape\": [\n 8\n ],\n \"dtype\": \"float32\",\n \"max_absolute_error\": 0.0,\n \"process_index\": 0,\n \"process_count\": 1\n}\nCompleted prediction: [0.34375, 0.8125, 1.28125, 1.75, 2.21875, 2.6875, 3.15625, 3.625]\n{\n \"expected\": \"cpu\",\n \"output_devices\": [\n \"TFRT_CPU_0\"\n ],\n \"platform\": \"cpu\",\n \"jax\": \"0.9.2\",\n \"python\": \"3.14.3\",\n \"shape\": [\n 8\n ],\n \"dtype\": \"float32\",\n \"max_absolute_error\": 0.0,\n \"process_index\": 0,\n \"process_count\": 1\n}\nCompleted prediction: [0.34375, 0.8125, 1.28125, 1.75, 2.21875, 2.6875, 3.15625, 3.625]\n{\n \"expected\": \"cpu\",\n \"output_devices\": [\n \"TFRT_CPU_0\"\n ],\n \"platform\": \"cpu\",\n \"jax\": \"0.9.2\",\n \"python\": \"3.14.3\",\n \"shape\": [\n 8\n ],\n \"dtype\": \"float32\",\n \"max_absolute_error\": 0.0,\n \"process_index\": 0,\n \"process_count\": 1\n}\nCompleted prediction: [0.34375, 0.8125, 1.28125, 1.75, 2.21875, 2.6875, 3.15625, 3.625]\nMismatched platform report rejected.\nChanged bias, first prediction: -0.03125\nRow reversal preserved values and requested placement.\nPASS: welcome-03\n", "url": "executions/welcome-03.json" }, "url": "lesson.html?lesson=welcome-03" } ] }, { "id": "arrays", "number": "01", "title": "Arrays & pure functions", "part": "Shared foundations", "prerequisitePhaseIds": [ "welcome" ], "hardware": "CPU", "project": "Audit masked batch moments and padding-sensitive denominators", "checkpoint": "Predict output shapes and explain why a pure function is easier to transform.", "source": "https://docs.jax.dev/en/latest/beginner_guide.html", "description": "Treat every axis as part of a contract. Work through normalization, broadcasting and immutable updates, then diagnose mistakes that still return plausible numbers.", "learningAdvice": "Draw the input and output shapes before running each change.", "projectId": "arrays-audit", "projectStageIds": [ "1", "2", "3", "4" ], "studyGuide": { "question": "Can you spot a wrong axis even when the program runs?", "why": "Array bugs often return plausible numbers. Use unequal dimensions and hand-computed rows so that a successful operation cannot hide a mistaken meaning. Carry that habit into preprocessing and model inputs.", "readiness": { "prompt": "A matrix has three observations and two features. You need one mean per feature. Which axis should be reduced, and how many values remain?", "answer": "Reduce the observation axis, axis 0. Two feature means remain. Reducing axis 1 instead returns three observation means, which answers a different question.", "lessonId": "arrays-01" }, "milestones": [ { "lessonIds": [ "arrays-01", "arrays-02" ], "title": "Write the data contract", "task": "Annotate observation and feature axes; calculate one normalized row by hand. Compare a feature bias with an observation bias using unequal dimensions." }, { "lessonIds": [ "arrays-03", "arrays-04" ], "title": "Make preprocessing safe to reuse", "task": "Pass parameters explicitly, preserve the caller’s arrays and check indexed accumulation separately from replacement. Carry shape, mask and valid-count checks into foundation-toolkit stage 2." } ], "transfer": { "prompt": "Normalize a matrix with a constant feature and one masked observation. What should the constant feature and the mean denominator do?", "approach": "Define the constant-feature policy before division, for example a unit fallback scale. Compute statistics only from valid observations; masking the final output alone leaves the fitted statistics contaminated. Reject a feature with no valid observations or declare an explicit policy." }, "pitfall": { "symptom": "Predictions look reasonable but adding a bias changes the wrong dimension.", "check": "Write the two shapes before broadcasting and test one row and one column independently." }, "completion": "Retain unequal-shape, constant-column and masked-row cases, a NumPy or hand reference, and proof that inputs were unchanged. Explain every axis in the output.", "furtherWork": "Sparse and ragged inputs need additional contract-focused labs; this phase deliberately teaches fixed dense arrays first." }, "projectNote": "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.", "additionalProjectIds": [ "foundation-toolkit" ], "lessons": [ { "id": "arrays-01", "title": "From NumPy to jax.numpy", "status": "authored", "objective": "Build feature-wise preprocessing with fitted statistics and verify behavior on new and constant data.", "prerequisites": [], "exercise": "Foundation · Standardize the original data using its standard deviation plus epsilon 1e-6. Predict the transformed rows and check feature means and scales. Explain why the scale is only approximately one with epsilon.", "evidence": "Keep the fitted mean/scale, hand-derived training values, NumPy comparison, constant-feature test, and together-versus-separate new-batch results. Explain why re-fitting on each inference batch changes the problem.", "check": "Why does x.mean(axis=0) return two values?", "artifacts": { "source": "phases/01-arrays/01-from-numpy-to-jax-numpy/docs/en.md", "script": "exercises/arrays-01.py", "notebook": "exercises/arrays-01.ipynb", "markdown": "exercises/arrays-01.md", "scriptSource": "phases/01-arrays/01-from-numpy-to-jax-numpy/code/main.py", "notebookSource": "phases/01-arrays/01-from-numpy-to-jax-numpy/notebooks/exercise.ipynb", "quizSource": "phases/01-arrays/01-from-numpy-to-jax-numpy/quiz.json", "evidenceSource": "phases/01-arrays/01-from-numpy-to-jax-numpy/outputs/evidence.md", "evidence": "exercises/arrays-01-evidence.md" }, "path": "phases/01-arrays/01-from-numpy-to-jax-numpy", "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 45, "contentHash": "f4c7777e2ef504900f821d3167363071298f8ac8996e7ef334ab848310b2a8e0", "visualArtifact": { "image": "figures/arrays-01.svg", "png": "figures/arrays-01.png", "data": "figures/arrays-01.json", "generatedAt": "2026-10-20T10:45:07.729437+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "f4c7777e2ef504900f821d3167363071298f8ac8996e7ef334ab848310b2a8e0", "kind": "executed" }, "execution": { "contentHash": "f4c7777e2ef504900f821d3167363071298f8ac8996e7ef334ab848310b2a8e0", "executedAt": "2026-10-20T10:45:10.892111+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Mean: [ 3. 14.]\nCentered:\n [[-2. -4.]\n [ 0. 0.]\n [ 2. 4.]]\nnew batch transformed: [[2.4494896 2.4494896]]\nPASS: arrays-01\n", "url": "executions/arrays-01.json" }, "url": "lesson.html?lesson=arrays-01" }, { "id": "arrays-02", "title": "Shapes, broadcasting, and dtypes", "status": "authored", "objective": "Specify broadcasting contracts and detect silent pairwise-loss and dtype mistakes.", "prerequisites": [ "arrays-01" ], "exercise": "Add a different scalar offset to each row using offsets [100., 200.]. Make its shape explicit.", "evidence": "Keep an axis sketch, asymmetric offset checks, aligned-versus-pairwise residual arrays, a rejected shape mismatch, and the float32 precision observation. Explain the intended residual contract.", "check": "Which offset shape adds one scalar to each row of a (2, 3) batch?", "artifacts": { "source": "phases/01-arrays/02-shapes-broadcasting-and-dtypes/docs/en.md", "script": "exercises/arrays-02.py", "notebook": "exercises/arrays-02.ipynb", "markdown": "exercises/arrays-02.md", "scriptSource": "phases/01-arrays/02-shapes-broadcasting-and-dtypes/code/main.py", "notebookSource": "phases/01-arrays/02-shapes-broadcasting-and-dtypes/notebooks/exercise.ipynb", "quizSource": "phases/01-arrays/02-shapes-broadcasting-and-dtypes/quiz.json", "evidenceSource": "phases/01-arrays/02-shapes-broadcasting-and-dtypes/outputs/evidence.md", "evidence": "exercises/arrays-02-evidence.md" }, "path": "phases/01-arrays/02-shapes-broadcasting-and-dtypes", "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 50, "contentHash": "28d613b670b5073f1d70575c8995cf41882ab452ee78c3e449b38a9a90c754fb", "visualArtifact": { "image": "figures/arrays-02.svg", "png": "figures/arrays-02.png", "data": "figures/arrays-02.json", "generatedAt": "2026-10-20T10:45:08.535041+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "28d613b670b5073f1d70575c8995cf41882ab452ee78c3e449b38a9a90c754fb", "kind": "executed", "mechanismImage": "figures/arrays-02-mechanism.svg", "mechanismPng": "figures/arrays-02-mechanism.png" }, "execution": { "contentHash": "28d613b670b5073f1d70575c8995cf41882ab452ee78c3e449b38a9a90c754fb", "executedAt": "2026-10-20T10:45:12.332982+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "[[11. 22. 33.]\n [14. 25. 36.]]\nShape: (2, 3) dtype: float32\naligned residuals: [0. 0. 0.]\npairwise residuals: [[ 0. 2. 4.]\n [-2. 0. 2.]\n [-4. -2. 0.]]\nfloat32 large + 1 − large: 0.0\nPASS: arrays-02\n", "url": "executions/arrays-02.json" }, "url": "lesson.html?lesson=arrays-02" }, { "id": "arrays-03", "title": "Pure functions and explicit inputs", "status": "authored", "objective": "Design an explicit parameter/data interface and return new state without mutating old containers.", "prerequisites": [ "arrays-02" ], "exercise": "Create new parameters with weight 3 while retaining the original bias. Show that the original parameters still produce the original predictions.", "evidence": "Keep both parameter configurations, their known predictions, a container-alias experiment, a non-mutating update, and an evaluation loss checked by hand. Explain which state the caller owns.", "check": "What makes predict easier to transform and test?", "artifacts": { "source": "phases/01-arrays/03-pure-functions-and-explicit-inputs/docs/en.md", "script": "exercises/arrays-03.py", "notebook": "exercises/arrays-03.ipynb", "markdown": "exercises/arrays-03.md", "scriptSource": "phases/01-arrays/03-pure-functions-and-explicit-inputs/code/main.py", "notebookSource": "phases/01-arrays/03-pure-functions-and-explicit-inputs/notebooks/exercise.ipynb", "quizSource": "phases/01-arrays/03-pure-functions-and-explicit-inputs/quiz.json", "evidenceSource": "phases/01-arrays/03-pure-functions-and-explicit-inputs/outputs/evidence.md", "evidence": "exercises/arrays-03-evidence.md" }, "path": "phases/01-arrays/03-pure-functions-and-explicit-inputs", "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 45, "contentHash": "2c4762c09f36cc0219ca90135c498786acb4d7c7391a9ae6802d606dc2dcb1a8", "visualArtifact": { "image": "figures/arrays-03.svg", "png": "figures/arrays-03.png", "data": "figures/arrays-03.json", "generatedAt": "2026-10-20T10:45:09.208653+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "2c4762c09f36cc0219ca90135c498786acb4d7c7391a9ae6802d606dc2dcb1a8", "kind": "executed" }, "execution": { "contentHash": "2c4762c09f36cc0219ca90135c498786acb4d7c7391a9ae6802d606dc2dcb1a8", "executedAt": "2026-10-20T10:45:13.515971+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "[1. 3. 5.]\naliased container weight: 9.0\nevaluation losses: 0.0 6.0\nPASS: arrays-03\n", "url": "executions/arrays-03.json" }, "url": "lesson.html?lesson=arrays-03" }, { "id": "arrays-04", "title": "Immutable updates and indexing", "status": "authored", "objective": "Implement functional indexed updates and distinguish accumulated contributions from shape-preserving masks.", "prerequisites": [ "arrays-03" ], "exercise": "Replace the last two elements with zero in a new array. Keep x unchanged.", "evidence": "Keep original and updated arrays, a repeated-index integer reference, the whole-matrix replacement check, and compact-versus-fixed-shape outputs. Explain replacement, accumulation, and name rebinding.", "check": "After y = x.at[1].set(20.), what happens to x?", "artifacts": { "source": "phases/01-arrays/04-immutable-updates-and-indexing/docs/en.md", "script": "exercises/arrays-04.py", "notebook": "exercises/arrays-04.ipynb", "markdown": "exercises/arrays-04.md", "scriptSource": "phases/01-arrays/04-immutable-updates-and-indexing/code/main.py", "notebookSource": "phases/01-arrays/04-immutable-updates-and-indexing/notebooks/exercise.ipynb", "quizSource": "phases/01-arrays/04-immutable-updates-and-indexing/quiz.json", "evidenceSource": "phases/01-arrays/04-immutable-updates-and-indexing/outputs/evidence.md", "evidence": "exercises/arrays-04-evidence.md" }, "path": "phases/01-arrays/04-immutable-updates-and-indexing", "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 45, "contentHash": "c29171d688888a8baa96c70ad58fb86f9daaf8c58b12cceed84e79375ed6a3c8", "visualArtifact": { "image": "figures/arrays-04.svg", "png": "figures/arrays-04.png", "data": "figures/arrays-04.json", "generatedAt": "2026-10-20T10:45:09.991353+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "c29171d688888a8baa96c70ad58fb86f9daaf8c58b12cceed84e79375ed6a3c8", "kind": "executed" }, "execution": { "contentHash": "c29171d688888a8baa96c70ad58fb86f9daaf8c58b12cceed84e79375ed6a3c8", "executedAt": "2026-10-20T10:45:15.394008+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Original: [1. 2. 3. 4.]\nSet: [ 1. 20. 3. 4.]\nAdd: [6. 2. 8. 4.]\nPASS: arrays-04\n", "url": "executions/arrays-04.json" }, "url": "lesson.html?lesson=arrays-04" } ] }, { "id": "transforms", "number": "02", "title": "JAX transformations", "part": "Shared foundations", "prerequisitePhaseIds": [ "arrays" ], "hardware": "CPU", "project": "Compare scalar, vmap, autodiff, and changed-shape JIT behavior", "checkpoint": "Compose grad, vmap, and jit; explain what each transformation changes.", "source": "https://docs.jax.dev/en/latest/beginner_guide.html", "description": "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.", "learningAdvice": "Keep eager and transformed outputs plus a deliberately broken assumption.", "projectId": "transforms-audit", "projectStageIds": [ "1", "2", "3", "4" ], "studyGuide": { "question": "What changes when you differentiate, batch or compile a function?", "why": "Build one trustworthy scalar calculation before composing transformations. Differentiation changes the quantity, batching adds an axis, and compilation changes execution. Keeping these roles distinct makes later training failures easier to locate.", "readiness": { "prompt": "For a squared residual, is the gradient of the average loss the same shape as the stack of per-example gradients?", "answer": "No. A stack retains the examples axis; the gradient of the average reduces that axis. Their mean agrees for the same equally weighted differentiable objective.", "lessonId": "transforms-02" }, "milestones": [ { "lessonIds": [ "first-gradient", "transforms-02" ], "title": "Establish the derivative", "task": "Predict derivatives at several inputs and compare with an analytic calculation. Check that the loss is a scalar before differentiating." }, { "lessonIds": [ "transforms-03", "transforms-04", "transforms-05" ], "title": "Compose without changing the question", "task": "Compare a Python loop, mapped gradients and their compiled version on the same inputs. Change numerical values separately from shapes and static configuration; record tracing and timing as different observations." } ], "transfer": { "prompt": "One minibatch has two examples and another has six. How should you combine their mean gradients?", "approach": "Weight each mean gradient by its example count, add the weighted gradients, then divide by eight. An unweighted average gives examples in the smaller batch too much influence. Verify against a gradient over the concatenated batch." }, "pitfall": { "symptom": "A Python conditional works eagerly but raises a tracer error under jit.", "check": "Identify the data-dependent truth test. Use a JAX control-flow operation with compatible branch outputs; do not mark changing array data static to suppress the symptom." }, "completion": "Use foundation-toolkit stage 3. Keep independent derivatives, batched/compiled parity and one changed signature. Explain why a faster warm call does not measure first-call cost.", "furtherWork": "Transformation combinations involving nondifferentiable boundaries require further examples; passing smooth fixtures does not validate arbitrary gradients." }, "projectNote": "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.", "additionalProjectIds": [ "foundation-toolkit" ], "lessons": [ { "id": "first-gradient", "title": "Your first gradient", "status": "authored", "objective": "Explain local sensitivity and the chain rule, then compare analytic, numerical, and automatic derivatives.", "prerequisites": [], "exercise": "Foundation · Define f(x) = x³. Derive f′(x), predict the values at −2, 0, and 3, then verify every prediction with JAX. Explain why the zero at x = 0 is separate from a minimum.", "evidence": "Keep predictions and analytic derivations at multiple inputs, the finite-difference step/error sweep with dtype, both argnums checks, and the scalar-objective repair. Explain why a zero gradient need not mark a minimum.", "check": "At x = 0, jax.grad(lambda x: x ** 3)(0.0) returns zero. What can you conclude?", "artifacts": { "source": "phases/02-transforms/01-your-first-gradient/docs/en.md", "script": "exercises/first-gradient.py", "notebook": "exercises/first-gradient.ipynb", "markdown": "exercises/first-gradient.md", "scriptSource": "phases/02-transforms/01-your-first-gradient/code/main.py", "notebookSource": "phases/02-transforms/01-your-first-gradient/notebooks/exercise.ipynb", "quizSource": "phases/02-transforms/01-your-first-gradient/quiz.json", "evidenceSource": "phases/02-transforms/01-your-first-gradient/outputs/evidence.md", "evidence": "exercises/first-gradient-evidence.md" }, "path": "phases/02-transforms/01-your-first-gradient", "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 55, "contentHash": "c382307c1e6ed518fc942a959c2d6583641585dc4e4e727b7f7c5c754b8a500c", "visualArtifact": { "image": "figures/first-gradient.svg", "png": "figures/first-gradient.png", "data": "figures/first-gradient.json", "generatedAt": "2026-10-20T10:45:10.748798+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "c382307c1e6ed518fc942a959c2d6583641585dc4e4e727b7f7c5c754b8a500c", "kind": "executed" }, "execution": { "contentHash": "c382307c1e6ed518fc942a959c2d6583641585dc4e4e727b7f7c5c754b8a500c", "executedAt": "2026-10-20T10:45:17.292166+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Finite-difference estimate: 5.999999999999339\n6.0\nchain rule: -1.0 -12.0\nchain rule: 0.0 6.0\nchain rule: 2.0 42.0\nh / estimate / error: 0.1 27.009973526000977 0.009973526000976562\nh / estimate / error: 0.01 27.000045776367188 4.57763671875e-05\nh / estimate / error: 0.001 26.9975643157959 0.0024356842041015625\nh / estimate / error: 1e-05 27.0843505859375 0.0843505859375\nh / estimate / error: 1e-07 0.0 27.0\nPASS: first-gradient\n", "url": "executions/first-gradient.json" }, "url": "lesson.html?lesson=first-gradient" }, { "id": "transforms-02", "title": "Losses and value_and_grad", "status": "authored", "objective": "Derive and audit a scalar loss, its parameter gradient, and stable versus unstable updates.", "prerequisites": [ "first-gradient" ], "exercise": "Foundation · Starting from w = 1, calculate the next weight and loss for α = 0.1 and α = 0.3. Verify both. Explain the overshoot using the error multiplier rather than concluding that autodiff is wrong.", "evidence": "Keep handwritten loss/gradient arithmetic, checks at three weights, small/large-rate update results, auxiliary residuals, and mean-versus-sum batch-replication results. Explain the update-scale change.", "check": "You duplicate every example and keep the same learning rate. What happens to a sum-loss SGD update compared with the original batch?", "artifacts": { "source": "phases/02-transforms/02-losses-and-value-and-grad/docs/en.md", "script": "exercises/transforms-02.py", "notebook": "exercises/transforms-02.ipynb", "markdown": "exercises/transforms-02.md", "scriptSource": "phases/02-transforms/02-losses-and-value-and-grad/code/main.py", "notebookSource": "phases/02-transforms/02-losses-and-value-and-grad/notebooks/exercise.ipynb", "quizSource": "phases/02-transforms/02-losses-and-value-and-grad/quiz.json", "evidenceSource": "phases/02-transforms/02-losses-and-value-and-grad/outputs/evidence.md", "evidence": "exercises/transforms-02-evidence.md" }, "path": "phases/02-transforms/02-losses-and-value-and-grad", "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 55, "contentHash": "694b4a5e071b2909d41390d78eb340bbdc40813a451e782cd7794277c06609af", "visualArtifact": { "image": "figures/transforms-02.svg", "png": "figures/transforms-02.png", "data": "figures/transforms-02.json", "generatedAt": "2026-10-20T10:45:11.729076+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "694b4a5e071b2909d41390d78eb340bbdc40813a451e782cd7794277c06609af", "kind": "executed" }, "execution": { "contentHash": "694b4a5e071b2909d41390d78eb340bbdc40813a451e782cd7794277c06609af", "executedAt": "2026-10-20T10:45:19.594199+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Loss: 4.6666669845581055 gradient: -9.333333969116211\nweight / loss / gradient: 0.0 18.666667938232422 -18.666667938232422\nweight / loss / gradient: 1.0 4.6666669845581055 -9.333333969116211\nweight / loss / gradient: 2.0 0.0 0.0\nPASS: transforms-02\n", "url": "executions/transforms-02.json" }, "url": "lesson.html?lesson=transforms-02" }, { "id": "transforms-03", "title": "Batch a function with vmap", "status": "authored", "objective": "Specify batching axes and connect per-example gradients to a reduced batch objective.", "prerequisites": [ "transforms-02" ], "exercise": "Foundation · Compute each prediction gradient with respect to the shared weight. Predict the shape and the derivative of the dot product, then verify against the input batch.", "evidence": "Keep the shape contract, loop/vmap/matrix equivalence, row/column layout checks, per-example gradient array, and mean-gradient identity under two target sets. Explain what each gradient shape means.", "check": "For B examples and D shared parameters, what shapes do per-example loss gradients and the mean-loss gradient have?", "artifacts": { "source": "phases/02-transforms/03-batch-a-function-with-vmap/docs/en.md", "script": "exercises/transforms-03.py", "notebook": "exercises/transforms-03.ipynb", "markdown": "exercises/transforms-03.md", "scriptSource": "phases/02-transforms/03-batch-a-function-with-vmap/code/main.py", "notebookSource": "phases/02-transforms/03-batch-a-function-with-vmap/notebooks/exercise.ipynb", "quizSource": "phases/02-transforms/03-batch-a-function-with-vmap/quiz.json", "evidenceSource": "phases/02-transforms/03-batch-a-function-with-vmap/outputs/evidence.md", "evidence": "exercises/transforms-03-evidence.md" }, "path": "phases/02-transforms/03-batch-a-function-with-vmap", "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 60, "contentHash": "c306f3ef08ecee0f35915e560ca61c38fd59adc47cce3ac28ddb6b1336b694da", "visualArtifact": { "image": "figures/transforms-03.svg", "png": "figures/transforms-03.png", "data": "figures/transforms-03.json", "generatedAt": "2026-10-20T10:45:12.424996+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "c306f3ef08ecee0f35915e560ca61c38fd59adc47cce3ac28ddb6b1336b694da", "kind": "executed", "mechanismImage": "figures/transforms-03-mechanism.svg", "mechanismPng": "figures/transforms-03-mechanism.png" }, "execution": { "contentHash": "c306f3ef08ecee0f35915e560ca61c38fd59adc47cce3ac28ddb6b1336b694da", "executedAt": "2026-10-20T10:45:21.429738+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "[1. 1. 8.]\nloop / vmap / matmul: [1. 1. 8.] [1. 1. 8.] [1. 1. 8.]\nPASS: transforms-03\n", "url": "executions/transforms-03.json" }, "url": "lesson.html?lesson=transforms-03" }, { "id": "transforms-04", "title": "Compile a function with jit", "status": "authored", "objective": "Validate compiled loss/gradient values and produce a correctly synchronized latency report.", "prerequisites": [ "transforms-03" ], "exercise": "Foundation · Compile value_and_grad(loss), check the loss and gradient against both the eager transform and the manual matrix formula, and explain how you would synchronize a benchmark of its two outputs.", "evidence": "Keep the independent loss/gradient calculation, backend/shapes/dtype, first-call interval, warmup/repeat policy, median/range samples, and synchronization of every output leaf. State what the benchmark excludes.", "check": "Why wait with block_until_ready when timing a computation?", "artifacts": { "source": "phases/02-transforms/04-compile-a-function-with-jit/docs/en.md", "script": "exercises/transforms-04.py", "notebook": "exercises/transforms-04.ipynb", "markdown": "exercises/transforms-04.md", "scriptSource": "phases/02-transforms/04-compile-a-function-with-jit/code/main.py", "notebookSource": "phases/02-transforms/04-compile-a-function-with-jit/notebooks/exercise.ipynb", "quizSource": "phases/02-transforms/04-compile-a-function-with-jit/quiz.json", "evidenceSource": "phases/02-transforms/04-compile-a-function-with-jit/outputs/evidence.md", "evidence": "exercises/transforms-04-evidence.md" }, "path": "phases/02-transforms/04-compile-a-function-with-jit", "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 55, "contentHash": "2bde24a4e956087fd03b0c6123d05749b5932e233dfd3164bb8d2674884aa209", "visualArtifact": { "image": "figures/transforms-04.svg", "png": "figures/transforms-04.png", "data": "figures/transforms-04.json", "generatedAt": "2026-10-20T10:45:13.160817+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "2bde24a4e956087fd03b0c6123d05749b5932e233dfd3164bb8d2674884aa209", "kind": "executed" }, "execution": { "contentHash": "2bde24a4e956087fd03b0c6123d05749b5932e233dfd3164bb8d2674884aa209", "executedAt": "2026-10-20T10:45:23.194590+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Loss: 0.46000000834465027\nFirst/repeat seconds: 0.013900000136345625 1.920759677886963e-05\neager median / min / max seconds: 2.6375288143754005e-05 2.3334287106990814e-05 7.895799353718758e-05\ncompiled median / min / max seconds: 4.312489181756973e-06 3.874767571687698e-06 5.416106432676315e-06\nbackend / shapes / dtype: cpu (4, 3) (3,) (4,) float32\ncompiled value-and-grad median seconds: 6.770249456167221e-06\nPASS: transforms-04\n", "url": "executions/transforms-04.json" }, "url": "lesson.html?lesson=transforms-04" }, { "id": "transforms-05", "title": "Tracing, static arguments, and recompilation", "status": "authored", "objective": "Explain the staging boundary and repair configuration versus runtime-control-flow mistakes.", "prerequisites": [ "transforms-04" ], "exercise": "Foundation · Extend reduce_values with a static max mode while preserving sum and mean. Reject unknown modes. Verify all three modes on [1, 2, 3] and on [−4, −1, −2].", "evidence": "Keep the small jaxpr observation, the expected tracer Boolean error, checks of both repaired branches, all static reduction modes including invalid-mode rejection, and shape/value-change cases. Do not infer compilation counts from timing alone.", "check": "A frequently changing runtime scalar controls a branch. What is the most appropriate first repair for a traced Python if?", "artifacts": { "source": "phases/02-transforms/05-tracing-static-arguments-and-recompilation/docs/en.md", "script": "exercises/transforms-05.py", "notebook": "exercises/transforms-05.ipynb", "markdown": "exercises/transforms-05.md", "scriptSource": "phases/02-transforms/05-tracing-static-arguments-and-recompilation/code/main.py", "notebookSource": "phases/02-transforms/05-tracing-static-arguments-and-recompilation/notebooks/exercise.ipynb", "quizSource": "phases/02-transforms/05-tracing-static-arguments-and-recompilation/quiz.json", "evidenceSource": "phases/02-transforms/05-tracing-static-arguments-and-recompilation/outputs/evidence.md", "evidence": "exercises/transforms-05-evidence.md" }, "path": "phases/02-transforms/05-tracing-static-arguments-and-recompilation", "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 60, "contentHash": "c6f1554d6740750d8f73e870475f9174bf1d0e1e07f237ae7c5117313a5c3189", "visualArtifact": { "image": "figures/transforms-05.svg", "png": "figures/transforms-05.png", "data": "figures/transforms-05.json", "generatedAt": "2026-10-20T10:45:13.855517+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "c6f1554d6740750d8f73e870475f9174bf1d0e1e07f237ae7c5117313a5c3189", "kind": "executed" }, "execution": { "contentHash": "c6f1554d6740750d8f73e870475f9174bf1d0e1e07f237ae7c5117313a5c3189", "executedAt": "2026-10-20T10:45:24.879206+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Sum: 6.0\nMean: 2.0\nTracing sum_squares; shape: (3,)\n{ lambda ; a:f32[3]. let\n b:f32[3] = mul a a\n c:f32[] = reduce_sum[axes=(0,) out_sharding=None] b\n in (c,) }\nExpected dynamic Boolean conversion error\nPASS: transforms-05\n", "url": "executions/transforms-05.json" }, "url": "lesson.html?lesson=transforms-05" } ] }, { "id": "state", "number": "03", "title": "State, randomness & control flow", "part": "Shared foundations", "prerequisitePhaseIds": [ "transforms" ], "hardware": "CPU", "project": "A deterministic batched simulation", "checkpoint": "Replay the same experiment with saved state and keys; explain the update order.", "source": "https://docs.jax.dev/en/latest/beginner_guide.html", "description": "Make randomness, parameters and loop state explicit. Learn how to repeat an experiment without accidentally reusing randomness or mutating a caller’s state.", "learningAdvice": "Record the state before and after a transition and explain every changed leaf.", "projectId": "foundation-toolkit", "projectStageIds": [ "4" ], "studyGuide": { "question": "Have you saved everything needed to determine the next step?", "why": "A program’s state includes random keys, counters and structured values as well as model parameters. Explicit ownership lets you reproduce a transition, run a batch safely, and later recover an interrupted experiment.", "readiness": { "prompt": "Two simulated examples receive the same random key and run the same sampling function. Should you expect independently allocated draws?", "answer": "No. The same key and operation replay the same draw in the tested environment. Split or derive distinct keys at the ownership boundary; unequal observed samples alone do not prove independence.", "lessonId": "state-01" }, "milestones": [ { "lessonIds": [ "state-01", "state-02" ], "title": "Name and structure the state", "task": "Draw the key ownership tree and identify parameter, optimizer and metadata leaves. Check both tree structure and leaf values." }, { "lessonIds": [ "state-03", "state-04" ], "title": "Audit repeated and conditional transitions", "task": "Implement a recurrence with fixed carry shapes. Save per-step observations as outputs, compare with a Python loop, then test both conditional branches and their boundary." } ], "transfer": { "prompt": "Resume a simulation halfway through with the saved position but a newly created key. What must you compare to determine whether it is a true continuation?", "approach": "Compare the next draw and next state, then the remaining trajectory, against an uninterrupted reference. Restore the original key and counter as well as position. A similar final average is not a continuation check." }, "pitfall": { "symptom": "A compiled loop fails when its history grows each step.", "check": "Keep carry structure, shapes and dtypes fixed. Return the per-step history through scan outputs." }, "completion": "Use foundation-toolkit stage 4. Save and reload state in a fresh process and verify the next transition and full remaining trajectory, not only the final mean.", "furtherWork": "Local deterministic state is separate from concurrent-writer or multi-host recovery; the recovery and operations phases extend these boundaries." }, "projectNote": "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.", "lessons": [ { "id": "state-01", "title": "Random keys without surprises", "status": "authored", "objective": "Split keys and show reproducibility without accidental key reuse.", "prerequisites": [], "exercise": "Split a fresh seed into four keys, then use vmap to draw two normal values from each. Replay from the seed and compare.", "evidence": "Keep a key-ownership tree, two-step replay including the final key, the per-ID reorder check, and the broken/repaired repeated-noise loop.", "check": "What happens if the same key is used twice for the same random operation?", "artifacts": { "source": "phases/03-state/01-random-keys-without-surprises/docs/en.md", "script": "exercises/state-01.py", "notebook": "exercises/state-01.ipynb", "markdown": "exercises/state-01.md", "scriptSource": "phases/03-state/01-random-keys-without-surprises/code/main.py", "notebookSource": "phases/03-state/01-random-keys-without-surprises/notebooks/exercise.ipynb", "quizSource": "phases/03-state/01-random-keys-without-surprises/quiz.json", "evidenceSource": "phases/03-state/01-random-keys-without-surprises/outputs/evidence.md", "evidence": "exercises/state-01-evidence.md" }, "path": "phases/03-state/01-random-keys-without-surprises", "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 60, "contentHash": "7fbf1d2ae5134d8bbe8ec52905411be3e7bc758705b93da7fdd357c07e9df24a", "visualArtifact": { "image": "figures/state-01.svg", "png": "figures/state-01.png", "data": "figures/state-01.json", "generatedAt": "2026-10-20T10:45:14.614519+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "7fbf1d2ae5134d8bbe8ec52905411be3e7bc758705b93da7fdd357c07e9df24a", "kind": "executed" }, "execution": { "contentHash": "7fbf1d2ae5134d8bbe8ec52905411be3e7bc758705b93da7fdd357c07e9df24a", "executedAt": "2026-10-20T10:45:26.921075+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Replay equal: True\nFresh draw equal: False\nReplay equal: True\nFresh draw equal: False\nPASS: state-01\n", "url": "executions/state-01.json" }, "url": "lesson.html?lesson=state-01" }, { "id": "state-02", "title": "Pytrees and structured parameters", "status": "authored", "objective": "Represent parameters and state in a nested structure and transform its leaves.", "prerequisites": [ "state-01" ], "exercise": "Compute the squared norm of all gradient leaves. Verify it by adding the weight and bias contributions directly.", "evidence": "Keep the parameter/gradient tree diagram, hand-derived leaves and update, flatten/unflatten round trip, nested-scale derivative and missing-key repair.", "check": "What structure does grad(loss)(params) have?", "artifacts": { "source": "phases/03-state/02-pytrees-and-structured-parameters/docs/en.md", "script": "exercises/state-02.py", "notebook": "exercises/state-02.ipynb", "markdown": "exercises/state-02.md", "scriptSource": "phases/03-state/02-pytrees-and-structured-parameters/code/main.py", "notebookSource": "phases/03-state/02-pytrees-and-structured-parameters/notebooks/exercise.ipynb", "quizSource": "phases/03-state/02-pytrees-and-structured-parameters/quiz.json", "evidenceSource": "phases/03-state/02-pytrees-and-structured-parameters/outputs/evidence.md", "evidence": "exercises/state-02-evidence.md" }, "path": "phases/03-state/02-pytrees-and-structured-parameters", "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 60, "contentHash": "0a7645bdac396ac8fd5de8325968a22ca7d46e85fe4e484664f1e46880575d33", "visualArtifact": { "image": "figures/state-02.svg", "png": "figures/state-02.png", "data": "figures/state-02.json", "generatedAt": "2026-10-20T10:45:15.422039+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "0a7645bdac396ac8fd5de8325968a22ca7d46e85fe4e484664f1e46880575d33", "kind": "executed" }, "execution": { "contentHash": "0a7645bdac396ac8fd5de8325968a22ca7d46e85fe4e484664f1e46880575d33", "executedAt": "2026-10-20T10:45:28.661784+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Before/after: 4.0 0.6399999260902405\nBefore/after: 4.0 0.6399999260902405\nExpected tree mismatch\nPASS: state-02\n", "url": "executions/state-02.json" }, "url": "lesson.html?lesson=state-02" }, { "id": "state-03", "title": "Compiled loops with lax.scan", "status": "authored", "objective": "Replace a fixed-length Python loop with scan and compare final states.", "prerequisites": [ "state-02" ], "exercise": "Make the decay an explicit scalar argument. Differentiate the final value with respect to decay and compare at 0.5 with the analytic derivative 2.75.", "evidence": "Keep the labeled carry/output transition, Python-reference trajectory, polynomial derivative, compiled value check and growing-carry failure with a fixed-shape repair.", "check": "What must remain consistent across scan steps?", "artifacts": { "source": "phases/03-state/03-compiled-loops-with-lax-scan/docs/en.md", "script": "exercises/state-03.py", "notebook": "exercises/state-03.ipynb", "markdown": "exercises/state-03.md", "scriptSource": "phases/03-state/03-compiled-loops-with-lax-scan/code/main.py", "notebookSource": "phases/03-state/03-compiled-loops-with-lax-scan/notebooks/exercise.ipynb", "quizSource": "phases/03-state/03-compiled-loops-with-lax-scan/quiz.json", "evidenceSource": "phases/03-state/03-compiled-loops-with-lax-scan/outputs/evidence.md", "evidence": "exercises/state-03-evidence.md" }, "path": "phases/03-state/03-compiled-loops-with-lax-scan", "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 65, "contentHash": "f5aeddbef90d417df5b3a131bd0b5bda95f7b7da91cb151fb394a0ae5478ecba", "visualArtifact": { "image": "figures/state-03.svg", "png": "figures/state-03.png", "data": "figures/state-03.json", "generatedAt": "2026-10-20T10:45:16.178098+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "f5aeddbef90d417df5b3a131bd0b5bda95f7b7da91cb151fb394a0ae5478ecba", "kind": "executed", "mechanismImage": "figures/state-03-mechanism.svg", "mechanismPng": "figures/state-03-mechanism.png" }, "execution": { "contentHash": "f5aeddbef90d417df5b3a131bd0b5bda95f7b7da91cb151fb394a0ae5478ecba", "executedAt": "2026-10-20T10:45:30.167199+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "History: [1. 1.5 1.75 1.875]\nFinal: 1.875\nHistory: [1. 1.5 1.75 1.875]\nFinal: 1.875\nExpected carry-shape mismatch\nPASS: state-03\n", "url": "executions/state-03.json" }, "url": "lesson.html?lesson=state-03" }, { "id": "state-04", "title": "Branches with lax.cond", "status": "authored", "objective": "Implement a data-dependent branch that works inside a compiled function.", "prerequisites": [ "state-03" ], "exercise": "Write a compiled function that applies x * 2 when x is positive and x - 2 otherwise. Check positive, negative, and zero inputs.", "evidence": "Keep the branch-contract diagram, eager/jit failure and repair, boundary cases, batched values and off-boundary derivatives, and the shape-mismatch diagnosis.", "check": "Can one cond branch return a scalar and the other a length-three array?", "artifacts": { "source": "phases/03-state/04-branches-with-lax-cond/docs/en.md", "script": "exercises/state-04.py", "notebook": "exercises/state-04.ipynb", "markdown": "exercises/state-04.md", "scriptSource": "phases/03-state/04-branches-with-lax-cond/code/main.py", "notebookSource": "phases/03-state/04-branches-with-lax-cond/notebooks/exercise.ipynb", "quizSource": "phases/03-state/04-branches-with-lax-cond/quiz.json", "evidenceSource": "phases/03-state/04-branches-with-lax-cond/outputs/evidence.md", "evidence": "exercises/state-04-evidence.md" }, "path": "phases/03-state/04-branches-with-lax-cond", "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 60, "contentHash": "8c61aa6f27d09a4f9e41421e0d3d9d738afdcbe20a3966f0e812a0685a0d4c70", "visualArtifact": { "image": "figures/state-04.svg", "png": "figures/state-04.png", "data": "figures/state-04.json", "generatedAt": "2026-10-20T10:45:16.868842+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "8c61aa6f27d09a4f9e41421e0d3d9d738afdcbe20a3966f0e812a0685a0d4c70", "kind": "executed" }, "execution": { "contentHash": "8c61aa6f27d09a4f9e41421e0d3d9d738afdcbe20a3966f0e812a0685a0d4c70", "executedAt": "2026-10-20T10:45:31.964561+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "3.0\n3.0\nExpected traced Python Boolean failure\nExpected branch shape mismatch\nPASS: state-04\n", "url": "executions/state-04.json" }, "url": "lesson.html?lesson=state-04" } ] }, { "id": "optimization", "number": "04", "title": "Math & optimization", "part": "Shared foundations", "prerequisitePhaseIds": [ "state" ], "hardware": "CPU", "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.", "source": "https://docs.jax.dev/en/latest/automatic-differentiation.html", "projectId": "optimizer-audit", "description": "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.", "learningAdvice": "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.", "coreLessonIds": [ "optimization-01", "optimization-02", "optimization-03", "optimization-04" ], "projectStageIds": [ "1", "2", "3", "4", "5" ], "projectNote": "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.", "studyGuide": { "question": "Can you explain an optimizer’s behavior from the objective and data?", "why": "Use a model small enough to solve independently. Connect residuals, gradients, curvature and update state before comparing optimizer names. A lower training loss is useful evidence, but it does not by itself establish better held-out predictions.", "readiness": { "prompt": "A model predicts one value for each observation. Targets accidentally have an extra trailing singleton axis. Why might mean squared error still run?", "answer": "Broadcasting can construct a pairwise residual matrix instead of one residual per observation. Assert identical prediction and target shapes, then compute one residual by hand before reducing.", "lessonId": "optimization-01" }, "milestones": [ { "lessonIds": [ "optimization-05", "optimization-01", "optimization-06", "optimization-07", "optimization-02" ], "title": "Audit the objective and derivative", "task": "Work through vectors, residuals, rank and directional derivatives. Use a linear solve and a finite-difference sweep as independent checks." }, { "lessonIds": [ "optimization-08", "optimization-03", "optimization-09", "optimization-10" ], "title": "Connect dynamics to evaluation", "task": "Predict a stable versus divergent quadratic update, inspect stable likelihood arithmetic and separate fitting, validation selection and final test reporting." }, { "lessonIds": [ "optimization-11", "optimization-04", "optimization-12" ], "title": "Make optimizer comparisons fair", "task": "Count-weight unequal batches, preserve optimizer moments and schedule position, then change feature scale while retaining the prediction contract." } ], "transfer": { "prompt": "Rescale one feature by a factor of ten. Is the same learning rate automatically a fair comparison?", "approach": "No. The parameterization changes directional curvature. Map coefficients to preserve predictions, verify the objective and compare trajectories under stated rates. Fit preprocessing on training data only; never use test performance to choose the rate." }, "pitfall": { "symptom": "Gradient descent diverges while all values initially remain finite.", "check": "Check the update sign, loss reduction, feature scales and largest quadratic curvature before trying a different seed." }, "completion": "Complete regression-audit, then optimizer-audit. Keep an analytic reference, stable/unstable trajectories, a resumed optimizer comparison and an untouched final evaluation split.", "furtherWork": "Constrained, nonsmooth and large-scale stochastic optimization are not covered comprehensively. Current ranking evidence is limited to declared synthetic objectives." }, "lessons": [ { "id": "optimization-05", "title": "Vectors, norms, and projections", "status": "authored", "objective": "Read a dot product geometrically and separate a vector into a projection and an orthogonal residual.", "prerequisites": [], "exercise": "Project \\(u=(2,-1,2)\\) onto \\(a=(1,0,1)\\). Check the residual by hand first.", "evidence": "Keep the hand-computed projection, perpendicular-residual check, rescaling experiment, zero-direction repair and changed-basis calculation.", "check": "Does doubling a nonzero direction double the projection onto its line?", "path": "phases/04-optimization/05-vectors-norms-and-projections", "artifacts": { "source": "phases/04-optimization/05-vectors-norms-and-projections/docs/en.md", "script": "exercises/optimization-05.py", "notebook": "exercises/optimization-05.ipynb", "markdown": "exercises/optimization-05.md", "scriptSource": "phases/04-optimization/05-vectors-norms-and-projections/code/main.py", "notebookSource": "phases/04-optimization/05-vectors-norms-and-projections/notebooks/exercise.ipynb", "quizSource": "phases/04-optimization/05-vectors-norms-and-projections/quiz.json", "evidenceSource": "phases/04-optimization/05-vectors-norms-and-projections/outputs/evidence.md", "evidence": "exercises/optimization-05-evidence.md" }, "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 60, "contentHash": "d4f1f65e9be5e5231003b756f6a7dcaefb895c32fab266192c9e5f10120a7db0", "visualArtifact": { "image": "figures/optimization-05.svg", "png": "figures/optimization-05.png", "data": "figures/optimization-05.json", "generatedAt": "2026-10-20T10:45:17.642573+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "d4f1f65e9be5e5231003b756f6a7dcaefb895c32fab266192c9e5f10120a7db0", "kind": "executed" }, "execution": { "contentHash": "d4f1f65e9be5e5231003b756f6a7dcaefb895c32fab266192c9e5f10120a7db0", "executedAt": "2026-10-20T10:45:33.435928+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "projection / residual: [3. 0.] [0. 4.]\nprojection / residual: [3. 0.] [0. 4.]\nPASS: optimization-05\n", "url": "executions/optimization-05.json" }, "url": "lesson.html?lesson=optimization-05" }, { "id": "optimization-01", "title": "Linear algebra and loss intuition", "status": "authored", "objective": "Connect a matrix-vector product to predictions and a scalar regression loss.", "prerequisites": [ "optimization-05" ], "exercise": "Set the bias to zero. Predict each residual and calculate the loss by hand, then compare it with JAX.", "evidence": "Keep the row-by-row predictions, analytic weight/bias derivatives, pairwise-broadcast counterexample, and checked-loss rejection/repair. Include the analytic vector-gradient derivation and mean-versus-sum comparison.", "check": "What shape should one scalar target per observation have for predictions of shape (n,)?", "path": "phases/04-optimization/01-linear-algebra-and-loss-intuition", "artifacts": { "source": "phases/04-optimization/01-linear-algebra-and-loss-intuition/docs/en.md", "script": "exercises/optimization-01.py", "notebook": "exercises/optimization-01.ipynb", "markdown": "exercises/optimization-01.md", "scriptSource": "phases/04-optimization/01-linear-algebra-and-loss-intuition/code/main.py", "notebookSource": "phases/04-optimization/01-linear-algebra-and-loss-intuition/notebooks/exercise.ipynb", "quizSource": "phases/04-optimization/01-linear-algebra-and-loss-intuition/quiz.json", "evidenceSource": "phases/04-optimization/01-linear-algebra-and-loss-intuition/outputs/evidence.md", "evidence": "exercises/optimization-01-evidence.md" }, "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 80, "contentHash": "c9130b38274b14e717a5c3ce998deefe45cf5e21cc528c20c45872f0424a7771", "visualArtifact": { "image": "figures/optimization-01.svg", "png": "figures/optimization-01.png", "data": "figures/optimization-01.json", "generatedAt": "2026-10-20T10:45:18.387856+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "c9130b38274b14e717a5c3ce998deefe45cf5e21cc528c20c45872f0424a7771", "kind": "executed" }, "execution": { "contentHash": "c9130b38274b14e717a5c3ce998deefe45cf5e21cc528c20c45872f0424a7771", "executedAt": "2026-10-20T10:45:35.408082+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Predictions: [3. 0. 2.]\nLoss: 0.0\nPredictions: [3. 0. 2.]\nLoss: 0.0\nAnalytic gradients: [-1.3333334 -1.3333334] -2.0\nWrong residual shape: (3, 3)\nPASS: optimization-01\n", "url": "executions/optimization-01.json" }, "url": "lesson.html?lesson=optimization-01" }, { "id": "optimization-06", "title": "Least squares, rank, and conditioning", "status": "authored", "objective": "Fit a linear model with a direct solver and diagnose nonunique or sensitive coefficients.", "prerequisites": [ "optimization-01" ], "exercise": "Add a constant bias column to inputs \\((-1,0,1)\\) and fit targets \\((-1,1,3)\\). Recover slope \\(2\\) and bias \\(1\\).", "evidence": "Save the hand solution, residual orthogonality, duplicate-column ambiguity, perturbation table and NumPy comparison.", "check": "Can zero training error prove that fitted weights are uniquely determined?", "path": "phases/04-optimization/06-least-squares-rank-and-conditioning", "artifacts": { "source": "phases/04-optimization/06-least-squares-rank-and-conditioning/docs/en.md", "script": "exercises/optimization-06.py", "notebook": "exercises/optimization-06.ipynb", "markdown": "exercises/optimization-06.md", "scriptSource": "phases/04-optimization/06-least-squares-rank-and-conditioning/code/main.py", "notebookSource": "phases/04-optimization/06-least-squares-rank-and-conditioning/notebooks/exercise.ipynb", "quizSource": "phases/04-optimization/06-least-squares-rank-and-conditioning/quiz.json", "evidenceSource": "phases/04-optimization/06-least-squares-rank-and-conditioning/outputs/evidence.md", "evidence": "exercises/optimization-06-evidence.md" }, "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 75, "contentHash": "29d8442cd8365bb5d24c7021a429284ada0b407c960fc4a572409a98d41782c2", "visualArtifact": { "image": "figures/optimization-06.svg", "png": "figures/optimization-06.png", "data": "figures/optimization-06.json", "generatedAt": "2026-10-20T10:45:19.363128+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "29d8442cd8365bb5d24c7021a429284ada0b407c960fc4a572409a98d41782c2", "kind": "executed" }, "execution": { "contentHash": "29d8442cd8365bb5d24c7021a429284ada0b407c960fc4a572409a98d41782c2", "executedAt": "2026-10-20T10:45:37.409448+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "weights / rank / MSE: [0.6666669 1.6666665] 2 0.11111113\nweights / rank / MSE: [0.6666669 1.6666665] 2 0.11111113\ncondition / target change / weight change: 200.00504 0.0010000002 0.14142144\nPASS: optimization-06\n", "url": "executions/optimization-06.json" }, "url": "lesson.html?lesson=optimization-06" }, { "id": "optimization-07", "title": "The chain rule, Jacobians, and directions", "status": "authored", "objective": "Compute a small Jacobian by hand and connect forward sensitivities to reverse-mode gradients.", "prerequisites": [ "optimization-06" ], "exercise": "Replace the scalar loss by \\(L(x)=f_0(x)+2f_1(x)\\). Predict its gradient at \\((2,3)\\) before using `grad`.", "evidence": "Keep the handwritten Jacobian, JVP and VJP checks, adjoint identity, finite-direction check and stop-gradient diagnosis.", "check": "Which object does a reverse-mode pullback map to input sensitivity?", "path": "phases/04-optimization/07-chain-rule-jacobians-and-directional-derivatives", "artifacts": { "source": "phases/04-optimization/07-chain-rule-jacobians-and-directional-derivatives/docs/en.md", "script": "exercises/optimization-07.py", "notebook": "exercises/optimization-07.ipynb", "markdown": "exercises/optimization-07.md", "scriptSource": "phases/04-optimization/07-chain-rule-jacobians-and-directional-derivatives/code/main.py", "notebookSource": "phases/04-optimization/07-chain-rule-jacobians-and-directional-derivatives/notebooks/exercise.ipynb", "quizSource": "phases/04-optimization/07-chain-rule-jacobians-and-directional-derivatives/quiz.json", "evidenceSource": "phases/04-optimization/07-chain-rule-jacobians-and-directional-derivatives/outputs/evidence.md", "evidence": "exercises/optimization-07-evidence.md" }, "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 80, "contentHash": "1aaac0936877ab47a75da09c427f592954229043d5aa10544c154d7bba1644c4", "visualArtifact": { "image": "figures/optimization-07.svg", "png": "figures/optimization-07.png", "data": "figures/optimization-07.json", "generatedAt": "2026-10-20T10:45:20.529670+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "1aaac0936877ab47a75da09c427f592954229043d5aa10544c154d7bba1644c4", "kind": "executed" }, "execution": { "contentHash": "1aaac0936877ab47a75da09c427f592954229043d5aa10544c154d7bba1644c4", "executedAt": "2026-10-20T10:45:39.979446+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Jv / loss gradient: [3. 1.] [46. 19.]\nJv / loss gradient: [3. 1.] [46. 19.]\nPASS: optimization-07\n", "url": "executions/optimization-07.json" }, "url": "lesson.html?lesson=optimization-07" }, { "id": "optimization-02", "title": "Gradient checking and numerical accuracy", "status": "authored", "objective": "Compare autodiff with finite differences over several step sizes.", "prerequisites": [ "optimization-07" ], "exercise": "Check the gradient again at w = [2., -1.]. Explain why this checks a minimum but does not prove correctness everywhere.", "evidence": "Save analytic and numerical derivatives, the step-size/error table with dtype, the nonquadratic transfer check, and the omitted-term failure/repair. Add the one-sided slopes showing why a symmetric probe cannot certify differentiability at a kink.", "check": "Does passing one gradient check prove a loss implementation is correct for all inputs?", "path": "phases/04-optimization/02-gradient-checking-and-numerical-accuracy", "artifacts": { "source": "phases/04-optimization/02-gradient-checking-and-numerical-accuracy/docs/en.md", "script": "exercises/optimization-02.py", "notebook": "exercises/optimization-02.ipynb", "markdown": "exercises/optimization-02.md", "scriptSource": "phases/04-optimization/02-gradient-checking-and-numerical-accuracy/code/main.py", "notebookSource": "phases/04-optimization/02-gradient-checking-and-numerical-accuracy/notebooks/exercise.ipynb", "quizSource": "phases/04-optimization/02-gradient-checking-and-numerical-accuracy/quiz.json", "evidenceSource": "phases/04-optimization/02-gradient-checking-and-numerical-accuracy/outputs/evidence.md", "evidence": "exercises/optimization-02-evidence.md" }, "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 80, "contentHash": "616a8090bb257062b73fe11d64b5f732112bb7155aecaec1a9d904e73d8fb26d", "visualArtifact": { "image": "figures/optimization-02.svg", "png": "figures/optimization-02.png", "data": "figures/optimization-02.json", "generatedAt": "2026-10-20T10:45:21.778706+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "616a8090bb257062b73fe11d64b5f732112bb7155aecaec1a9d904e73d8fb26d", "kind": "executed" }, "execution": { "contentHash": "616a8090bb257062b73fe11d64b5f732112bb7155aecaec1a9d904e73d8fb26d", "executedAt": "2026-10-20T10:45:42.659689+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Finite difference: [-2.9999971 2.4999976]\nAutodiff: [-3. 2.5]\nFinite difference: [-2.9999971 2.4999976]\nAutodiff: [-3. 2.5]\nh / max error: 0.1 1.9073486328125e-06\nh / max error: 0.01 2.86102294921875e-06\nh / max error: 0.0001 0.0004978179931640625\nh / max error: 1e-08 3.0\nPASS: optimization-02\n", "url": "executions/optimization-02.json" }, "url": "lesson.html?lesson=optimization-02" }, { "id": "optimization-08", "title": "Curvature, Hessians, and learning rates", "status": "authored", "objective": "Use curvature to explain learning-rate stability, slow directions and why a zero gradient need not be a minimum.", "prerequisites": [ "optimization-02" ], "exercise": "For curvatures \\((2,8)\\), derive the stable interval and check one step at rate \\(0.1\\) from \\((1,1)\\).", "evidence": "Keep the Hessian and Hessian-vector checks, stable/divergent trajectories, rescaling comparison and saddle counterexample.", "check": "A smooth function has zero gradient at a point. What can we conclude?", "path": "phases/04-optimization/08-curvature-hessians-and-learning-rates", "artifacts": { "source": "phases/04-optimization/08-curvature-hessians-and-learning-rates/docs/en.md", "script": "exercises/optimization-08.py", "notebook": "exercises/optimization-08.ipynb", "markdown": "exercises/optimization-08.md", "scriptSource": "phases/04-optimization/08-curvature-hessians-and-learning-rates/code/main.py", "notebookSource": "phases/04-optimization/08-curvature-hessians-and-learning-rates/notebooks/exercise.ipynb", "quizSource": "phases/04-optimization/08-curvature-hessians-and-learning-rates/quiz.json", "evidenceSource": "phases/04-optimization/08-curvature-hessians-and-learning-rates/outputs/evidence.md", "evidence": "exercises/optimization-08-evidence.md" }, "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 75, "contentHash": "dabbf2599d986824d9d48912adf8a7ae87c89a4062fd291d1111a54b3405ce50", "visualArtifact": { "image": "figures/optimization-08.svg", "png": "figures/optimization-08.png", "data": "figures/optimization-08.json", "generatedAt": "2026-10-20T10:45:23.256983+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "dabbf2599d986824d9d48912adf8a7ae87c89a4062fd291d1111a54b3405ce50", "kind": "executed" }, "execution": { "contentHash": "dabbf2599d986824d9d48912adf8a7ae87c89a4062fd291d1111a54b3405ce50", "executedAt": "2026-10-20T10:45:46.372064+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "final weights / loss: [ 0.12157664 -0. ] 0.0073904395\nfinal weights / loss: [ 0.12157664 -0. ] 0.0073904395\nPASS: optimization-08\n", "url": "executions/optimization-08.json" }, "url": "lesson.html?lesson=optimization-08" }, { "id": "optimization-03", "title": "Write gradient descent yourself", "status": "authored", "objective": "Fit a regression model using explicit parameter updates.", "prerequisites": [ "optimization-08" ], "exercise": "Run ten steps with learning rate 1.1. Predict whether the error contracts before executing; show evidence of divergence.", "evidence": "Keep the two hand-computed steps, rate stability derivation, edge/divergent runs, scaled-curvature transfer and uphill-update repair. Include the tiny-update counterexample and explain your stopping criteria.", "check": "What makes rate 1.1 unstable for this particular quadratic?", "path": "phases/04-optimization/03-write-gradient-descent-yourself", "artifacts": { "source": "phases/04-optimization/03-write-gradient-descent-yourself/docs/en.md", "script": "exercises/optimization-03.py", "notebook": "exercises/optimization-03.ipynb", "markdown": "exercises/optimization-03.md", "scriptSource": "phases/04-optimization/03-write-gradient-descent-yourself/code/main.py", "notebookSource": "phases/04-optimization/03-write-gradient-descent-yourself/notebooks/exercise.ipynb", "quizSource": "phases/04-optimization/03-write-gradient-descent-yourself/quiz.json", "evidenceSource": "phases/04-optimization/03-write-gradient-descent-yourself/outputs/evidence.md", "evidence": "exercises/optimization-03-evidence.md" }, "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 80, "contentHash": "23c8351c018c9c1948e3c6eb785f9548d587e9f92d0143f8942926c0e73dfd37", "visualArtifact": { "image": "figures/optimization-03.svg", "png": "figures/optimization-03.png", "data": "figures/optimization-03.json", "generatedAt": "2026-10-20T10:45:24.128758+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "23c8351c018c9c1948e3c6eb785f9548d587e9f92d0143f8942926c0e73dfd37", "kind": "executed" }, "execution": { "contentHash": "23c8351c018c9c1948e3c6eb785f9548d587e9f92d0143f8942926c0e73dfd37", "executedAt": "2026-10-20T10:45:48.289484+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Final weight: 1.9999969005584717\nFinal loss: 9.606537787476555e-12\nFinal weight: 1.9999969005584717\nFinal loss: 9.606537787476555e-12\nFirst two post-update losses: [2.5600002 1.6384 ]\nPASS: optimization-03\n", "url": "executions/optimization-03.json" }, "url": "lesson.html?lesson=optimization-03" }, { "id": "optimization-09", "title": "Probability, likelihood, and stable losses", "status": "authored", "objective": "Derive binary cross-entropy from a probability model and compute it safely from logits.", "prerequisites": [ "optimization-03" ], "exercise": "Compute the binary losses for positive labels at probabilities \\(0.8\\) and \\(0.2\\), and explain which prediction is better supported.", "evidence": "Keep the hand log-loss, analytic gradient, extreme-logit failure/repair, multiclass shift check and empirical-probability calculation.", "check": "Why compute binary cross-entropy from logits?", "path": "phases/04-optimization/09-probability-likelihood-and-stable-losses", "artifacts": { "source": "phases/04-optimization/09-probability-likelihood-and-stable-losses/docs/en.md", "script": "exercises/optimization-09.py", "notebook": "exercises/optimization-09.ipynb", "markdown": "exercises/optimization-09.md", "scriptSource": "phases/04-optimization/09-probability-likelihood-and-stable-losses/code/main.py", "notebookSource": "phases/04-optimization/09-probability-likelihood-and-stable-losses/notebooks/exercise.ipynb", "quizSource": "phases/04-optimization/09-probability-likelihood-and-stable-losses/quiz.json", "evidenceSource": "phases/04-optimization/09-probability-likelihood-and-stable-losses/outputs/evidence.md", "evidence": "exercises/optimization-09-evidence.md" }, "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 80, "contentHash": "4bd1d215e2058da7f81dab777779167ff3d3789614687d8c6fb570d27d09fdb5", "visualArtifact": { "image": "figures/optimization-09.svg", "png": "figures/optimization-09.png", "data": "figures/optimization-09.json", "generatedAt": "2026-10-20T10:45:25.280717+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "4bd1d215e2058da7f81dab777779167ff3d3789614687d8c6fb570d27d09fdb5", "kind": "executed" }, "execution": { "contentHash": "4bd1d215e2058da7f81dab777779167ff3d3789614687d8c6fb570d27d09fdb5", "executedAt": "2026-10-20T10:45:52.039722+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "uncertain loss / confident-wrong loss: 0.6931472 100.0\nuncertain loss / confident-wrong loss: 0.6931472 100.0\nPASS: optimization-09\n", "url": "executions/optimization-09.json" }, "url": "lesson.html?lesson=optimization-09" }, { "id": "optimization-10", "title": "Regularization and honest validation", "status": "authored", "objective": "Derive a ridge solution, distinguish data loss from penalized loss, and select a penalty using held-out data.", "prerequisites": [ "optimization-09" ], "exercise": "For a target slope of \\(4\\) and penalty \\(\\lambda=1\\), derive the fitted slope and verify its gradient.", "evidence": "Keep the ridge gradient/solve comparison, penalty sweep with separate train/validation losses, duplication check and bias-mask repair.", "check": "Which loss should select the penalty in this exercise?", "path": "phases/04-optimization/10-regularization-and-validation", "artifacts": { "source": "phases/04-optimization/10-regularization-and-validation/docs/en.md", "script": "exercises/optimization-10.py", "notebook": "exercises/optimization-10.ipynb", "markdown": "exercises/optimization-10.md", "scriptSource": "phases/04-optimization/10-regularization-and-validation/code/main.py", "notebookSource": "phases/04-optimization/10-regularization-and-validation/notebooks/exercise.ipynb", "quizSource": "phases/04-optimization/10-regularization-and-validation/quiz.json", "evidenceSource": "phases/04-optimization/10-regularization-and-validation/outputs/evidence.md", "evidence": "exercises/optimization-10-evidence.md" }, "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 75, "contentHash": "6bf7a85d3ed2a0796b63581a9cb039bef40354dc743aa040000223ad9679dc24", "visualArtifact": { "image": "figures/optimization-10.svg", "png": "figures/optimization-10.png", "data": "figures/optimization-10.json", "generatedAt": "2026-10-20T10:45:26.508946+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "6bf7a85d3ed2a0796b63581a9cb039bef40354dc743aa040000223ad9679dc24", "kind": "executed" }, "execution": { "contentHash": "6bf7a85d3ed2a0796b63581a9cb039bef40354dc743aa040000223ad9679dc24", "executedAt": "2026-10-20T10:45:55.153882+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "train / validation / objective: 1.0 0.0 3.0\ntrain / validation / objective: 1.0 0.0 3.0\npenalty / train / validation: [0. 0.5 2. ] [0. 1. 4.] [4. 0. 4.]\nPASS: optimization-10\n", "url": "executions/optimization-10.json" }, "url": "lesson.html?lesson=optimization-10" }, { "id": "optimization-11", "title": "Minibatches, expectation, and gradient noise", "status": "authored", "objective": "Explain an unbiased minibatch gradient and verify how sampling and reduction affect its variance.", "prerequisites": [ "optimization-10" ], "exercise": "Enumerate all \\(64\\) ordered batches of size \\(3\\). Check their mean and variance against the independent-sampling formula.", "evidence": "Keep the exact per-example mean, enumerated batch-variance calculation, uneven-batch repair, PRNG replay and gradient-accumulation check.", "check": "When does averaging batch means reproduce the dataset mean?", "path": "phases/04-optimization/11-minibatches-expectation-and-gradient-noise", "artifacts": { "source": "phases/04-optimization/11-minibatches-expectation-and-gradient-noise/docs/en.md", "script": "exercises/optimization-11.py", "notebook": "exercises/optimization-11.ipynb", "markdown": "exercises/optimization-11.md", "scriptSource": "phases/04-optimization/11-minibatches-expectation-and-gradient-noise/code/main.py", "notebookSource": "phases/04-optimization/11-minibatches-expectation-and-gradient-noise/notebooks/exercise.ipynb", "quizSource": "phases/04-optimization/11-minibatches-expectation-and-gradient-noise/quiz.json", "evidenceSource": "phases/04-optimization/11-minibatches-expectation-and-gradient-noise/outputs/evidence.md", "evidence": "exercises/optimization-11-evidence.md" }, "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 75, "contentHash": "885230e1c6579d0de2f25823226239c1f009704193cc86771883ac02aa1b67a7", "visualArtifact": { "image": "figures/optimization-11.svg", "png": "figures/optimization-11.png", "data": "figures/optimization-11.json", "generatedAt": "2026-10-20T10:45:27.716408+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "885230e1c6579d0de2f25823226239c1f009704193cc86771883ac02aa1b67a7", "kind": "executed" }, "execution": { "contentHash": "885230e1c6579d0de2f25823226239c1f009704193cc86771883ac02aa1b67a7", "executedAt": "2026-10-20T10:45:58.425895+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "full gradient / single variance / pair variance: -2.5 1.25 0.625\nsample indices / sample gradient: [3 2] -3.5\nfull gradient / single variance / pair variance: -2.5 1.25 0.625\nsample indices / sample gradient: [3 2] -3.5\nPASS: optimization-11\n", "url": "executions/optimization-11.json" }, "url": "lesson.html?lesson=optimization-11" }, { "id": "optimization-04", "title": "Optimize with Optax", "status": "authored", "objective": "Replace the handwritten update with an optimizer and track its state.", "prerequisites": [ "optimization-11" ], "exercise": "At the original zero parameters, compare the first Optax SGD update with an explicit tree-map gradient descent step.", "evidence": "Save the analytic first update, manual/Optax equivalence, retained/reset momentum comparison, continuation replay and double-subtraction repair. Include the zero-gradient momentum step and explain why retained state still moves the parameter.", "check": "What should a training step retain besides updated parameters for a stateful optimizer?", "path": "phases/04-optimization/04-optimize-with-optax", "artifacts": { "source": "phases/04-optimization/04-optimize-with-optax/docs/en.md", "script": "exercises/optimization-04.py", "notebook": "exercises/optimization-04.ipynb", "markdown": "exercises/optimization-04.md", "scriptSource": "phases/04-optimization/04-optimize-with-optax/code/main.py", "notebookSource": "phases/04-optimization/04-optimize-with-optax/notebooks/exercise.ipynb", "quizSource": "phases/04-optimization/04-optimize-with-optax/quiz.json", "evidenceSource": "phases/04-optimization/04-optimize-with-optax/outputs/evidence.md", "evidence": "exercises/optimization-04-evidence.md" }, "source": "https://optax.readthedocs.io/en/latest/getting_started.html", "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 80, "contentHash": "fcfca2398f2a506bc0c8c7cd1392b3ce8d316fda0fcf48ab5a0a5af8167f0406", "visualArtifact": { "image": "figures/optimization-04.svg", "png": "figures/optimization-04.png", "data": "figures/optimization-04.json", "generatedAt": "2026-10-20T10:45:28.871552+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "fcfca2398f2a506bc0c8c7cd1392b3ce8d316fda0fcf48ab5a0a5af8167f0406", "kind": "executed" }, "execution": { "contentHash": "fcfca2398f2a506bc0c8c7cd1392b3ce8d316fda0fcf48ab5a0a5af8167f0406", "executedAt": "2026-10-20T10:46:01.015273+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Parameters: {'bias': Array(1., dtype=float32), 'weight': Array(1.9999996, dtype=float32)}\nLoss: 5.126058404787345e-14\nParameters: {'bias': Array(1., dtype=float32), 'weight': Array(1.9999996, dtype=float32)}\nLoss: 5.126058404787345e-14\nFirst parameters: {'bias': Array(0.4, dtype=float32), 'weight': Array(0.29333338, dtype=float32)}\nPASS: optimization-04\n", "url": "executions/optimization-04.json" }, "url": "lesson.html?lesson=optimization-04" }, { "id": "optimization-12", "title": "Adam, clipping, and learning-rate schedules", "status": "authored", "objective": "Explain Adam’s moments, inspect clipping and schedule order, and compare optimizers with a controlled update budget.", "prerequisites": [ "optimization-04" ], "exercise": "Compare global norm clipping with coordinate-wise clipping on \\((3,4)\\). Check which preserves the ratio of the coordinates.", "evidence": "Keep two hand-checked Adam steps, clipping-direction checks, schedule values and boundary, equal-budget loss traces, and the optimization decision memo.", "check": "Does clipping gradients before Adam guarantee that final updates obey the same norm threshold?", "path": "phases/04-optimization/12-adam-clipping-and-learning-rate-schedules", "artifacts": { "source": "phases/04-optimization/12-adam-clipping-and-learning-rate-schedules/docs/en.md", "script": "exercises/optimization-12.py", "notebook": "exercises/optimization-12.ipynb", "markdown": "exercises/optimization-12.md", "scriptSource": "phases/04-optimization/12-adam-clipping-and-learning-rate-schedules/code/main.py", "notebookSource": "phases/04-optimization/12-adam-clipping-and-learning-rate-schedules/notebooks/exercise.ipynb", "quizSource": "phases/04-optimization/12-adam-clipping-and-learning-rate-schedules/quiz.json", "evidenceSource": "phases/04-optimization/12-adam-clipping-and-learning-rate-schedules/outputs/evidence.md", "evidence": "exercises/optimization-12-evidence.md" }, "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 90, "contentHash": "b53c9001cf8f5566eab55a47c9ceaa664afbbb5b019c4df9da6f0bdba176da8f", "visualArtifact": { "image": "figures/optimization-12.svg", "png": "figures/optimization-12.png", "data": "figures/optimization-12.json", "generatedAt": "2026-10-20T10:45:30.516857+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "b53c9001cf8f5566eab55a47c9ceaa664afbbb5b019c4df9da6f0bdba176da8f", "kind": "executed", "mechanismImage": "figures/optimization-12-mechanism.svg", "mechanismPng": "figures/optimization-12-mechanism.png" }, "execution": { "contentHash": "b53c9001cf8f5566eab55a47c9ceaa664afbbb5b019c4df9da6f0bdba176da8f", "executedAt": "2026-10-20T10:46:04.755758+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "step / update: 1 [-0.09999993 0.09999993]\nstep / update: 2 [-0.09334469 0.00893815]\nstep / update: 1 [-0.09999993 0.09999993]\nstep / update: 2 [-0.09334469 0.00893815]\nschedule rates: [0.1 0.0775 0.055 0.0325 0.01 0.01 ]\nsgd initial / final loss: 5.5 0.0029602647\nmomentum initial / final loss: 5.5 0.024945451\nadam initial / final loss: 5.5 0.00012754407\nPASS: optimization-12\n", "url": "executions/optimization-12.json" }, "url": "lesson.html?lesson=optimization-12" } ] }, { "id": "networks", "number": "05", "title": "Neural network training", "part": "Training systems", "prerequisitePhaseIds": [ "optimization" ], "hardware": "CPU first; accelerator for larger experiments", "project": "Build and audit a neural classifier", "checkpoint": "Keep train and evaluation behavior separate and reproduce metrics from a saved configuration.", "source": "https://flax.readthedocs.io/", "projectId": "mlp-classifier", "description": "Build the same ideas into small neural networks. Understand hidden layers, parameter state, stable objectives and evaluation before moving to a small image classifier.", "learningAdvice": "Explain a loss curve together with held-out errors; a lower training loss alone is not sufficient.", "studyGuide": { "question": "Did the network learn the task, or exploit an evaluation mistake?", "why": "Move from explicit matrix operations to stateful modules while preserving independent checks. Learn to distinguish capacity, optimization, data quality and evaluation errors before adding layers.", "readiness": { "prompt": "A classifier predicts the majority class on every example and reports high accuracy. What information is missing?", "answer": "Class counts and per-class errors are missing. A majority baseline can score highly on imbalanced data while missing every rare example. Inspect the confusion matrix and the denominator for each class.", "lessonId": "networks-04" }, "milestones": [ { "lessonIds": [ "networks-01", "networks-02" ], "title": "Connect equations to module state", "task": "Check a forward pass and unreduced loss independently, then identify parameter and nonparameter state in the NNX model." }, { "lessonIds": [ "networks-03", "networks-04" ], "title": "Train and evaluate under separate contracts", "task": "Keep held-out examples out of updates, verify evaluation leaves state unchanged and read confusion rows as true-class denominators." }, { "lessonIds": [ "networks-05" ], "title": "Diagnose one failure at a time", "task": "Use recorded input ranges, losses, gradient norms and update sizes to distinguish a scale problem from a data or shape problem." } ], "transfer": { "prompt": "Training loss falls, but recall for a rare class is zero. Write the next three checks before changing the model.", "approach": "Inspect class counts and labels; compare a constant predictor and per-class confusion; then inspect score distributions and choose any threshold on validation data. Report final metrics on held-out data without using it to tune the repair." }, "pitfall": { "symptom": "Evaluation results change even though you did not call an optimizer.", "check": "Compare state snapshots and train/eval flags; stochastic or running-statistic behavior may still be active." }, "completion": "Use mlp-classifier before the larger image harness. Retain independent forward/loss checks, two runs with declared seeds, state-pure evaluation and class-specific failure analysis.", "furtherWork": "Natural-image dataset ingestion, richer augmentation studies and pretrained vision transfer remain larger follow-up labs; synthetic images are separate from real-photo generalization." }, "projectNote": "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.", "additionalProjectIds": [ "reproducible-experiments" ], "lessons": [ { "id": "networks-01", "title": "Build a tiny multilayer perceptron", "status": "authored", "objective": "Implement a small model with explicit parameters before adding a model library.", "prerequisites": [], "exercise": "Complement XOR class labels by negating the output layer; verify logits and predictions on new coordinates. Diagnose pairwise loss broadcasting.", "evidence": "Keep the 33-parameter shape calculation, independent NumPy cross-entropy, train and held-out seeds, initial/final loss, nearby held-out accuracy, and one repaired score/label shape failure.", "check": "Explain why removing tanh makes two affine layers collapse to a linear XOR boundary, then verify a changed label convention.", "path": "phases/05-networks/01-build-a-tiny-multilayer-perceptron", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/05-networks/01-build-a-tiny-multilayer-perceptron/docs/en.md", "script": "exercises/networks-01.py", "notebook": "exercises/networks-01.ipynb", "markdown": "exercises/networks-01.md", "scriptSource": "phases/05-networks/01-build-a-tiny-multilayer-perceptron/code/main.py", "notebookSource": "phases/05-networks/01-build-a-tiny-multilayer-perceptron/notebooks/exercise.ipynb", "quizSource": "phases/05-networks/01-build-a-tiny-multilayer-perceptron/quiz.json", "evidenceSource": "phases/05-networks/01-build-a-tiny-multilayer-perceptron/outputs/evidence.md", "evidence": "exercises/networks-01-evidence.md" }, "minutes": 85, "contentHash": "ad4066cdd0c2addbba6417f189b99f4fa4b7de66a73aeb3ab2fbc5ebdd65dbc0", "visualArtifact": { "image": "figures/networks-01.svg", "png": "figures/networks-01.png", "data": "figures/networks-01.json", "generatedAt": "2026-10-20T10:45:32.084291+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "ad4066cdd0c2addbba6417f189b99f4fa4b7de66a73aeb3ab2fbc5ebdd65dbc0", "kind": "executed" }, "execution": { "contentHash": "ad4066cdd0c2addbba6417f189b99f4fa4b7de66a73aeb3ab2fbc5ebdd65dbc0", "executedAt": "2026-10-20T10:46:11.350755+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Training loss: 0.7047799825668335 -> 0.0015878621488809586 held-out accuracy: 1.0\nTraining loss: 0.7047799825668335 -> 0.0015878621488809586 held-out accuracy: 1.0\nConsistent hidden permutation preserves logits; one-sided permutation changes them.\nPASS: networks-01\n", "url": "executions/networks-01.json" }, "url": "lesson.html?lesson=networks-01" }, { "id": "networks-02", "title": "Model and state with Flax NNX", "status": "authored", "objective": "Express a model and inspect the distinction between parameters and other state.", "prerequisites": [ "networks-01" ], "exercise": "Separate nnx.Param and nonparameter state; make an independent copied-state snapshot and diagnose shared-variable mutation.", "evidence": "Keep the NumPy forward comparison, count of 13 trainable scalars and one diagnostic counter, counter independence after copying state, and logits for batch size one and five.", "check": "Explain why exclusion from Param gradients does not prevent forward state mutation, and demonstrate safe snapshot independence.", "path": "phases/05-networks/02-model-and-state-with-flax-nnx", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/05-networks/02-model-and-state-with-flax-nnx/docs/en.md", "script": "exercises/networks-02.py", "notebook": "exercises/networks-02.ipynb", "markdown": "exercises/networks-02.md", "scriptSource": "phases/05-networks/02-model-and-state-with-flax-nnx/code/main.py", "notebookSource": "phases/05-networks/02-model-and-state-with-flax-nnx/notebooks/exercise.ipynb", "quizSource": "phases/05-networks/02-model-and-state-with-flax-nnx/quiz.json", "evidenceSource": "phases/05-networks/02-model-and-state-with-flax-nnx/outputs/evidence.md", "evidence": "exercises/networks-02-evidence.md" }, "minutes": 80, "contentHash": "81334e7f41bd4ff3205536759c986bf631b1fb32e7e2a387c45df4c8f6cbe8f7", "visualArtifact": { "image": "figures/networks-02.svg", "png": "figures/networks-02.png", "data": "figures/networks-02.json", "generatedAt": "2026-10-20T10:45:33.696240+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "81334e7f41bd4ff3205536759c986bf631b1fb32e7e2a387c45df4c8f6cbe8f7", "kind": "executed" }, "execution": { "contentHash": "81334e7f41bd4ff3205536759c986bf631b1fb32e7e2a387c45df4c8f6cbe8f7", "executedAt": "2026-10-20T10:46:15.263948+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Trainable scalars: 13; other state: 1; clone counters: 1 2\nTrainable scalars: 13; other state: 1; clone counters: 1 2\nClone logits shifted by one; original parameters unchanged.\nPASS: networks-02\n", "url": "executions/networks-02.json" }, "url": "lesson.html?lesson=networks-02" }, { "id": "networks-03", "title": "A compiled train and evaluation step", "status": "authored", "objective": "Build train and eval steps and check that evaluation does not update weights.", "prerequisites": [ "networks-02" ], "exercise": "Compile separate training and evaluation functions, replay 80 updates, and reconstruct held-out metrics from equal-size chunks without changing state.", "evidence": "Keep pre/post parameter snapshots around evaluation, optimizer step count 80, independent held-out loss and accuracy, complete replay comparison, and a deliberately bad held-out update on an isolated clone.", "check": "Identify which weights a pre-update loss describes and prove evaluation changes neither parameters nor optimizer step count.", "path": "phases/05-networks/03-a-compiled-train-and-evaluation-step", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/05-networks/03-a-compiled-train-and-evaluation-step/docs/en.md", "script": "exercises/networks-03.py", "notebook": "exercises/networks-03.ipynb", "markdown": "exercises/networks-03.md", "scriptSource": "phases/05-networks/03-a-compiled-train-and-evaluation-step/code/main.py", "notebookSource": "phases/05-networks/03-a-compiled-train-and-evaluation-step/notebooks/exercise.ipynb", "quizSource": "phases/05-networks/03-a-compiled-train-and-evaluation-step/quiz.json", "evidenceSource": "phases/05-networks/03-a-compiled-train-and-evaluation-step/outputs/evidence.md", "evidence": "exercises/networks-03-evidence.md" }, "minutes": 95, "contentHash": "89b4eb211777ba4d0940434d6a4e300cbf46dababad5dc486afe2c0ca678105c", "visualArtifact": { "image": "figures/networks-03.svg", "png": "figures/networks-03.png", "data": "figures/networks-03.json", "generatedAt": "2026-10-20T10:45:35.572198+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "89b4eb211777ba4d0940434d6a4e300cbf46dababad5dc486afe2c0ca678105c", "kind": "executed" }, "execution": { "contentHash": "89b4eb211777ba4d0940434d6a4e300cbf46dababad5dc486afe2c0ca678105c", "executedAt": "2026-10-20T10:46:21.256922+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Updates: 80 held-out loss: 0.018426384776830673 accuracy: 1.0\nUpdates: 80 held-out loss: 0.018426384776830673 accuracy: 1.0\nCount-weighted held-out metrics: [0.01842638 0.99999995] constructed accuracy: 0.5416666666666666\nPASS: networks-03\n", "url": "executions/networks-03.json" }, "url": "lesson.html?lesson=networks-03" }, { "id": "networks-04", "title": "Train a CNN on a small dataset", "status": "authored", "objective": "Train a classifier with a fixed data split and report held-out metrics.", "prerequisites": [ "networks-03" ], "exercise": "Verify a convolution patch manually, train the noisy-bar CNN on a fixed split, transfer to 10-by-10 input shapes, and repair NCHW/NHWC confusion.", "evidence": "Keep the 8-by-8-to-6-by-6 shape diagram, NumPy patch dot product, seeds 20/21, independent held-out cross-entropy, 24-observation confusion matrix, constant-class baseline, and layout failure repair.", "check": "Explain what global spatial mean pooling discards and verify that confusion counts and trace reproduce the reported held-out accuracy.", "path": "phases/05-networks/04-train-a-cnn-on-a-small-dataset", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/05-networks/04-train-a-cnn-on-a-small-dataset/docs/en.md", "script": "exercises/networks-04.py", "notebook": "exercises/networks-04.ipynb", "markdown": "exercises/networks-04.md", "scriptSource": "phases/05-networks/04-train-a-cnn-on-a-small-dataset/code/main.py", "notebookSource": "phases/05-networks/04-train-a-cnn-on-a-small-dataset/notebooks/exercise.ipynb", "quizSource": "phases/05-networks/04-train-a-cnn-on-a-small-dataset/quiz.json", "evidenceSource": "phases/05-networks/04-train-a-cnn-on-a-small-dataset/outputs/evidence.md", "evidence": "exercises/networks-04-evidence.md" }, "minutes": 110, "contentHash": "d1e8c9915efd949a1601e7b22b00a7c21ba365e4e699d9cf35d95e316d8f69ce", "visualArtifact": { "image": "figures/networks-04.svg", "png": "figures/networks-04.png", "data": "figures/networks-04.json", "generatedAt": "2026-10-20T10:45:37.932032+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "d1e8c9915efd949a1601e7b22b00a7c21ba365e4e699d9cf35d95e316d8f69ce", "kind": "executed", "mechanismImage": "figures/networks-04-mechanism.svg", "mechanismPng": "figures/networks-04-mechanism.png" }, "execution": { "contentHash": "d1e8c9915efd949a1601e7b22b00a7c21ba365e4e699d9cf35d95e316d8f69ce", "executedAt": "2026-10-20T10:46:26.728813+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Held-out loss: 0.015700625 accuracy: 1.0 confusion:\n [[12 0]\n [ 0 12]]\nHeld-out loss: 0.015700625 accuracy: 1.0 confusion:\n [[12 0]\n [ 0 12]]\nTrue-class counts: [12 12] row-normalized confusion: [[1. 0.]\n [0. 1.]]\nAbsent class recall: undefined, not zero.\nPASS: networks-04\n", "url": "executions/networks-04.json" }, "url": "lesson.html?lesson=networks-04" }, { "id": "networks-05", "title": "Debug unstable learning", "status": "authored", "objective": "Introduce a bad learning rate and diagnose it using loss and gradient evidence.", "prerequisites": [ "networks-04" ], "exercise": "Derive the curvature-dependent stable rate, reproduce an alternating divergent trajectory, and distinguish scaled features and corrupted inputs.", "evidence": "Keep the analytic error recurrence, per-step weight/loss/gradient norm for rates 0.1 and 0.3, observed 3.24 loss multiplier, changed-scale comparison, clipping bound, and one repaired nonfinite-input case.", "check": "Derive why scaling features and targets by ten changes this squared-loss curvature by 100 without treating that bound as universal.", "path": "phases/05-networks/05-debug-unstable-learning", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/05-networks/05-debug-unstable-learning/docs/en.md", "script": "exercises/networks-05.py", "notebook": "exercises/networks-05.ipynb", "markdown": "exercises/networks-05.md", "scriptSource": "phases/05-networks/05-debug-unstable-learning/code/main.py", "notebookSource": "phases/05-networks/05-debug-unstable-learning/notebooks/exercise.ipynb", "quizSource": "phases/05-networks/05-debug-unstable-learning/quiz.json", "evidenceSource": "phases/05-networks/05-debug-unstable-learning/outputs/evidence.md", "evidence": "exercises/networks-05-evidence.md" }, "minutes": 85, "contentHash": "4dc504459d1c38a540eab0c4b283d33182a60234995d845e318339e82219ce30", "visualArtifact": { "image": "figures/networks-05.svg", "png": "figures/networks-05.png", "data": "figures/networks-05.json", "generatedAt": "2026-10-20T10:45:39.045241+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "4dc504459d1c38a540eab0c4b283d33182a60234995d845e318339e82219ce30", "kind": "executed" }, "execution": { "contentHash": "4dc504459d1c38a540eab0c4b283d33182a60234995d845e318339e82219ce30", "executedAt": "2026-10-20T10:46:28.762700+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Stable final weight: 1.999999761581421 unstable pre-update [weight,loss,|gradient|]:\n [[ 0.00000000e+00 1.86666679e+01 1.86666679e+01]\n [ 5.60000038e+00 6.04800110e+01 3.36000061e+01]\n [-4.48000145e+00 1.95955292e+02 6.04800148e+01]\n [ 1.36640034e+01 6.34895264e+02 1.08864037e+02]\n [-1.89952106e+01 2.05706152e+03 1.95955322e+02]\n [ 3.97913895e+01 6.66488281e+03 3.52719635e+02]]\nStable final weight: 1.999999761581421 unstable pre-update [weight,loss,|gradient|]:\n [[ 0.00000000e+00 1.86666679e+01 1.86666679e+01]\n [ 5.60000038e+00 6.04800110e+01 3.36000061e+01]\n [-4.48000145e+00 1.95955292e+02 6.04800148e+01]\n [ 1.36640034e+01 6.34895264e+02 1.08864037e+02]\n [-1.89952106e+01 2.05706152e+03 1.95955322e+02]\n [ 3.97913895e+01 6.66488281e+03 3.52719635e+02]]\nReference / compensated / unchanged-rate updates: 1.8833334 1.8833332 8.416667\nPASS: networks-05\n", "url": "executions/networks-05.json" }, "url": "lesson.html?lesson=networks-05" } ] }, { "id": "recovery", "number": "06", "title": "Data & checkpoint recovery", "part": "Training systems", "prerequisitePhaseIds": [ "networks" ], "hardware": "CPU first; TPU for system measurements", "project": "Resume optimizer and PRNG state across an interrupted training update", "checkpoint": "Compare resumed and uninterrupted runs under a stated determinism tolerance.", "source": "https://orbax.readthedocs.io/en/latest/", "description": "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.", "learningAdvice": "Compare sample IDs, optimizer memory, random state and losses after restore.", "projectId": "recovery-audit", "projectStageIds": [ "1", "2", "3", "4" ], "projectNote": "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.", "additionalProjectIds": [ "engineering-release", "sharded-training", "reproducible-experiments" ], "studyGuide": { "question": "Can an interrupted run make exactly the same next update?", "why": "Recovery is a contract between data order, model state and publication. Tracking explains which run produced an artifact; a checkpoint supplies the values needed to continue it. Neither replaces the other.", "readiness": { "prompt": "The weights restore exactly, but the next loss differs. Which state should you inspect before blaming numerical nondeterminism?", "answer": "Compare optimizer moments and count, random keys, data permutation/cursor, preprocessing and configuration. The same weights with a different next batch are a different transition.", "lessonId": "recovery-03" }, "milestones": [ { "lessonIds": [ "recovery-01", "recovery-02" ], "title": "Preserve record identity and position", "task": "Inspect record IDs through shuffling and batching. State the final-batch policy and save iterator state only at the intended consumption boundary." }, { "lessonIds": [ "recovery-03", "recovery-04", "recovery-05" ], "title": "Prove a committed continuation", "task": "Save complete state, restore in a fresh process and compare future IDs and updates across an epoch boundary. Reject an incomplete or incompatible checkpoint before publishing it." }, { "lessonIds": [ "recovery-06" ], "title": "Connect continuation to lineage", "task": "Retrieve an MLflow run and its artifact, compare data/source identities and metric-step conventions, and distinguish a registry selection from active deployment." } ], "transfer": { "prompt": "The checkpoint is valid, but the dataset was reordered without changing its filename. Is the saved cursor sufficient?", "approach": "No. The cursor is meaningful relative to a dataset identity and ordering configuration. Verify content/version identity first; reject mismatches rather than silently resuming on different records." }, "pitfall": { "symptom": "The advertised checkpoint step is newer than the restored state.", "check": "Inspect completion and publication order. An asynchronous save request is not proof that the artifact is durable or restorable." }, "completion": "Use sharded-training recovery stages and the MLflow extension. Keep uninterrupted and restarted trajectories, a corrupt/incompatible rejection, and a retrieved artifact tied to its run.", "furtherWork": "Network storage failures, concurrent checkpoint writers and true multi-host recovery still need separate failure-injection qualification." }, "lessons": [ { "id": "recovery-01", "title": "Design an input pipeline", "status": "authored", "objective": "Build and audit a deterministic batched host pipeline.", "prerequisites": [], "exercise": "Pad the final three-label batch to four elements with zero. Show that an unmasked mean differs, then derive the correct masked mean.", "evidence": "Keep the exact epoch/example ID order, uneven final-batch sizes, hand-derived masked versus unmasked label mean, changed seed order and duplicate/drop audit.", "check": "What information is missing from a model-only checkpoint for mid-epoch recovery?", "path": "phases/06-recovery/01-design-an-input-pipeline", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/06-recovery/01-design-an-input-pipeline/docs/en.md", "script": "exercises/recovery-01.py", "notebook": "exercises/recovery-01.ipynb", "markdown": "exercises/recovery-01.md", "scriptSource": "phases/06-recovery/01-design-an-input-pipeline/code/main.py", "notebookSource": "phases/06-recovery/01-design-an-input-pipeline/notebooks/exercise.ipynb", "quizSource": "phases/06-recovery/01-design-an-input-pipeline/quiz.json", "evidenceSource": "phases/06-recovery/01-design-an-input-pipeline/outputs/evidence.md", "evidence": "exercises/recovery-01-evidence.md" }, "minutes": 70, "contentHash": "01efd384f4d1b167b7564b54c62321b1138e40d4d092dcc9975d133f0321e6ce", "visualArtifact": { "image": "figures/recovery-01.svg", "png": "figures/recovery-01.png", "data": "figures/recovery-01.json", "generatedAt": "2026-10-20T10:45:39.678672+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "01efd384f4d1b167b7564b54c62321b1138e40d4d092dcc9975d133f0321e6ce", "kind": "executed" }, "execution": { "contentHash": "01efd384f4d1b167b7564b54c62321b1138e40d4d092dcc9975d133f0321e6ce", "executedAt": "2026-10-20T10:46:29.735773+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Batch sizes: [4, 4, 3]\nExample order: [9, 10, 4, 1, 0, 7, 6, 8, 2, 5, 3]\nPreparation seconds (local observation only): 0.00836424995213747\nDataset SHA256: 1a81451a8fb32a1bf0ad6307ea19198153ab1603b37174c76c4d498e343e6892\nBatch sizes: [4, 4, 3]\nExample order: [9, 10, 4, 1, 0, 7, 6, 8, 2, 5, 3]\nPreparation seconds (local observation only): 1.874985173344612e-05\nDataset SHA256: 1a81451a8fb32a1bf0ad6307ea19198153ab1603b37174c76c4d498e343e6892\nEpoch one order: [4, 0, 5, 10, 6, 2, 1, 9, 7, 3, 8]\nOmitted IDs: [2, 3, 5]\nCopy preserves fingerprint; reordered source and edited labels invalidate it.\nPASS: recovery-01\n", "url": "executions/recovery-01.json" }, "url": "lesson.html?lesson=recovery-01" }, { "id": "recovery-02", "title": "Load, batch, and checkpoint order with Grain", "status": "authored", "objective": "Build a Grain source/sampler/batch pipeline and verify its record order.", "prerequisites": [ "recovery-01" ], "exercise": "Consume two batches, snapshot, and replay the third batch with a new compatible iterator.", "evidence": "Keep Grain sampler configuration and exact example IDs for full and partial batches; save iterator bytes, recreate the loader, and compare the next batch before/after restore. Record the incompatible configuration failure.", "check": "What makes a fresh loader resume at the saved next batch?", "path": "phases/06-recovery/02-load-batch-and-prefetch-with-grain", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/06-recovery/02-load-batch-and-prefetch-with-grain/docs/en.md", "script": "exercises/recovery-02.py", "notebook": "exercises/recovery-02.ipynb", "markdown": "exercises/recovery-02.md", "scriptSource": "phases/06-recovery/02-load-batch-and-prefetch-with-grain/code/main.py", "notebookSource": "phases/06-recovery/02-load-batch-and-prefetch-with-grain/notebooks/exercise.ipynb", "quizSource": "phases/06-recovery/02-load-batch-and-prefetch-with-grain/quiz.json", "evidenceSource": "phases/06-recovery/02-load-batch-and-prefetch-with-grain/outputs/evidence.md", "evidence": "exercises/recovery-02-evidence.md" }, "minutes": 75, "contentHash": "6742a4d4033d726195b9cfacf994dd74246ec6fbcd882f42ccd13651b834e4e6", "visualArtifact": { "image": "figures/recovery-02.svg", "png": "figures/recovery-02.png", "data": "figures/recovery-02.json", "generatedAt": "2026-10-20T10:45:41.460755+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "6742a4d4033d726195b9cfacf994dd74246ec6fbcd882f42ccd13651b834e4e6", "kind": "executed" }, "execution": { "contentHash": "6742a4d4033d726195b9cfacf994dd74246ec6fbcd882f42ccd13651b834e4e6", "executedAt": "2026-10-20T10:46:37.199821+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "First IDs: [8, 6, 7, 9]\nNext IDs: [0, 5, 1, 2]\nIterator state bytes: 329\nRestored next batch matches\nFirst IDs: [8, 6, 7, 9]\nNext IDs: [0, 5, 1, 2]\nIterator state bytes: 329\nRestored next batch matches\nSame seed restarts; it does not resume\nEach epoch covers all twelve IDs\nNext-epoch IDs replayed: [[2, 0, 11, 1], [5, 7, 6, 8]]\nExpected sampler compatibility rejection\nPASS: recovery-02\n", "url": "executions/recovery-02.json" }, "url": "lesson.html?lesson=recovery-02" }, { "id": "recovery-03", "title": "Save model, optimizer, and random state", "status": "authored", "objective": "Checkpoint complete training state with Orbax and replay the next update.", "prerequisites": [ "recovery-02" ], "exercise": "Advance to step five, save to a new directory, restore and verify step six against uninterrupted execution.", "evidence": "Keep Orbax save/wait/restore commands, snapshot step, Adam count/moments, exact PRNG words, numerical next-update replay and the parameter-only restart counterexample.", "check": "Which checkpoint supports replaying the next Adam update with stochastic targets?", "path": "phases/06-recovery/03-save-model-optimizer-and-random-state", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/06-recovery/03-save-model-optimizer-and-random-state/docs/en.md", "script": "exercises/recovery-03.py", "notebook": "exercises/recovery-03.ipynb", "markdown": "exercises/recovery-03.md", "scriptSource": "phases/06-recovery/03-save-model-optimizer-and-random-state/code/main.py", "notebookSource": "phases/06-recovery/03-save-model-optimizer-and-random-state/notebooks/exercise.ipynb", "quizSource": "phases/06-recovery/03-save-model-optimizer-and-random-state/quiz.json", "evidenceSource": "phases/06-recovery/03-save-model-optimizer-and-random-state/outputs/evidence.md", "evidence": "exercises/recovery-03-evidence.md" }, "minutes": 100, "contentHash": "7a9557db16fcc99d2f36eb35d15264536a412380c94f1a1ecada3fae051b2b13", "visualArtifact": { "image": "figures/recovery-03.svg", "png": "figures/recovery-03.png", "data": "figures/recovery-03.json", "generatedAt": "2026-10-20T10:45:43.384535+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "7a9557db16fcc99d2f36eb35d15264536a412380c94f1a1ecada3fae051b2b13", "kind": "executed" }, "execution": { "contentHash": "7a9557db16fcc99d2f36eb35d15264536a412380c94f1a1ecada3fae051b2b13", "executedAt": "2026-10-20T10:46:42.170123+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Restored step: 3\nOptimizer count: 3\nNext update and loss agree; checkpoint exists: True\nRestored step: 3\nOptimizer count: 3\nNext update and loss agree; checkpoint exists: True\nModel-only checkpoint produces a different Adam update\nReset key changes the noisy objective and continuation stream\nFull restore Adam count: 4; reset memory count: 1; next parameters differ.\nExpected changed configuration rejection\nExpected exact key-data mismatch\nPASS: recovery-03\n", "url": "executions/recovery-03.json" }, "url": "lesson.html?lesson=recovery-03" }, { "id": "recovery-04", "title": "Recover data position and resume", "status": "authored", "objective": "Persist model/optimizer/key and Grain iterator position at one boundary.", "prerequisites": [ "recovery-03" ], "exercise": "Repeat the protocol after one batch and compare the next five batches through step six. Save to a distinct directory.", "evidence": "Keep the committed model/optimizer/key/iterator snapshot, uninterrupted and resumed batch-ID ledger across the epoch boundary, exact keys/counts, declared floating tolerances and inconsistent model-only/reader-only restart failures.", "check": "Which evidence supports a consistent restart boundary?", "path": "phases/06-recovery/04-recover-data-position-and-resume", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/06-recovery/04-recover-data-position-and-resume/docs/en.md", "script": "exercises/recovery-04.py", "notebook": "exercises/recovery-04.ipynb", "markdown": "exercises/recovery-04.md", "scriptSource": "phases/06-recovery/04-recover-data-position-and-resume/code/main.py", "notebookSource": "phases/06-recovery/04-recover-data-position-and-resume/notebooks/exercise.ipynb", "quizSource": "phases/06-recovery/04-recover-data-position-and-resume/quiz.json", "evidenceSource": "phases/06-recovery/04-recover-data-position-and-resume/outputs/evidence.md", "evidence": "exercises/recovery-04-evidence.md" }, "minutes": 110, "contentHash": "24da684f527e7b04d4809cc6a23ec0be5882e57eed5b0bb5aa7f9b6ded984682", "visualArtifact": { "image": "figures/recovery-04.svg", "png": "figures/recovery-04.png", "data": "figures/recovery-04.json", "generatedAt": "2026-10-20T10:45:45.668045+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "24da684f527e7b04d4809cc6a23ec0be5882e57eed5b0bb5aa7f9b6ded984682", "kind": "executed", "mechanismImage": "figures/recovery-04-mechanism.svg", "mechanismPng": "figures/recovery-04-mechanism.png" }, "execution": { "contentHash": "24da684f527e7b04d4809cc6a23ec0be5882e57eed5b0bb5aa7f9b6ded984682", "executedAt": "2026-10-20T10:46:51.862987+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Restored next four ID batches: [[4, 10, 11, 3], [2, 0, 11, 1], [5, 7, 6, 8], [4, 3, 10, 9]]\nRestored next four losses: [0.5496272444725037, 0.8390463590621948, 0.07169045507907867, 0.32427603006362915]\nFull state agrees at step six\nRestored next four ID batches: [[4, 10, 11, 3], [2, 0, 11, 1], [5, 7, 6, 8], [4, 3, 10, 9]]\nRestored next four losses: [0.5496272444725037, 0.8390463590621948, 0.07169045507907867, 0.32427603006362915]\nFull state agrees at step six\nMissing iterator restore repeats earlier IDs and changes the next loss\nCorrect IDs alone do not restore the training transition\nRejected incompatible seed\nRejected incompatible sha256\nOriginal run contract remains valid.\nExpected batch-configuration mismatch\nPASS: recovery-04\n", "url": "executions/recovery-04.json" }, "url": "lesson.html?lesson=recovery-04" }, { "id": "recovery-05", "title": "Asynchronous checkpoints and failure boundaries", "status": "authored", "objective": "Distinguish requested, completed, verified and published checkpoint state. Continue a functional update while an asynchronous writer is alive.", "prerequisites": [ "recovery-04" ], "exercise": "Verify and publish the already written second checkpoint. Restore it and compare the next functional transition against the live state. Keep the original snapshot unchanged.", "evidence": "Keep the requested/completed/accepted step log, the restored snapshot comparison, interrupted-publication result and surfaced finalization error. Explain why directory presence alone is separate from acceptance.", "check": "A new checkpoint directory exists but the completion wait raises an exception. Which recovery decision follows this lesson’s protocol?", "path": "phases/06-recovery/05-asynchronous-checkpoints-and-failure-boundaries", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/06-recovery/05-asynchronous-checkpoints-and-failure-boundaries/docs/en.md", "script": "exercises/recovery-05.py", "notebook": "exercises/recovery-05.ipynb", "markdown": "exercises/recovery-05.md", "scriptSource": "phases/06-recovery/05-asynchronous-checkpoints-and-failure-boundaries/code/main.py", "notebookSource": "phases/06-recovery/05-asynchronous-checkpoints-and-failure-boundaries/notebooks/exercise.ipynb", "quizSource": "phases/06-recovery/05-asynchronous-checkpoints-and-failure-boundaries/quiz.json", "evidenceSource": "phases/06-recovery/05-asynchronous-checkpoints-and-failure-boundaries/outputs/evidence.md", "evidence": "exercises/recovery-05-evidence.md" }, "minutes": 90, "contentHash": "53af93b1c6c37ad3ae60af050b3fbce1aec596ad453fe6c06ad32042e883b6b8", "visualArtifact": { "image": "figures/recovery-05.svg", "png": "figures/recovery-05.png", "data": "figures/recovery-05.json", "generatedAt": "2026-10-20T10:45:46.880275+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "53af93b1c6c37ad3ae60af050b3fbce1aec596ad453fe6c06ad32042e883b6b8", "kind": "executed" }, "execution": { "contentHash": "53af93b1c6c37ad3ae60af050b3fbce1aec596ad453fe6c06ad32042e883b6b8", "executedAt": "2026-10-20T10:46:54.698510+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Checkpoint workspace: /var/folders/mp/_mq7srqx2y10p5nhjz9m7hj40000gn/T/jax-async-lesson-9ibeu4bo\n{\n \"save_call_seconds\": 0.0035320832394063473,\n \"remaining_work_and_wait_seconds\": 0.017399583011865616,\n \"accepted_step\": 1,\n \"live_step\": 2\n}\nUnpublished step 2 exists; application recovery still selects verified step 1.\nCheckpoint workspace: /var/folders/mp/_mq7srqx2y10p5nhjz9m7hj40000gn/T/jax-async-lesson-xq01n_nx\n{\n \"save_call_seconds\": 0.0024104579351842403,\n \"remaining_work_and_wait_seconds\": 0.009226541966199875,\n \"accepted_step\": 1,\n \"live_step\": 2\n}\nUnpublished step 2 exists; application recovery still selects verified step 1.\nExpected background failure surfaced: injected finalization callback failure\nStep two verified and accepted; next transition agrees.\nMissing target rejected; accepted checkpoint unchanged.\nPASS: recovery-05\n", "url": "executions/recovery-05.json" }, "url": "lesson.html?lesson=recovery-05" }, { "id": "recovery-06", "title": "Track experiments and lineage with MLflow", "status": "authored", "objective": "Track, retrieve and compare two JAX experiments using explicit data identity, metric steps and downloaded model artifacts.", "prerequisites": [ "recovery-05" ], "exercise": "Add a third, deliberately wrong model artifact and show why the smallest training loss alone is insufficient to select a release.", "evidence": "Keep both MLflow run IDs, source/data hashes, complete loss histories, the selection metric and predictions independently recomputed from the downloaded artifact.", "check": "A run is tagged champion. What does that establish?", "path": "phases/06-recovery/06-track-experiments-and-lineage-with-mlflow", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/06-recovery/06-track-experiments-and-lineage-with-mlflow/docs/en.md", "script": "exercises/recovery-06.py", "notebook": "exercises/recovery-06.ipynb", "markdown": "exercises/recovery-06.md", "scriptSource": "phases/06-recovery/06-track-experiments-and-lineage-with-mlflow/code/main.py", "notebookSource": "phases/06-recovery/06-track-experiments-and-lineage-with-mlflow/notebooks/exercise.ipynb", "quizSource": "phases/06-recovery/06-track-experiments-and-lineage-with-mlflow/quiz.json", "evidenceSource": "phases/06-recovery/06-track-experiments-and-lineage-with-mlflow/outputs/evidence.md", "evidence": "exercises/recovery-06-evidence.md" }, "minutes": 70, "contentHash": "2045fd46cce599b92c2b22fca98a4f64f22b5ca7fa8a52502d6adfac11e0dfdc", "visualArtifact": { "image": "figures/recovery-06.svg", "png": "figures/recovery-06.png", "data": "figures/recovery-06.json", "generatedAt": "2026-10-20T10:45:49.856694+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "2045fd46cce599b92c2b22fca98a4f64f22b5ca7fa8a52502d6adfac11e0dfdc", "kind": "executed" }, "execution": { "contentHash": "2045fd46cce599b92c2b22fca98a4f64f22b5ca7fa8a52502d6adfac11e0dfdc", "executedAt": "2026-10-20T10:47:02.065912+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Validation MSE by learning rate: {0.02: 0.11994461715221405, 0.15: 5.049471951679152e-10}\nVerified 30 logged steps per run; selected and reloaded the better validation candidate.\nValidation difference: 0.11994461664726686\nWrong zero model validation MSE: 6.0\nChanged target changes dataset identity.\nPASS: recovery-06\n", "url": "executions/recovery-06.json" }, "url": "lesson.html?lesson=recovery-06" } ] }, { "id": "transformers", "number": "07", "title": "Transformers & language models", "part": "Training systems", "prerequisitePhaseIds": [ "recovery" ], "hardware": "CPU for small blocks; TPU or GPU for training experiments", "project": "A checkpointed causal language model", "checkpoint": "Verify masking, sequence handling, restoration, and generation before comparing speed.", "source": "https://flax.readthedocs.io/", "description": "Build attention and causal masking from small arrays, then connect tokenization, mixed precision, training and generation.", "learningAdvice": "Inspect what each token can attend to and count only valid targets in the loss.", "projectId": "text-harness", "projectStageIds": [ "1", "2", "3", "4" ], "projectNote": "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.", "studyGuide": { "question": "Can you prove that a predicted token cannot see its future?", "why": "Attention is a weighted mixing operation whose axes and masks are testable. Build a causal model from these contracts before interpreting a generated sentence as evidence of learning.", "readiness": { "prompt": "A query has equal scores for three allowed keys. What are the attention weights, and what happens if one key is masked?", "answer": "The weights are equal thirds. Masking one key leaves equal halves on the two allowed keys and zero on the masked key. Normalize over keys, not queries.", "lessonId": "transformers-01" }, "milestones": [ { "lessonIds": [ "transformers-01", "transformers-02" ], "title": "Audit attention and targets", "task": "Check one attention row by hand, token/target shifts, padding denominators and segment boundaries. Perturb future tokens and inspect earlier outputs." }, { "lessonIds": [ "transformers-03", "transformers-04" ], "title": "Train, resume and decode one small model", "task": "Preserve residual shapes and sensitive arithmetic, compare the next resumed update and state the decoding policy. Continue into the text harness for cache/full-prefix parity." } ], "transfer": { "prompt": "Two documents are packed into one sequence. Causal masking passes, but one document still changes the other’s predictions. What is missing?", "approach": "Causality only blocks future positions. Require matching segment identities too, align targets within each document, and exclude padding or invalid boundaries from the loss. Test a change to the other document while holding the target document fixed." }, "pitfall": { "symptom": "Padding produces NaNs in attention.", "check": "Count permitted keys in every query row. A fully masked softmax needs an explicit invalid-query policy; its output and loss must not become valid training evidence." }, "completion": "Use text-harness stages 1–4. Keep independent attention arithmetic, future-token invariance, resumed Adam state and full-prefix/cache parity with declared precision.", "furtherWork": "Real-corpus curation, tokenizer tradeoff studies, long-context evaluation and live generative quality need further labs beyond the tiny byte-model harness." }, "additionalProjectIds": [ "embedding-lineage" ], "lessons": [ { "id": "transformers-01", "title": "Attention from small pieces", "status": "authored", "objective": "Implement scaled dot-product attention and test its output shape.", "prerequisites": [], "exercise": "Permute the three key/value pairs together, then permute only values. Verify invariance for the paired permutation and a changed result for the mismatched pairing.", "evidence": "Save the Q/K/V axis diagram, uniform hand calculation, NumPy nonuniform reference, paired permutation check, and wrong-axis repair.", "check": "Which dimension should softmax normalize for one query to combine all key/value pairs?", "path": "phases/07-transformers/01-attention-from-small-pieces", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/07-transformers/01-attention-from-small-pieces/docs/en.md", "script": "exercises/transformers-01.py", "notebook": "exercises/transformers-01.ipynb", "markdown": "exercises/transformers-01.md", "scriptSource": "phases/07-transformers/01-attention-from-small-pieces/code/main.py", "notebookSource": "phases/07-transformers/01-attention-from-small-pieces/notebooks/exercise.ipynb", "quizSource": "phases/07-transformers/01-attention-from-small-pieces/quiz.json", "evidenceSource": "phases/07-transformers/01-attention-from-small-pieces/outputs/evidence.md", "evidence": "exercises/transformers-01-evidence.md" }, "minutes": 90, "contentHash": "a7197251c1f17b230004129a7936f5275c610353013941c3ed2c8d8d3a0b2346", "visualArtifact": { "image": "figures/transformers-01.svg", "png": "figures/transformers-01.png", "data": "figures/transformers-01.json", "generatedAt": "2026-10-20T10:45:51.311855+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "a7197251c1f17b230004129a7936f5275c610353013941c3ed2c8d8d3a0b2346", "kind": "executed" }, "execution": { "contentHash": "a7197251c1f17b230004129a7936f5275c610353013941c3ed2c8d8d3a0b2346", "executedAt": "2026-10-20T10:47:04.943776+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Weights / output: [[0.33333334 0.33333334 0.33333334]\n [0.33333334 0.33333334 0.33333334]] [[2. 2.]\n [2. 2.]]\nWeights / output: [[0.33333334 0.33333334 0.33333334]\n [0.33333334 0.33333334 0.33333334]] [[2. 2.]\n [2. 2.]]\nPASS: transformers-01\n", "url": "executions/transformers-01.json" }, "url": "lesson.html?lesson=transformers-01" }, { "id": "transformers-02", "title": "Masking, tokenization, and sequence packing", "status": "authored", "objective": "Construct a causal mask and prevent packed sequences from attending across boundaries.", "prerequisites": [ "transformers-01" ], "exercise": "Encode abcab with a three-character vocabulary. Create shifted inputs/targets, and a loss-eligibility vector for packed segments [0,0,0,1,1] that excludes a cross-segment next-token target.", "evidence": "Keep token/target alignment, the two masks, future perturbation, segment-leakage comparison and all-masked-row repair. Verify that masks and ignored loss positions enforce the same boundaries.", "check": "What must a packed causal mask enforce?", "path": "phases/07-transformers/02-masking-tokenization-and-sequence-packing", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/07-transformers/02-masking-tokenization-and-sequence-packing/docs/en.md", "script": "exercises/transformers-02.py", "notebook": "exercises/transformers-02.ipynb", "markdown": "exercises/transformers-02.md", "scriptSource": "phases/07-transformers/02-masking-tokenization-and-sequence-packing/code/main.py", "notebookSource": "phases/07-transformers/02-masking-tokenization-and-sequence-packing/notebooks/exercise.ipynb", "quizSource": "phases/07-transformers/02-masking-tokenization-and-sequence-packing/quiz.json", "evidenceSource": "phases/07-transformers/02-masking-tokenization-and-sequence-packing/outputs/evidence.md", "evidence": "exercises/transformers-02-evidence.md" }, "minutes": 90, "contentHash": "d62730a9b6f73ecf78a5b1ea0d799061469f8f60316f7715e6cd70ed70a0f606", "visualArtifact": { "image": "figures/transformers-02.svg", "png": "figures/transformers-02.png", "data": "figures/transformers-02.json", "generatedAt": "2026-10-20T10:45:52.168439+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "d62730a9b6f73ecf78a5b1ea0d799061469f8f60316f7715e6cd70ed70a0f606", "kind": "executed" }, "execution": { "contentHash": "d62730a9b6f73ecf78a5b1ea0d799061469f8f60316f7715e6cd70ed70a0f606", "executedAt": "2026-10-20T10:47:08.194358+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Causal means: [2. 3. 4.666667 7.5 ]\nCausal means: [2. 3. 4.666667 7.5 ]\nPASS: transformers-02\n", "url": "executions/transformers-02.json" }, "url": "lesson.html?lesson=transformers-02" }, { "id": "transformers-03", "title": "A Transformer block and mixed precision", "status": "authored", "objective": "Build a residual attention block and compare numerical behavior across precision choices.", "prerequisites": [ "transformers-02" ], "exercise": "Differentiate the mean squared block output with respect to every projection. Check each gradient leaf’s shape and finite values, then perturb one parameter coordinate and compare a central difference.", "evidence": "Keep the architecture diagram, zero-projection identity, causal perturbation, independent normalization, finite-difference coordinate check, precision discrepancy and residual-width repair.", "check": "If all projection matrices are zero in this block, what should the output be?", "path": "phases/07-transformers/03-a-transformer-block-and-mixed-precision", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/07-transformers/03-a-transformer-block-and-mixed-precision/docs/en.md", "script": "exercises/transformers-03.py", "notebook": "exercises/transformers-03.ipynb", "markdown": "exercises/transformers-03.md", "scriptSource": "phases/07-transformers/03-a-transformer-block-and-mixed-precision/code/main.py", "notebookSource": "phases/07-transformers/03-a-transformer-block-and-mixed-precision/notebooks/exercise.ipynb", "quizSource": "phases/07-transformers/03-a-transformer-block-and-mixed-precision/quiz.json", "evidenceSource": "phases/07-transformers/03-a-transformer-block-and-mixed-precision/outputs/evidence.md", "evidence": "exercises/transformers-03-evidence.md" }, "minutes": 90, "contentHash": "e1476923ec8af45541c1ffde358df2ce7e0100dd828c265a75ad1db8f790568d", "visualArtifact": { "image": "figures/transformers-03.svg", "png": "figures/transformers-03.png", "data": "figures/transformers-03.json", "generatedAt": "2026-10-20T10:45:54.021705+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "e1476923ec8af45541c1ffde358df2ce7e0100dd828c265a75ad1db8f790568d", "kind": "executed", "mechanismImage": "figures/transformers-03-mechanism.svg", "mechanismPng": "figures/transformers-03-mechanism.png" }, "execution": { "contentHash": "e1476923ec8af45541c1ffde358df2ce7e0100dd828c265a75ad1db8f790568d", "executedAt": "2026-10-20T10:47:18.530225+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Block shape: (4, 8)\nBlock shape: (4, 8)\nbfloat16 max absolute error: 0.013455629348754883\nFloat32 reduction policy max absolute error: 0.013455629348754883\nPASS: transformers-03\n", "url": "executions/transformers-03.json" }, "url": "lesson.html?lesson=transformers-03" }, { "id": "transformers-04", "title": "Train, checkpoint, and generate", "status": "authored", "objective": "Train a tiny causal model, restore it, and generate from a fixed seed.", "prerequisites": [ "transformers-03" ], "exercise": "Generate twelve new tokens from prompt bcab. Derive the expected cycle first; verify the string. Explain why success on this periodic example is separate from natural-language performance.", "evidence": "Save the target alignment, uniform-loss derivation, training/held-out metrics with corpus limitations, snapshot continuation check, generated cycle and metadata/prompt rejection repairs.", "check": "Why is saving only parameters insufficient to replay the next Adam update?", "path": "phases/07-transformers/04-train-checkpoint-and-generate", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/07-transformers/04-train-checkpoint-and-generate/docs/en.md", "script": "exercises/transformers-04.py", "notebook": "exercises/transformers-04.ipynb", "markdown": "exercises/transformers-04.md", "scriptSource": "phases/07-transformers/04-train-checkpoint-and-generate/code/main.py", "notebookSource": "phases/07-transformers/04-train-checkpoint-and-generate/notebooks/exercise.ipynb", "quizSource": "phases/07-transformers/04-train-checkpoint-and-generate/quiz.json", "evidenceSource": "phases/07-transformers/04-train-checkpoint-and-generate/outputs/evidence.md", "evidence": "exercises/transformers-04-evidence.md" }, "minutes": 120, "contentHash": "a3d44f0f5adadb3a73c2b36d11aac31b162247575dca4562ca7fbc7ae0db0211", "visualArtifact": { "image": "figures/transformers-04.svg", "png": "figures/transformers-04.png", "data": "figures/transformers-04.json", "generatedAt": "2026-10-20T10:45:58.053923+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "a3d44f0f5adadb3a73c2b36d11aac31b162247575dca4562ca7fbc7ae0db0211", "kind": "executed" }, "execution": { "contentHash": "a3d44f0f5adadb3a73c2b36d11aac31b162247575dca4562ca7fbc7ae0db0211", "executedAt": "2026-10-20T10:47:31.761160+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Initial / held loss / held accuracy: 0.9807573556900024 0.0017881905660033226 1.0\nGreedy continuation: abcabcabcabc\nInitial / held loss / held accuracy: 0.9807573556900024 0.0017881905660033226 1.0\nGreedy continuation: abcabcabcabc\nPASS: transformers-04\n", "url": "executions/transformers-04.json" }, "url": "lesson.html?lesson=transformers-04" } ] }, { "id": "performance", "number": "08", "title": "Performance diagnosis", "part": "Training systems", "prerequisitePhaseIds": [ "networks" ], "hardware": "CPU for timing concepts; TPU or GPU for profiling", "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.", "source": "https://openxla.org/xprof", "description": "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.", "learningAdvice": "Keep workload shapes, warm-up policy and synchronized timings with your device report.", "projectId": "performance-audit", "projectStageIds": [ "1", "2", "3", "4" ], "projectNote": "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.", "studyGuide": { "question": "What measured evidence identifies the bottleneck?", "why": "Optimize a verified computation and a named timing boundary. Trace events, compiler estimates and wall-clock samples answer different questions; connect them before proposing a change.", "readiness": { "prompt": "A JAX call returns an array handle almost immediately. Does a timer ending there necessarily measure completed computation?", "answer": "No. Synchronize the result inside the measured boundary. Report first-call setup separately from warmed execution and keep inputs, dtype and output checks fixed.", "lessonId": "performance-01" }, "milestones": [ { "lessonIds": [ "performance-01", "performance-02" ], "title": "Choose a comparable measurement", "task": "Record warmup, synchronization and sample distributions. Change shapes and host reads separately so you can identify their effects." }, { "lessonIds": [ "performance-03", "performance-04" ], "title": "Explain the evidence", "task": "Locate one annotated operation in a real trace, distinguish nested from additive time and state the memory level behind a roofline estimate." }, { "lessonIds": [ "performance-05" ], "title": "Test a tradeoff", "task": "Compare rematerialized and ordinary gradients before timing. Distinguish saved residuals, compiler estimates and observed peak memory." } ], "transfer": { "prompt": "A kernel becomes twice as fast but the full step barely changes. What does the result suggest?", "approach": "First measure the kernel’s original share of the step. Even eliminating a small component has limited total benefit. Also inspect copies, synchronization and changed compilation; a local kernel timer cannot explain the complete pipeline alone." }, "pitfall": { "symptom": "A claimed speedup disappears in repeated runs.", "check": "Check equivalent work, compilation reuse, synchronization, sample spread and competing system activity; retain the raw samples." }, "completion": "Use the trace and timing stage of sharded-training. Keep before/after correctness, named measurement boundaries, trace locations and distributions with a justified accept/reject decision.", "furtherWork": "CPU traces do not characterize TPU/GPU communication or memory. Target traces and sustained workload measurements remain required for accelerator claims." }, "additionalProjectIds": [ "sharded-training" ], "lessons": [ { "id": "performance-01", "title": "Benchmark asynchronous work correctly", "status": "authored", "objective": "Warm up a computation and synchronize results before reporting timing.", "prerequisites": [], "exercise": "Repeat the synchronized benchmark with 96 rows, then record an independently checked row-energy workload.", "evidence": "Keep benchmark.py, complete observed JSON/sample output, workload shapes/dtypes, numerical tolerance/error, first-call versus warmed boundaries, changed-workload report and repaired timer explanation.", "check": "Distinguish call-plus-wait wall time from compiler/kernel time; reproduce a new workload with an independent reference.", "path": "phases/08-performance/01-benchmark-asynchronous-work-correctly", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/08-performance/01-benchmark-asynchronous-work-correctly/docs/en.md", "script": "exercises/performance-01.py", "notebook": "exercises/performance-01.ipynb", "markdown": "exercises/performance-01.md", "scriptSource": "phases/08-performance/01-benchmark-asynchronous-work-correctly/code/main.py", "notebookSource": "phases/08-performance/01-benchmark-asynchronous-work-correctly/notebooks/exercise.ipynb", "quizSource": "phases/08-performance/01-benchmark-asynchronous-work-correctly/quiz.json", "evidenceSource": "phases/08-performance/01-benchmark-asynchronous-work-correctly/outputs/evidence.md", "evidence": "exercises/performance-01-evidence.md" }, "minutes": 80, "contentHash": "4b0835048c71e4db833b4b8988c5210075885e908194780b4db2d0cad29c33ed", "visualArtifact": { "image": "figures/performance-01.svg", "png": "figures/performance-01.png", "data": "figures/performance-01.json", "generatedAt": "2026-10-20T10:45:58.968076+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "4b0835048c71e4db833b4b8988c5210075885e908194780b4db2d0cad29c33ed", "kind": "executed", "mechanismImage": "figures/performance-01-mechanism.svg", "mechanismPng": "figures/performance-01-mechanism.png" }, "execution": { "contentHash": "4b0835048c71e4db833b4b8988c5210075885e908194780b4db2d0cad29c33ed", "executedAt": "2026-10-20T10:47:33.645730+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "{\n \"backend\": \"TFRT_CPU_0\",\n \"jax\": \"0.9.2\",\n \"numpy\": \"2.4.4\",\n \"python\": \"3.14.3\",\n \"input_shape\": [\n 64,\n 32\n ],\n \"weight_shape\": [\n 32,\n 16\n ],\n \"dtype\": \"float32\",\n \"boundary\": \"input already placed; call plus output synchronization\",\n \"first_call_seconds\": 0.051288416143506765,\n \"warm_samples_seconds\": [\n 5.9833284467458725e-05,\n 2.5249551981687546e-05,\n 1.200009137392044e-05,\n 1.429067924618721e-05,\n 1.00000761449337e-05,\n 8.584000170230865e-06,\n 8.875038474798203e-06\n ],\n \"warm_median_seconds\": 1.200009137392044e-05,\n \"max_abs_error\": 2.384185791015625e-07\n}\n{\n \"backend\": \"TFRT_CPU_0\",\n \"jax\": \"0.9.2\",\n \"numpy\": \"2.4.4\",\n \"python\": \"3.14.3\",\n \"input_shape\": [\n 64,\n 32\n ],\n \"weight_shape\": [\n 32,\n 16\n ],\n \"dtype\": \"float32\",\n \"boundary\": \"input already placed; call plus output synchronization\",\n \"first_call_seconds\": 0.02849070867523551,\n \"warm_samples_seconds\": [\n 0.00011466676369309425,\n 0.0001077083870768547,\n 4.1958410292863846e-05,\n 3.5458244383335114e-05,\n 1.2999866157770157e-05,\n 1.1667143553495407e-05,\n 1.3041310012340546e-05\n ],\n \"warm_median_seconds\": 3.5458244383335114e-05,\n \"max_abs_error\": 2.384185791015625e-07\n}\nHandle-only seconds: 1.1832918971776962e-05\nComplete-call samples: [1.8625054508447647e-05, 1.2416858226060867e-05, 1.2083910405635834e-05]\nStructured samples: [6.495881825685501e-05, 3.074994310736656e-05, 3.1624920666217804e-05]\nChanged workload: (96, 32) [6.645871326327324e-05, 5.883257836103439e-05, 5.4459087550640106e-05, 2.1249987185001373e-05, 1.1249911040067673e-05]\nEnergy shape and samples: (64,) [5.758274346590042e-05, 4.183407872915268e-05, 2.150004729628563e-05, 1.8625054508447647e-05, 1.8666032701730728e-05]\n{'boundary': 'placed inputs; warm call plus output wait', 'samples': [4.9999915063381195e-05, 5.924981087446213e-05, 5.816575139760971e-05, 5.9042125940322876e-05, 6.141699850559235e-05]}\nPASS: performance-01\n", "url": "executions/performance-01.json" }, "url": "lesson.html?lesson=performance-01" }, { "id": "performance-02", "title": "Diagnose recompilation and host synchronization", "status": "authored", "objective": "Find a shape-changing call or host read and connect it to observed overhead.", "prerequisites": [ "performance-01" ], "exercise": "Extend the signature ledger to sixteen values and diagnose a static branch, masked final batch and runtime Python branch.", "evidence": "Keep the controlled call ledger, actual trace observations labeled as traces, static/dynamic explanation, observed host-consumption intervals, masked-versus-unmasked counterexample and reproduced/repaired tracing failure.", "check": "Explain which argument changes affect specialization; justify masked fixed shapes and host reporting boundaries with numerical evidence.", "path": "phases/08-performance/02-diagnose-recompilation-and-host-synchronization", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/08-performance/02-diagnose-recompilation-and-host-synchronization/docs/en.md", "script": "exercises/performance-02.py", "notebook": "exercises/performance-02.ipynb", "markdown": "exercises/performance-02.md", "scriptSource": "phases/08-performance/02-diagnose-recompilation-and-host-synchronization/code/main.py", "notebookSource": "phases/08-performance/02-diagnose-recompilation-and-host-synchronization/notebooks/exercise.ipynb", "quizSource": "phases/08-performance/02-diagnose-recompilation-and-host-synchronization/quiz.json", "evidenceSource": "phases/08-performance/02-diagnose-recompilation-and-host-synchronization/outputs/evidence.md", "evidence": "exercises/performance-02-evidence.md" }, "minutes": 80, "contentHash": "fc7b7357392cc925f95ee951dd96a30ce905c4b45496c33435a6bcc337bbc07a", "visualArtifact": { "image": "figures/performance-02.svg", "png": "figures/performance-02.png", "data": "figures/performance-02.json", "generatedAt": "2026-10-20T10:45:59.629728+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "fc7b7357392cc925f95ee951dd96a30ce905c4b45496c33435a6bcc337bbc07a", "kind": "executed" }, "execution": { "contentHash": "fc7b7357392cc925f95ee951dd96a30ce905c4b45496c33435a6bcc337bbc07a", "executedAt": "2026-10-20T10:47:35.719968+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "{\n \"jax\": \"0.9.2\",\n \"backend\": \"TFRT_CPU_0\",\n \"calls\": [\n {\n \"label\": \"initial signature\",\n \"shape\": [\n 8\n ],\n \"dtype\": \"float32\",\n \"observed_trace_delta\": 1,\n \"result\": 140.0\n },\n {\n \"label\": \"same signature changed data\",\n \"shape\": [\n 8\n ],\n \"dtype\": \"float32\",\n \"observed_trace_delta\": 0,\n \"result\": 204.0\n },\n {\n \"label\": \"same signature changed dynamic gain\",\n \"shape\": [\n 8\n ],\n \"dtype\": \"float32\",\n \"observed_trace_delta\": 0,\n \"result\": 560.0\n },\n {\n \"label\": \"different input shape\",\n \"shape\": [\n 12\n ],\n \"dtype\": \"float32\",\n \"observed_trace_delta\": 1,\n \"result\": 506.0\n }\n ],\n \"python_trace_events\": [\n {\n \"shape\": [\n 8\n ],\n \"dtype\": \"float32\"\n },\n {\n \"shape\": [\n 12\n ],\n \"dtype\": \"float32\"\n }\n ]\n}\n{\n \"jax\": \"0.9.2\",\n \"backend\": \"TFRT_CPU_0\",\n \"calls\": [\n {\n \"label\": \"initial signature\",\n \"shape\": [\n 8\n ],\n \"dtype\": \"float32\",\n \"observed_trace_delta\": 1,\n \"result\": 140.0\n },\n {\n \"label\": \"same signature changed data\",\n \"shape\": [\n 8\n ],\n \"dtype\": \"float32\",\n \"observed_trace_delta\": 0,\n \"result\": 204.0\n },\n {\n \"label\": \"same signature changed dynamic gain\",\n \"shape\": [\n 8\n ],\n \"dtype\": \"float32\",\n \"observed_trace_delta\": 0,\n \"result\": 560.0\n },\n {\n \"label\": \"different input shape\",\n \"shape\": [\n 12\n ],\n \"dtype\": \"float32\",\n \"observed_trace_delta\": 1,\n \"result\": 506.0\n }\n ],\n \"python_trace_events\": [\n {\n \"shape\": [\n 8\n ],\n \"dtype\": \"float32\"\n },\n {\n \"shape\": [\n 12\n ],\n \"dtype\": \"float32\"\n }\n ]\n}\nStatic branch trace observations: [{'shape': [8], 'square': True}, {'shape': [8], 'square': False}]\nPer-step host reads seconds: 0.00012137508019804955\nFinal-only wait seconds: 3.5999808460474014e-05\nAdded ledger row: {'label': 'new length sixteen', 'shape': [16], 'dtype': 'float32', 'observed_trace_delta': 1, 'result': 1240.0}\nMasked versus unmasked padded objective: 55.0 58.0\nObserved runtime-Python-branch tracing failure\nPASS: performance-02\n", "url": "executions/performance-02.json" }, "url": "lesson.html?lesson=performance-02" }, { "id": "performance-03", "title": "Read an XProf trace", "status": "authored", "objective": "Capture an actual warmed CPU trace with named steps. Identify nested intervals and avoid double-counting their durations.", "prerequisites": [ "performance-02" ], "exercise": "Convert every model interval to milliseconds and summarize its minimum, median and maximum. Keep the individual samples and explain why four observations is distinct from a tail-latency service guarantee.", "evidence": "Keep the actual XPlane and Perfetto paths, named interval durations and nesting check, units and an interpretation separating injected input delay from completed computation.", "check": "A step annotation lasts longer than its model annotation. What can this alone establish?", "path": "phases/08-performance/03-read-an-xprof-trace", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/08-performance/03-read-an-xprof-trace/docs/en.md", "script": "exercises/performance-03.py", "notebook": "exercises/performance-03.ipynb", "markdown": "exercises/performance-03.md", "scriptSource": "phases/08-performance/03-read-an-xprof-trace/code/main.py", "notebookSource": "phases/08-performance/03-read-an-xprof-trace/notebooks/exercise.ipynb", "quizSource": "phases/08-performance/03-read-an-xprof-trace/quiz.json", "evidenceSource": "phases/08-performance/03-read-an-xprof-trace/outputs/evidence.md", "evidence": "exercises/performance-03-evidence.md" }, "minutes": 100, "contentHash": "8035a6f01d72cddc68632f648e4e09b96b1e34f69f29a1d0bd99144f3f418f5f", "visualArtifact": { "image": "figures/performance-03.svg", "png": "figures/performance-03.png", "data": "figures/performance-03.json", "generatedAt": "2026-10-20T10:46:00.292509+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "8035a6f01d72cddc68632f648e4e09b96b1e34f69f29a1d0bd99144f3f418f5f", "kind": "executed" }, "execution": { "contentHash": "8035a6f01d72cddc68632f648e4e09b96b1e34f69f29a1d0bd99144f3f418f5f", "executedAt": "2026-10-20T10:47:37.536684+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Trace directory: /private/var/folders/mp/_mq7srqx2y10p5nhjz9m7hj40000gn/T/jax-course-trace-l3tdro07\n{\n \"input_wait_us\": [\n 1263,\n 1253,\n 1269,\n 1271\n ],\n \"model_and_wait_us\": [\n 108,\n 137,\n 74,\n 140\n ],\n \"step_us\": [\n 1393,\n 1406,\n 1354,\n 1423\n ]\n}\nFour named steps verified; host annotations are not isolated device kernels.\nTrace directory: /private/var/folders/mp/_mq7srqx2y10p5nhjz9m7hj40000gn/T/jax-course-trace-g8x32sim\n{\n \"input_wait_us\": [\n 1263,\n 1279,\n 1265,\n 1275\n ],\n \"model_and_wait_us\": [\n 82,\n 136,\n 62,\n 122\n ],\n \"step_us\": [\n 1350,\n 1429,\n 1335,\n 1407\n ]\n}\nFour named steps verified; host annotations are not isolated device kernels.\nUnattributed outer time, microseconds: [5, 14, 8, 10]\nWrong double-counted totals: [2695, 2844, 2662, 2804]\n{'minimum_ms': 0.062, 'median_ms': 0.10200000000000001, 'maximum_ms': 0.136}\nHypothetical union: 12 microseconds; summed duration: 19.\nPASS: performance-03\n", "url": "executions/performance-03.json" }, "url": "lesson.html?lesson=performance-03" }, { "id": "performance-04", "title": "Reason with HLO and the roofline model", "status": "authored", "objective": "Locate matrix multiplication and activation in StableHLO. Derive the operation and minimum byte counts for a declared matrix contract.", "prerequisites": [ "performance-03" ], "exercise": "Derive the work and minimum traffic for shapes \\((48,24)\\) and \\((24,12)\\). Compare your hand calculation with the contract function and report the hypothetical limiting budget.", "evidence": "Keep StableHLO, synchronized CPU samples and independent output parity. Derive the byte/FLOP model and label the roofline peak and bandwidth as hypothetical.", "check": "The lowered program contains a dot operation and a roofline model predicts a high ceiling. What is still needed before claiming high accelerator throughput?", "path": "phases/08-performance/04-reason-with-hlo-and-the-roofline-model", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/08-performance/04-reason-with-hlo-and-the-roofline-model/docs/en.md", "script": "exercises/performance-04.py", "notebook": "exercises/performance-04.ipynb", "markdown": "exercises/performance-04.md", "scriptSource": "phases/08-performance/04-reason-with-hlo-and-the-roofline-model/code/main.py", "notebookSource": "phases/08-performance/04-reason-with-hlo-and-the-roofline-model/notebooks/exercise.ipynb", "quizSource": "phases/08-performance/04-reason-with-hlo-and-the-roofline-model/quiz.json", "evidenceSource": "phases/08-performance/04-reason-with-hlo-and-the-roofline-model/outputs/evidence.md", "evidence": "exercises/performance-04-evidence.md" }, "minutes": 95, "contentHash": "24ad8bac2a5da5324fe40b66a37dbc8ed9430798cde46b66845ad3981944cf31", "visualArtifact": { "image": "figures/performance-04.svg", "png": "figures/performance-04.png", "data": "figures/performance-04.json", "generatedAt": "2026-10-20T10:46:01.117672+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "24ad8bac2a5da5324fe40b66a37dbc8ed9430798cde46b66845ad3981944cf31", "kind": "executed" }, "execution": { "contentHash": "24ad8bac2a5da5324fe40b66a37dbc8ed9430798cde46b66845ad3981944cf31", "executedAt": "2026-10-20T10:47:39.421946+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "StableHLO operations present: dot_general and tanh\nCompiler estimates (backend-specific): {'flops': 8192.0, 'bytes accessed1{}': 512.0, 'bytes accessedout{}': 2048.0, 'transcendentals': 256.0, 'bytes accessed': 5632.0, 'utilization1{}': 1.0, 'utilization0{}': 2.0, 'bytes accessed0{}': 3072.0}\nLocal CPU samples seconds: [2.1084211766719818e-05, 7.845880463719368e-05, 5.645817145705223e-05, 4.0166545659303665e-05, 3.654090687632561e-05]\nMatmul-only accounting: {'flops': 8192, 'minimum_bytes': 3584, 'flops_per_byte': 2.2857142857142856}\nHypothetical ceilings GFLOP/s: [26.666666666666664, 53.33333333333333, 100.0, 100.0, 100.0]\nStableHLO operations present: dot_general and tanh\nCompiler estimates (backend-specific): {'utilization1{}': 1.0, 'bytes accessed0{}': 3072.0, 'utilization0{}': 2.0, 'transcendentals': 256.0, 'bytes accessedout{}': 2048.0, 'bytes accessed1{}': 512.0, 'flops': 8192.0, 'bytes accessed': 5632.0}\nLocal CPU samples seconds: [2.9209069907665253e-05, 2.4625100195407867e-05, 1.6333069652318954e-05, 1.2042000889778137e-05, 1.1958181858062744e-05]\nMatmul-only accounting: {'flops': 8192, 'minimum_bytes': 3584, 'flops_per_byte': 2.2857142857142856}\nHypothetical ceilings GFLOP/s: [26.666666666666664, 53.33333333333333, 100.0, 100.0, 100.0]\nCeilings with twice minimum traffic GFLOP/s: [13.333333333333332, 26.666666666666664, 53.33333333333333, 100.0, 100.0]\nChanged contract: 27648 8064 3.4285714285714284 bandwidth-limited in the hypothetical model\nArithmetic unchanged; modeled traffic halved and intensity doubled.\nPASS: performance-04\n", "url": "executions/performance-04.json" }, "url": "lesson.html?lesson=performance-04" }, { "id": "performance-05", "title": "Memory, rematerialization, and optimization tradeoffs", "status": "authored", "objective": "Explain why reverse-mode differentiation needs forward intermediates. Compare ordinary and rematerialized gradients with an independent reverse recurrence.", "prerequisites": [ "performance-04" ], "exercise": "Use the direction \\(v=(1,\\ldots,1)\\) and central differences to independently check a directional derivative of the original objective at the existing input. Compare it with the gradient dot direction and report the perturbation.", "evidence": "Keep the plain/rematerialized gradient and HVP comparisons, saved-residual counts, compiler-memory observations and synchronized timings. Explain why fewer residuals do not guarantee lower measured memory.", "check": "Rematerialization reduces the number of saved residual descriptions but the compiled temporary-byte estimates are equal. What is the supported conclusion?", "path": "phases/08-performance/05-memory-rematerialization-and-optimization-tradeoffs", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/08-performance/05-memory-rematerialization-and-optimization-tradeoffs/docs/en.md", "script": "exercises/performance-05.py", "notebook": "exercises/performance-05.ipynb", "markdown": "exercises/performance-05.md", "scriptSource": "phases/08-performance/05-memory-rematerialization-and-optimization-tradeoffs/code/main.py", "notebookSource": "phases/08-performance/05-memory-rematerialization-and-optimization-tradeoffs/notebooks/exercise.ipynb", "quizSource": "phases/08-performance/05-memory-rematerialization-and-optimization-tradeoffs/quiz.json", "evidenceSource": "phases/08-performance/05-memory-rematerialization-and-optimization-tradeoffs/outputs/evidence.md", "evidence": "exercises/performance-05-evidence.md" }, "minutes": 100, "contentHash": "0ed4e1e9692f1cf9c00068841c968e0f59192c11c421262b792093399bfce201", "visualArtifact": { "image": "figures/performance-05.svg", "png": "figures/performance-05.png", "data": "figures/performance-05.json", "generatedAt": "2026-10-20T10:46:02.014854+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "0ed4e1e9692f1cf9c00068841c968e0f59192c11c421262b792093399bfce201", "kind": "executed" }, "execution": { "contentHash": "0ed4e1e9692f1cf9c00068841c968e0f59192c11c421262b792093399bfce201", "executedAt": "2026-10-20T10:47:42.865487+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Saved residual descriptions:\nplain \nf32[16,16] from a constant\nf32[16] output of tanh from :17:20 (layer)\nf32[16] output of sub from :17:20 (layer)\nf32[16] output of tanh from :17:20 (layer)\nf32[16] output of sub from :17:20 (layer)\nf32[16] output of tanh from :17:20 (layer)\nf32[16] output of sub from :17:20 (layer)\nf32[16] output of tanh from :17:20 (layer)\nf32[16] output of sub from :17:20 (layer)\nremat \nf32[16,16] from a constant\nf32[16] from the argument v\nf32[16] output of tanh from :17:20 (layer)\nf32[16] output of tanh from :17:20 (layer)\nf32[16] output of tanh from :17:20 (layer)\nf32[16] output of tanh from :17:20 (layer)\nCompiler memory estimates: [{'argument_size_in_bytes': 64, 'output_size_in_bytes': 64, 'temp_size_in_bytes': 256, 'alias_size_in_bytes': 0}, {'argument_size_in_bytes': 64, 'output_size_in_bytes': 64, 'temp_size_in_bytes': 256, 'alias_size_in_bytes': 0}]\nSynchronized gradient samples: [[0.0001016668975353241, 8.600018918514252e-05, 0.00013170810416340828, 5.487492308020592e-05, 3.0375085771083832e-05], [3.7041958421468735e-05, 2.6874709874391556e-05, 1.2292061001062393e-05, 9.582843631505966e-06, 9.292270988225937e-06]]\nIndependent reverse recurrence agrees with both gradients.\nplain \nf32[16,16] from a constant\nf32[16] output of tanh from :17:20 (layer)\nf32[16] output of sub from :17:20 (layer)\nf32[16] output of tanh from :17:20 (layer)\nf32[16] output of sub from :17:20 (layer)\nf32[16] output of tanh from :17:20 (layer)\nf32[16] output of sub from :17:20 (layer)\nf32[16] output of tanh from :17:20 (layer)\nf32[16] output of sub from :17:20 (layer)\nremat \nf32[16,16] from a constant\nf32[16] from the argument v\nf32[16] output of tanh from :17:20 (layer)\nf32[16] output of tanh from :17:20 (layer)\nf32[16] output of tanh from :17:20 (layer)\nf32[16] output of tanh from :17:20 (layer)\nCompiler memory estimates: [{'argument_size_in_bytes': 64, 'output_size_in_bytes': 64, 'temp_size_in_bytes': 256, 'alias_size_in_bytes': 0}, {'argument_size_in_bytes': 64, 'output_size_in_bytes': 64, 'temp_size_in_bytes': 256, 'alias_size_in_bytes': 0}]\nSynchronized gradient samples: [[0.0001016668975353241, 8.600018918514252e-05, 0.00013170810416340828, 5.487492308020592e-05, 3.0375085771083832e-05], [3.7041958421468735e-05, 2.6874709874391556e-05, 1.2292061001062393e-05, 9.582843631505966e-06, 9.292270988225937e-06]]\nSaved residual descriptions:\nplain \nf32[16,16] from a constant\nf32[16] output of tanh from :89:20 (layer)\nf32[16] output of sub from :89:20 (layer)\nf32[16] output of tanh from :89:20 (layer)\nf32[16] output of sub from :89:20 (layer)\nf32[16] output of tanh from :89:20 (layer)\nf32[16] output of sub from :89:20 (layer)\nf32[16] output of tanh from :89:20 (layer)\nf32[16] output of sub from :89:20 (layer)\nremat \nf32[16,16] from a constant\nf32[16] from the argument v\nf32[16] output of tanh from :89:20 (layer)\nf32[16] output of tanh from :89:20 (layer)\nf32[16] output of tanh from :89:20 (layer)\nf32[16] output of tanh from :89:20 (layer)\nCompiler memory estimates: [{'argument_size_in_bytes': 64, 'output_size_in_bytes': 64, 'temp_size_in_bytes': 256, 'alias_size_in_bytes': 0}, {'argument_size_in_bytes': 64, 'output_size_in_bytes': 64, 'temp_size_in_bytes': 256, 'alias_size_in_bytes': 0}]\nSynchronized gradient samples: [[3.495905548334122e-05, 2.570822834968567e-05, 1.4999881386756897e-05, 1.4166813343763351e-05, 1.6415957361459732e-05], [2.570776268839836e-05, 2.408307045698166e-05, 8.791731670498848e-05, 7.383385673165321e-05, 5.5417418479919434e-05]]\nIndependent reverse recurrence agrees with both gradients.\nWhole-objective residuals: ['f32[16,16] from a constant', 'f32[16] from the argument v']\nWhole-objective memory estimate: CompiledMemoryStats(generated_code_size_in_bytes=0, argument_size_in_bytes=64, output_size_in_bytes=64, alias_size_in_bytes=0, temp_size_in_bytes=256, host_generated_code_size_in_bytes=0, host_argument_size_in_bytes=0, host_output_size_in_bytes=0, host_alias_size_in_bytes=0, host_temp_size_in_bytes=0)\nDirectional derivative finite/automatic: 0.007404538337141275 0.00740453926846385 epsilon 0.001\nHigher-order directional derivatives agree.\nPASS: performance-05\n", "url": "executions/performance-05.json" }, "url": "lesson.html?lesson=performance-05" } ] }, { "id": "distributed", "number": "09", "title": "Distributed training", "part": "Specializations", "prerequisitePhaseIds": [ "recovery", "performance" ], "hardware": "Logical CPU devices for sharding practice; multi-host work needs a real cluster", "project": "A resilient sharded training experiment", "checkpoint": "Explain partitioning, communication costs, and recovery behavior with measured evidence.", "source": "https://docs.jax.dev/en/latest/", "description": "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.", "learningAdvice": "Compare against a single-device reference; measure real accelerator performance separately.", "projectId": "sharded-training", "projectStageIds": [ "1", "2", "3", "4" ], "projectNote": "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.", "studyGuide": { "question": "Is the distributed update still optimizing the same objective?", "why": "Start with the global mathematical operation, then specify ownership and communication. Logical CPU devices help expose placement and reduction mistakes before you spend accelerator time.", "readiness": { "prompt": "Two devices see one and three valid examples. Should their mean gradients receive equal weight?", "answer": "No. Weight by valid example counts. Sum local gradient numerators and valid counts, then divide once. Padding and empty partitions must not silently change the denominator.", "lessonId": "distributed-02" }, "milestones": [ { "lessonIds": [ "distributed-01", "distributed-02" ], "title": "Verify placement and the global update", "task": "Inspect actual shard indices and compare each relevant output with an independent unpartitioned calculation, including padded and unequal partitions." }, { "lessonIds": [ "distributed-03", "distributed-04" ], "title": "Connect communication and recovery", "task": "Inspect collectives without inferring network speed from payload size. Restore data position and optimizer state before comparing subsequent global steps." } ], "transfer": { "prompt": "A final batch leaves one partition empty. Should it contribute a zero mean, be omitted, or fail?", "approach": "Do not compute an undefined local mean. Accumulate a zero numerator and zero valid count for the empty partition, reduce globally and reject a globally empty batch. Compare the result with the valid unpartitioned examples." }, "pitfall": { "symptom": "The gradient grows with device count.", "check": "Check whether local means were summed without count normalization. Verify one complete update before interpreting a training curve." }, "completion": "Use sharded-training. Retain actual placements and collectives, independent uneven/empty-partition checks and full-state replay. Label logical CPU devices explicitly.", "furtherWork": "Real network collectives, multi-controller ingestion and host-loss recovery are unqualified. The optional launch instructions are preparation, not execution receipts." }, "lessons": [ { "id": "distributed-01", "title": "Arrays, meshes, and sharding", "status": "authored", "objective": "Partition an array over a device mesh and inspect its local shards.", "prerequisites": [], "exercise": "Compare row and column layouts of a (12,8) array, and repair a ten-row mean with padding plus an explicit mask.", "evidence": "Padding repairs the shape but adds artificial rows. The mask keeps them out of both numerator and denominator. Checking only divisibility would miss this statistical error.", "check": "With row sharding, why does a global column sum need results from other devices?", "path": "phases/09-distributed/01-arrays-meshes-and-sharding", "runtime": { "kind": "virtual-cpu", "deviceCount": 4 }, "hardware": "4 logical CPU devices on one host", "artifacts": { "source": "phases/09-distributed/01-arrays-meshes-and-sharding/docs/en.md", "script": "exercises/distributed-01.py", "notebook": "exercises/distributed-01.ipynb", "markdown": "exercises/distributed-01.md", "scriptSource": "phases/09-distributed/01-arrays-meshes-and-sharding/code/main.py", "notebookSource": "phases/09-distributed/01-arrays-meshes-and-sharding/notebooks/exercise.ipynb", "quizSource": "phases/09-distributed/01-arrays-meshes-and-sharding/quiz.json", "evidenceSource": "phases/09-distributed/01-arrays-meshes-and-sharding/outputs/evidence.md", "evidence": "exercises/distributed-01-evidence.md" }, "minutes": 80, "contentHash": "099854f4a238be06a63c9181a954d2abd88a02db6567e62e94e58a2c2852e3c1", "visualArtifact": { "image": "figures/distributed-01.svg", "png": "figures/distributed-01.png", "data": "figures/distributed-01.json", "generatedAt": "2026-10-20T10:46:02.815169+00:00", "environment": { "backend": "cpu", "deviceCount": 4, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "099854f4a238be06a63c9181a954d2abd88a02db6567e62e94e58a2c2852e3c1", "kind": "executed" }, "execution": { "contentHash": "099854f4a238be06a63c9181a954d2abd88a02db6567e62e94e58a2c2852e3c1", "executedAt": "2026-10-20T10:47:44.700275+00:00", "environment": { "backend": "cpu", "deviceCount": 4, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Row-shard shapes: [(2, 4), (2, 4), (2, 4), (2, 4)]\nColumn-shard shapes: [(8, 1), (8, 1), (8, 1), (8, 1)]\nColumn sums: [112. 120. 128. 136.]\nRow-shard shapes: [(2, 4), (2, 4), (2, 4), (2, 4)]\nColumn-shard shapes: [(8, 1), (8, 1), (8, 1), (8, 1)]\nColumn sums: [112. 120. 128. 136.]\nResharded values preserved\nRow sums: [ 6. 22. 38. 54. 70. 86. 102. 118.]\nExpected indivisible batch\nExpected indivisible batch\nPASS: distributed-01\n", "url": "executions/distributed-01.json" }, "url": "lesson.html?lesson=distributed-01" }, { "id": "distributed-02", "title": "A sharded training step", "status": "authored", "objective": "Run a training step with explicit data and parameter placement.", "prerequisites": [ "distributed-01" ], "exercise": "Run a changed target/initial-state update, then diagnose the factor-of-four local-mean bug and repair uneven-batch aggregation.", "evidence": "Each local gradient already divides by its local count. Summing four equal local means omits division by four. For unequal counts, an unweighted mean also misrepresents examples; sum count-weighted local means and divide by total count.", "check": "Four equal-size shards compute local mean gradients. Which aggregation gives the global mean gradient?", "path": "phases/09-distributed/02-a-sharded-training-step", "runtime": { "kind": "virtual-cpu", "deviceCount": 4 }, "hardware": "4 logical CPU devices on one host", "artifacts": { "source": "phases/09-distributed/02-a-sharded-training-step/docs/en.md", "script": "exercises/distributed-02.py", "notebook": "exercises/distributed-02.ipynb", "markdown": "exercises/distributed-02.md", "scriptSource": "phases/09-distributed/02-a-sharded-training-step/code/main.py", "notebookSource": "phases/09-distributed/02-a-sharded-training-step/notebooks/exercise.ipynb", "quizSource": "phases/09-distributed/02-a-sharded-training-step/quiz.json", "evidenceSource": "phases/09-distributed/02-a-sharded-training-step/outputs/evidence.md", "evidence": "exercises/distributed-02-evidence.md" }, "minutes": 90, "contentHash": "1c8eea315a4b010f2d42b9d63799ef7d0db1b4f49a05b7582498d773c890dae2", "visualArtifact": { "image": "figures/distributed-02.svg", "png": "figures/distributed-02.png", "data": "figures/distributed-02.json", "generatedAt": "2026-10-20T10:46:03.549327+00:00", "environment": { "backend": "cpu", "deviceCount": 4, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "1c8eea315a4b010f2d42b9d63799ef7d0db1b4f49a05b7582498d773c890dae2", "kind": "executed" }, "execution": { "contentHash": "1c8eea315a4b010f2d42b9d63799ef7d0db1b4f49a05b7582498d773c890dae2", "executedAt": "2026-10-20T10:47:46.399473+00:00", "environment": { "backend": "cpu", "deviceCount": 4, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Initial loss: 4.875\nGradient: [-5.25 -0.75] updated weights: [0.52500004 0.075 ]\nUpdated loss: 2.6784372329711914\nInitial loss: 4.875\nGradient: [-5.25 -0.75] updated weights: [0.52500004 0.075 ]\nUpdated loss: 2.6784372329711914\nPermutation preserves the global update\nTen-step weights: [ 1.704636 -0.7047408]\nSummed local means are four times too large; weighted aggregation repaired\nSummed local means are four times too large; weighted aggregation repaired\nPASS: distributed-02\n", "url": "executions/distributed-02.json" }, "url": "lesson.html?lesson=distributed-02" }, { "id": "distributed-03", "title": "Communication-efficient algorithms", "status": "authored", "objective": "Compare all-reduce and reduce-scatter against the same independent global gradient. Explain global result shape versus per-device storage.", "prerequisites": [ "distributed-02" ], "exercise": "Repeat the gradient comparison after reversing both observation rows and labels and changing the initial weights. Verify every all-reduce and reconstructed local replica, then explain why reversing only the labels changes the problem.", "evidence": "Keep global NumPy gradient checks, local shard shapes, lowered collectives, changed-value replicas and complete scatter/gather payload accounting. Separate the idealized ring model from observed CPU timings.", "check": "A program replaces all-reduce with reduce-scatter and immediately gathers the full vector again. What does the idealized ring model predict?", "path": "phases/09-distributed/03-communication-efficient-algorithms", "runtime": { "kind": "virtual-cpu", "deviceCount": 4 }, "hardware": "4 logical CPU devices; target cluster execution unverified", "artifacts": { "source": "phases/09-distributed/03-communication-efficient-algorithms/docs/en.md", "script": "exercises/distributed-03.py", "notebook": "exercises/distributed-03.ipynb", "markdown": "exercises/distributed-03.md", "scriptSource": "phases/09-distributed/03-communication-efficient-algorithms/code/main.py", "notebookSource": "phases/09-distributed/03-communication-efficient-algorithms/notebooks/exercise.ipynb", "quizSource": "phases/09-distributed/03-communication-efficient-algorithms/quiz.json", "evidenceSource": "phases/09-distributed/03-communication-efficient-algorithms/outputs/evidence.md", "evidence": "exercises/distributed-03-evidence.md" }, "minutes": 100, "contentHash": "8513af7c9ad7fd790134999a25b276112f6429a15cdbb2fd3f2528075d7372fa", "visualArtifact": { "image": "figures/distributed-03.svg", "png": "figures/distributed-03.png", "data": "figures/distributed-03.json", "generatedAt": "2026-10-20T10:46:04.411433+00:00", "environment": { "backend": "cpu", "deviceCount": 4, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "8513af7c9ad7fd790134999a25b276112f6429a15cdbb2fd3f2528075d7372fa", "kind": "executed", "mechanismImage": "figures/distributed-03-mechanism.svg", "mechanismPng": "figures/distributed-03-mechanism.png" }, "execution": { "contentHash": "8513af7c9ad7fd790134999a25b276112f6429a15cdbb2fd3f2528075d7372fa", "executedAt": "2026-10-20T10:47:48.161184+00:00", "environment": { "backend": "cpu", "deviceCount": 4, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Gradient reference: [0.0029296851716935635, 0.00219726306386292, 0.0014648411888629198, 0.0007324193138629198, -2.6193447411060333e-09, -0.000732424552552402, -0.001464846427552402, -0.002197268418967724]\nLocal result shapes: replicated (8,), reduce-scatter (2,), gathered (8,)\nIdealized ring bytes per rank: 48.0 24.0\nCompleted CPU median microseconds: {'all_reduce': 99.08340871334076, 'reduce_scatter': 102.49996557831764, 'all_gather_only': 108.41712355613708}\nGradient reference: [0.0029296851716935635, 0.00219726306386292, 0.0014648411888629198, 0.0007324193138629198, -2.6193447411060333e-09, -0.000732424552552402, -0.001464846427552402, -0.002197268418967724]\nLocal result shapes: replicated (8,), reduce-scatter (2,), gathered (8,)\nIdealized ring bytes per rank: 48.0 24.0\nCompleted CPU median microseconds: {'all_reduce': 109.4592735171318, 'reduce_scatter': 109.74984616041183, 'all_gather_only': 86.12498641014099}\nSharded momentum matches the independent elementwise update.\nJoint permutation preserves the gradient; label-only reversal does not.\nIdealized payload/storage budget, not observed network traffic: [{'ranks': 2, 'all_reduce_bytes': 1048576, 'scatter_gather_bytes': 1048576, 'shard_storage_bytes': 524288}, {'ranks': 4, 'all_reduce_bytes': 1572864, 'scatter_gather_bytes': 1572864, 'shard_storage_bytes': 262144}, {'ranks': 8, 'all_reduce_bytes': 1835008, 'scatter_gather_bytes': 1835008, 'shard_storage_bytes': 131072}]\nPASS: distributed-03\n", "url": "executions/distributed-03.json" }, "url": "lesson.html?lesson=distributed-03" }, { "id": "distributed-04", "title": "Resilient distributed training", "status": "authored", "objective": "Restore a complete training state with explicit target sharding. Compare next sample IDs, optimizer state and losses across an epoch boundary.", "prerequisites": [ "distributed-03" ], "exercise": "Extend uninterrupted and restored execution by two additional steps. Verify sample IDs, complete state and every local parameter replica. Explain which comparison would fail if a single replica held a different parameter vector.", "evidence": "Keep the checkpoint path, saved/restored full state, subsequent sample IDs and losses through a new epoch. Diagnose missing momentum and a changed shuffle key; state the untested multi-host boundary.", "check": "The first resumed batch and update match, but sample IDs differ at the next epoch. Which saved value should you inspect first?", "path": "phases/09-distributed/04-resilient-distributed-training", "runtime": { "kind": "virtual-cpu", "deviceCount": 4 }, "hardware": "4 logical CPU devices; target cluster execution unverified", "artifacts": { "source": "phases/09-distributed/04-resilient-distributed-training/docs/en.md", "script": "exercises/distributed-04.py", "notebook": "exercises/distributed-04.ipynb", "markdown": "exercises/distributed-04.md", "scriptSource": "phases/09-distributed/04-resilient-distributed-training/code/main.py", "notebookSource": "phases/09-distributed/04-resilient-distributed-training/notebooks/exercise.ipynb", "quizSource": "phases/09-distributed/04-resilient-distributed-training/quiz.json", "evidenceSource": "phases/09-distributed/04-resilient-distributed-training/outputs/evidence.md", "evidence": "exercises/distributed-04-evidence.md" }, "minutes": 105, "contentHash": "fba8433e52316e9b4754f027642448996809914e15f1853638a5018bfa9eaeb8", "visualArtifact": { "image": "figures/distributed-04.svg", "png": "figures/distributed-04.png", "data": "figures/distributed-04.json", "generatedAt": "2026-10-20T10:46:05.852257+00:00", "environment": { "backend": "cpu", "deviceCount": 4, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "fba8433e52316e9b4754f027642448996809914e15f1853638a5018bfa9eaeb8", "kind": "executed" }, "execution": { "contentHash": "fba8433e52316e9b4754f027642448996809914e15f1853638a5018bfa9eaeb8", "executedAt": "2026-10-20T10:47:51.553258+00:00", "environment": { "backend": "cpu", "deviceCount": 4, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Recorded backend/processes/devices: cpu 1 4\nRestored step/cursor: 1 8\nNext sample IDs across epoch boundary: [[12, 5, 4, 1, 8, 13, 2, 11], [9, 0, 7, 15, 11, 13, 8, 14], [1, 6, 5, 4, 3, 12, 2, 10]]\nUninterrupted losses: [0.5069587826728821, 0.19637340307235718, 0.1798383891582489]\nRestored losses: [0.5069587826728821, 0.19637340307235718, 0.1798383891582489]\nCheckpoint: /var/folders/mp/_mq7srqx2y10p5nhjz9m7hj40000gn/T/jax-distributed-recovery-8l8pqueq/step-one\nRecorded backend/processes/devices: cpu 1 4\nRestored step/cursor: 1 8\nNext sample IDs across epoch boundary: [[12, 5, 4, 1, 8, 13, 2, 11], [9, 0, 7, 15, 11, 13, 8, 14], [1, 6, 5, 4, 3, 12, 2, 10]]\nUninterrupted losses: [0.5069587826728821, 0.19637340307235718, 0.1798383891582489]\nRestored losses: [0.5069587826728821, 0.19637340307235718, 0.1798383891582489]\nCheckpoint: /var/folders/mp/_mq7srqx2y10p5nhjz9m7hj40000gn/T/jax-distributed-recovery-9h7rtdxi/step-one\nSame next sample IDs, different weights after missing momentum: 0.0195157527923584\nFive resumed transitions and all addressable parameter replicas agree.\nCurrent-epoch batch matches; next-epoch batch exposes the wrong key.\nPASS: distributed-04\n", "url": "executions/distributed-04.json" }, "url": "lesson.html?lesson=distributed-04" } ] }, { "id": "science", "number": "10", "title": "Scientific computing", "part": "Specializations", "prerequisitePhaseIds": [ "optimization" ], "hardware": "CPU; target accelerator execution requires separate evidence", "project": "A differentiable inverse problem", "checkpoint": "Recover a known physical parameter and report error, stability, and gradient checks.", "source": "https://docs.kidger.site/diffrax/", "description": "Build and differentiate a physical simulator, recover a positive rate from observations, and separate numerical accuracy, identifiability and measured execution cost.", "learningAdvice": "Complete gradients, explicit state and scan first. Start with analytic references, check discretization and sensitivity, then fit parameters and defend held-out predictions.", "projectId": "scientific-inverse", "studyGuide": { "question": "Are you fitting the physical model or the solver’s error?", "why": "Use an analytic system to separate simulation correctness, numerical discretization and parameter inference. An accurate autodiff derivative can still belong to an inaccurate discretized model.", "readiness": { "prompt": "A solver’s output is differentiable. Does that establish that its gradient matches the continuous differential equation?", "answer": "No. Autodiff follows the implemented discrete steps. Compare a discrete reference and refine the step size toward an analytic or higher-accuracy continuous reference.", "lessonId": "science-03" }, "milestones": [ { "lessonIds": [ "science-01", "science-02" ], "title": "Validate dynamics and axes", "task": "Compare one integration step independently, refine the grid and separate independent trajectory axes from dependent time steps." }, { "lessonIds": [ "science-03", "science-04" ], "title": "Fit an identifiable inverse problem", "task": "Check objective derivatives, choose observations that constrain the unknown parameter and evaluate predictions at held-out times. Separate noise sensitivity from discretization bias." } ], "transfer": { "prompt": "All observations are taken at the initial time of a decay process. Can they identify its decay rate?", "approach": "They identify initial amplitude but contain no elapsed-time decay signal. Collect later observations and inspect parameter sensitivity; more optimizer steps cannot create missing information." }, "pitfall": { "symptom": "The fitted rate changes as the time step is refined.", "check": "Compare solver error and parameter estimates across resolutions before attributing the shift to physical variation." }, "completion": "Use scientific-inverse. Keep discrete and continuous references, a refinement study, derivative checks, recovered parameters and a held-out/noise sensitivity report.", "furtherWork": "Adaptive solvers, event handling, stiff systems, PDEs and domain-specific physical validation remain outside the scalar ODE lab." }, "projectNote": "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.", "lessons": [ { "id": "science-01", "title": "A functional physical simulation", "status": "authored", "objective": "Implement a pure scalar state transition with units and a fixed time grid. Separate a discrete implementation oracle from continuous approximation error.", "prerequisites": [], "exercise": "Change the rate to \\(0.3\\), the initial excess to \\(5\\), and use twenty steps of \\(0.1\\) seconds. Verify the full discrete trajectory and bound the endpoint error against the exponential.", "evidence": "Full geometric/exponential comparisons, fixed-horizon refinement table, unstable-step failure and repair. Keep the environment, observed outputs and your explanation of the figure.", "check": "The code matches the geometric-sequence reference exactly but differs from the exponential. What have we established?", "path": "phases/10-science/01-a-functional-physical-simulation", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/10-science/01-a-functional-physical-simulation/docs/en.md", "script": "exercises/science-01.py", "notebook": "exercises/science-01.ipynb", "markdown": "exercises/science-01.md", "scriptSource": "phases/10-science/01-a-functional-physical-simulation/code/main.py", "notebookSource": "phases/10-science/01-a-functional-physical-simulation/notebooks/exercise.ipynb", "quizSource": "phases/10-science/01-a-functional-physical-simulation/quiz.json", "evidenceSource": "phases/10-science/01-a-functional-physical-simulation/outputs/evidence.md", "evidence": "exercises/science-01-evidence.md" }, "minutes": 85, "contentHash": "8486d3e81207351a75dfa4705d96508dd920bb659388c6ab92450b3530072db8", "visualArtifact": { "image": "figures/science-01.svg", "png": "figures/science-01.png", "data": "figures/science-01.json", "generatedAt": "2026-10-20T10:46:07.165936+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "8486d3e81207351a75dfa4705d96508dd920bb659388c6ab92450b3530072db8", "kind": "executed" }, "execution": { "contentHash": "8486d3e81207351a75dfa4705d96508dd920bb659388c6ab92450b3530072db8", "executedAt": "2026-10-20T10:47:54.140127+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Euler final: 0.3570125000000001 analytic final: 0.493193927883213\nEndpoint errors: [0.13618143 0.06399208 0.03107265 0.0153167 ]\nLast refinement ratio: 2.0286780018480166\nEuler final: 0.3570125000000001 analytic final: 0.493193927883213\nEndpoint errors: [0.13618143 0.06399208 0.03107265 0.0153167 ]\nLast refinement ratio: 2.0286780018480166\nStable but sign-alternating: [ 2. -0.8 0.32 -0.128 0.0512]\nUnstable numerical endpoint: 3.543121999999993\nChanged rate, initial state and grid verified\nEnergy checks distinguish growth, but not sign preservation\nSame-horizon errors: [np.float64(0.050590770107151406), np.float64(0.02471618055802971), np.float64(0.012218898970361547)]\nPASS: science-01\n", "url": "executions/science-01.json" }, "url": "lesson.html?lesson=science-01" }, { "id": "science-02", "title": "Vectorize trajectories and scan time", "status": "authored", "objective": "Construct and verify a fourth-order fixed-step transition. Distinguish batch-major from time-major trajectory layouts.", "prerequisites": [ "science-01" ], "exercise": "Add a fourth initial condition, verify the changed output shape and exponential values, then verify that reordering initial conditions reorders rows only.", "evidence": "Polynomial/exponential references, batch/time shape diagram, permutation and final-only checks. Keep the environment, observed outputs and your explanation of the figure.", "check": "What does a shape of (3, 41) mean in the batch-major solver?", "path": "phases/10-science/02-vectorize-trajectories-and-scan-time", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/10-science/02-vectorize-trajectories-and-scan-time/docs/en.md", "script": "exercises/science-02.py", "notebook": "exercises/science-02.ipynb", "markdown": "exercises/science-02.md", "scriptSource": "phases/10-science/02-vectorize-trajectories-and-scan-time/code/main.py", "notebookSource": "phases/10-science/02-vectorize-trajectories-and-scan-time/notebooks/exercise.ipynb", "quizSource": "phases/10-science/02-vectorize-trajectories-and-scan-time/quiz.json", "evidenceSource": "phases/10-science/02-vectorize-trajectories-and-scan-time/outputs/evidence.md", "evidence": "exercises/science-02-evidence.md" }, "minutes": 100, "contentHash": "234d7211d2bf84a53616cb343de4501461839ddc5882b35fa1fe0f6f3f80225e", "visualArtifact": { "image": "figures/science-02.svg", "png": "figures/science-02.png", "data": "figures/science-02.json", "generatedAt": "2026-10-20T10:48:19.511737+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "234d7211d2bf84a53616cb343de4501461839ddc5882b35fa1fe0f6f3f80225e", "kind": "executed", "mechanismImage": "figures/science-02-mechanism.svg", "mechanismPng": "figures/science-02-mechanism.png" }, "execution": { "contentHash": "234d7211d2bf84a53616cb343de4501461839ddc5882b35fa1fe0f6f3f80225e", "executedAt": "2026-10-20T10:47:57.386723+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Shape: (3, 41) endpoints: [0.12329848 0.24659697 0.49319394]\nscan(vectors) transposed equals vmap(scan)\nShape: (3, 41) endpoints: [0.12329848 0.24659697 0.49319394]\nscan(vectors) transposed equals vmap(scan)\nRK4 polynomial maximum difference: 1.4432899320127035e-15\nPaired-rate endpoints: [0.33516002 0.24659697 0.22160636]\nPermutation preserves trajectory identity\nOnly final states returned: [0.12329848 0.24659697 0.49319394]\nGrid axes: rate, initial condition, time (3, 3, 41)\nPASS: science-02\n", "url": "executions/science-02.json" }, "url": "lesson.html?lesson=science-02" }, { "id": "science-03", "title": "Differentiate through a solver", "status": "authored", "objective": "Derive an independent discrete solver sensitivity. Check an observation-loss gradient at a nonoptimal parameter.", "prerequisites": [ "science-02" ], "exercise": "Change the initial value to \\(3\\) and the rate to \\(0.4\\). Verify the endpoint gradient against the continuous derivative, then check the derivative with respect to the initial value as well.", "evidence": "Polynomial derivative, finite-difference window, sensitivity refinement and detached-gradient diagnosis. Keep the environment, observed outputs and your explanation of the figure.", "check": "An autodiff derivative matches the discrete RK4 formula, but both differ from the exact ODE derivative. What should you test next?", "path": "phases/10-science/03-differentiate-through-a-solver", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/10-science/03-differentiate-through-a-solver/docs/en.md", "script": "exercises/science-03.py", "notebook": "exercises/science-03.ipynb", "markdown": "exercises/science-03.md", "scriptSource": "phases/10-science/03-differentiate-through-a-solver/code/main.py", "notebookSource": "phases/10-science/03-differentiate-through-a-solver/notebooks/exercise.ipynb", "quizSource": "phases/10-science/03-differentiate-through-a-solver/quiz.json", "evidenceSource": "phases/10-science/03-differentiate-through-a-solver/outputs/evidence.md", "evidence": "exercises/science-03-evidence.md" }, "minutes": 105, "contentHash": "6b175d90141c2052127d85736ccf34c11f3f035c2884507fed92e387b08e7e68", "visualArtifact": { "image": "figures/science-03.svg", "png": "figures/science-03.png", "data": "figures/science-03.json", "generatedAt": "2026-10-20T10:46:09.193417+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "6b175d90141c2052127d85736ccf34c11f3f035c2884507fed92e387b08e7e68", "kind": "executed" }, "execution": { "contentHash": "6b175d90141c2052127d85736ccf34c11f3f035c2884507fed92e387b08e7e68", "executedAt": "2026-10-20T10:48:03.834483+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Autodiff: -0.9863878096749447 discrete oracle: -0.9863878096749429 continuous oracle: -0.9863878557664257\nLoss: 0.04837454286198381 gradient: 0.2686922266772409 finite difference: 0.26869222316237146\nAutodiff: -0.9863878096749447 discrete oracle: -0.9863878096749429 continuous oracle: -0.9863878557664257\nLoss: 0.04837454286198381 gradient: 0.2686922266772409 finite difference: 0.26869222316237146\nSensitivity errors: [np.float64(0.0002448795513753099), np.float64(1.3192617195456613e-05), np.float64(7.654351459329689e-07), np.float64(4.6091481076260266e-08)]\nFinite-difference absolute errors: [3.515348206362123e-05, 3.5151551858181307e-07, 3.514869451048952e-09, 3.5118241648035564e-11, 1.1179057679555626e-11]\nRate and initial-state sensitivities: -2.695973769291141 0.44932896460456245\nWeighted derivative reference: 0.30367311533053687\nDetached autodiff is zero; actual value sensitivity: 0.26869222316237146\nPASS: science-03\n", "url": "executions/science-03.json" }, "url": "lesson.html?lesson=science-03" }, { "id": "science-04", "title": "Fit parameters and scale the simulation", "status": "authored", "objective": "Fit a constrained physical parameter through an executed simulator. Separate training fit, held-out prediction and solver refinement.", "prerequisites": [ "science-03" ], "exercise": "Repeat inference for a new true rate \\(0.4\\) using analytically generated observations, starting at \\(0.9\\). Show that the recovered parameter and held-out predictions change consistently.", "evidence": "Full loss/rate histories, held-out and refined-grid metrics, noisy reference search, timing boundary and limitations. Keep the environment, observed outputs and your explanation of the figure.", "check": "You fit only observations at time zero and every rate gives zero loss. What is the problem?", "path": "phases/10-science/04-fit-parameters-and-scale-the-simulation", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/10-science/04-fit-parameters-and-scale-the-simulation/docs/en.md", "script": "exercises/science-04.py", "notebook": "exercises/science-04.ipynb", "markdown": "exercises/science-04.md", "scriptSource": "phases/10-science/04-fit-parameters-and-scale-the-simulation/code/main.py", "notebookSource": "phases/10-science/04-fit-parameters-and-scale-the-simulation/notebooks/exercise.ipynb", "quizSource": "phases/10-science/04-fit-parameters-and-scale-the-simulation/quiz.json", "evidenceSource": "phases/10-science/04-fit-parameters-and-scale-the-simulation/outputs/evidence.md", "evidence": "exercises/science-04-evidence.md" }, "minutes": 115, "contentHash": "bb80ff68a6e9f9b6cf42f0a888fc8bd47dd6ac73b99118ee76dcfd9e4baad8e2", "visualArtifact": { "image": "figures/science-04.svg", "png": "figures/science-04.png", "data": "figures/science-04.json", "generatedAt": "2026-10-20T10:46:15.322632+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "bb80ff68a6e9f9b6cf42f0a888fc8bd47dd6ac73b99118ee76dcfd9e4baad8e2", "kind": "executed" }, "execution": { "contentHash": "bb80ff68a6e9f9b6cf42f0a888fc8bd47dd6ac73b99118ee76dcfd9e4baad8e2", "executedAt": "2026-10-20T10:48:32.222051+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Initial/fitted rate: 1.4 0.7000000090128304\nInitial/final loss: 0.11676630453561089 1.9484774169068916e-32\nHeld-out RMSE: 4.572895313141112e-16 finer-grid RMSE: 7.282689372293896e-09\nInitial-time-only gradient: 0.0\nInitial/fitted rate: 1.4 0.7000000090128304\nInitial/final loss: 0.11676630453561089 1.9484774169068916e-32\nHeld-out RMSE: 4.572895313141112e-16 finer-grid RMSE: 7.282689372293896e-09\nInitial-time-only gradient: 0.0\nTrue rate: 0.7 perfect coarse Euler rate: 0.5906238205625731\nbatch, median milliseconds, milliseconds per trajectory: [(1, 0.007125083357095718, 0.007125083357095718), (16, 0.007792375981807709, 0.0004870234988629818), (256, 0.030416063964366913, 0.00011881274986080825)]\nChanged true/recovered rate: 0.4 0.40000000054230433\nRecovered initial amplitude: 2.8000000000000003\nNoisy grid-search estimate: 0.7 residual MSE: 2.6000611439697964e-05\nPASS: science-04\n", "url": "executions/science-04.json" }, "url": "lesson.html?lesson=science-04" } ] }, { "id": "probability", "number": "11", "title": "Probabilistic modeling", "part": "Specializations", "prerequisitePhaseIds": [ "optimization" ], "hardware": "CPU; transparent JAX samplers with an optional NumPyro ecosystem bridge", "project": "A Bayesian model with diagnostic evidence", "checkpoint": "Report sampler diagnostics and distinguish uncertainty from predictive error.", "source": "https://num.pyro.ai/en/stable/", "description": "Specify probability models, derive exact Bayesian regression, implement and diagnose HMC, and compare variational approximations with predictive evidence.", "learningAdvice": "Start with likelihoods, stable losses and explicit random keys. Check a known posterior before trusting a sampler, then separate inference error from model mismatch.", "projectId": "bayesian-regression", "studyGuide": { "question": "Which uncertainty does your model actually represent?", "why": "Use an exactly solvable Bayesian regression to check inference. Separate uncertainty about weights, uncertainty about a future observation, and mismatch between the assumed likelihood and the data.", "readiness": { "prompt": "Should an interval for a new noisy observation equal an interval for the latent mean prediction?", "answer": "No. The observation includes additional likelihood noise. Under the independent Gaussian model, add its variance to the latent predictive variance before taking a square root.", "lessonId": "probability-02" }, "milestones": [ { "lessonIds": [ "probability-01", "probability-02" ], "title": "Build an exact reference", "task": "Keep distribution units and normalization, derive a posterior in a small solvable case and compare latent versus observation uncertainty." }, { "lessonIds": [ "probability-03", "probability-04" ], "title": "Challenge approximate inference", "task": "Compare sampler moments with the exact target, inspect multiple-chain diagnostics and measure the covariance lost by a diagonal approximation. Use predictive residuals to challenge the likelihood itself." } ], "transfer": { "prompt": "Chains agree closely, but prediction residuals bend systematically with the input. What should be changed first?", "approach": "Investigate the model’s mean function and observation assumptions. Agreement between chains concerns exploration of the chosen posterior; it does not certify that the chosen model describes the data." }, "pitfall": { "symptom": "Intervals are narrow and diagnostics look reassuring despite poor predictions.", "check": "Check whether you plotted latent or observation uncertainty, and whether the likelihood or mean function misses structure." }, "completion": "Use bayesian-regression. Retain analytic posterior checks, sampler/approximation diagnostics and a predictive model-mismatch case, with uncertainty assumptions stated.", "furtherWork": "General hierarchical models and robust modern diagnostic workflows need further examples. The instructional sampler is not a replacement for validated inference libraries." }, "projectNote": "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.", "lessons": [ { "id": "probability-01", "title": "Distributions and sampling", "status": "authored", "objective": "Distinguish samples, density, variance and standard error Check normal moments against independent analytic expectations", "prerequisites": [], "exercise": "Change the sensor to mean \\(-1\\), standard deviation \\(0.5\\), and draw \\(30000\\) readings from a new key. Predict the variance and the standard error of the sample mean, then verify both moments.", "evidence": "Seeded arrays, independent density values, moment discrepancies and underflow diagnosis. Keep the environment, observed outputs and your explanation of the figure.", "check": "Which statement explains why a large Monte Carlo sample can have a precise mean while individual readings remain variable?", "path": "phases/11-probability/01-distributions-and-sampling", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/11-probability/01-distributions-and-sampling/docs/en.md", "script": "exercises/probability-01.py", "notebook": "exercises/probability-01.ipynb", "markdown": "exercises/probability-01.md", "scriptSource": "phases/11-probability/01-distributions-and-sampling/code/main.py", "notebookSource": "phases/11-probability/01-distributions-and-sampling/notebooks/exercise.ipynb", "quizSource": "phases/11-probability/01-distributions-and-sampling/quiz.json", "evidenceSource": "phases/11-probability/01-distributions-and-sampling/outputs/evidence.md", "evidence": "exercises/probability-01-evidence.md" }, "minutes": 90, "contentHash": "07740e69631bef1b6b0ef282ad669820ddd7ea9ea6b8a9a681b0cf8b2661e1e5", "visualArtifact": { "image": "figures/probability-01.svg", "png": "figures/probability-01.png", "data": "figures/probability-01.json", "generatedAt": "2026-10-20T10:46:16.734000+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "07740e69631bef1b6b0ef282ad669820ddd7ea9ea6b8a9a681b0cf8b2661e1e5", "kind": "executed" }, "execution": { "contentHash": "07740e69631bef1b6b0ef282ad669820ddd7ea9ea6b8a9a681b0cf8b2661e1e5", "executedAt": "2026-10-20T10:48:36.575595+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "sample mean, variance: 2.0067408084869385 9.035282135009766\nstable equal-mixture log density: -1000.3799438476562\nsample mean, variance: 2.0067408084869385 9.035282135009766\nstable equal-mixture log density: -1000.3799438476562\npeak density: 3.9894230365753174\n-5.0 -2.379885673522949\nPASS: probability-01\n", "url": "executions/probability-01.json" }, "url": "lesson.html?lesson=probability-01" }, { "id": "probability-02", "title": "Build a Bayesian regression", "status": "authored", "objective": "State the prior, likelihood and observation-noise assumptions Derive posterior precision and check a closed-form special case", "prerequisites": [ "probability-01" ], "exercise": "Strengthen the prior by changing its scale from two to one half. Predict the direction of change in the weight norm and check it. Explain why a tighter posterior need not mean a better model.", "evidence": "Precision derivation, exact fractional reference, host solve and predictive variance plot. Keep the environment, observed outputs and your explanation of the figure.", "check": "Which variance describes a new noisy reading at a given input?", "path": "phases/11-probability/02-build-a-bayesian-regression", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/11-probability/02-build-a-bayesian-regression/docs/en.md", "script": "exercises/probability-02.py", "notebook": "exercises/probability-02.ipynb", "markdown": "exercises/probability-02.md", "scriptSource": "phases/11-probability/02-build-a-bayesian-regression/code/main.py", "notebookSource": "phases/11-probability/02-build-a-bayesian-regression/notebooks/exercise.ipynb", "quizSource": "phases/11-probability/02-build-a-bayesian-regression/quiz.json", "evidenceSource": "phases/11-probability/02-build-a-bayesian-regression/outputs/evidence.md", "evidence": "exercises/probability-02-evidence.md" }, "minutes": 110, "contentHash": "cf43427cba47d5d98dec188e6cbc0cb753cab53adb42fdf09c71c880b031f540", "visualArtifact": { "image": "figures/probability-02.svg", "png": "figures/probability-02.png", "data": "figures/probability-02.json", "generatedAt": "2026-10-20T10:46:18.108464+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "cf43427cba47d5d98dec188e6cbc0cb753cab53adb42fdf09c71c880b031f540", "kind": "executed" }, "execution": { "contentHash": "cf43427cba47d5d98dec188e6cbc0cb753cab53adb42fdf09c71c880b031f540", "executedAt": "2026-10-20T10:48:42.156598+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "posterior mean: [0.923077 1.7777778] covariance: [[0.30769232 0. ]\n [0. 0.44444445]]\nlatent variance: [0.30769232 2.0854702 ] observation variance: [1.3076923 3.0854702]\nposterior mean: [0.923077 1.7777778] covariance: [[0.30769232 0. ]\n [0. 0.44444445]]\nlatent variance: [0.30769232 2.0854702 ] observation variance: [1.3076923 3.0854702]\n[0.16 0.23529412]\nPASS: probability-02\n", "url": "executions/probability-02.json" }, "url": "lesson.html?lesson=probability-02" }, { "id": "probability-03", "title": "Hamiltonian Monte Carlo and sampler diagnostics", "status": "authored", "objective": "Implement leapfrog with the correct half-step boundaries Check reversibility and an independent analytic gradient", "prerequisites": [ "probability-02" ], "exercise": "Increase the step size to 1.2 with the same leapfrog count. Record acceptance and positive energy errors. Explain why keeping more rejected samples is not a repair for unstable integration.", "evidence": "Reversibility check, four traces, two-dimensional moments, key replay and classical split R-hat with limitations. Keep the environment, observed outputs and your explanation of the figure.", "check": "Four chains have classical split R-hat close to one. What can we conclude?", "path": "phases/11-probability/03-hamiltonian-monte-carlo-and-sampler-diagnostics", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/11-probability/03-hamiltonian-monte-carlo-and-sampler-diagnostics/docs/en.md", "script": "exercises/probability-03.py", "notebook": "exercises/probability-03.ipynb", "markdown": "exercises/probability-03.md", "scriptSource": "phases/11-probability/03-hamiltonian-monte-carlo-and-sampler-diagnostics/code/main.py", "notebookSource": "phases/11-probability/03-hamiltonian-monte-carlo-and-sampler-diagnostics/notebooks/exercise.ipynb", "quizSource": "phases/11-probability/03-hamiltonian-monte-carlo-and-sampler-diagnostics/quiz.json", "evidenceSource": "phases/11-probability/03-hamiltonian-monte-carlo-and-sampler-diagnostics/outputs/evidence.md", "evidence": "exercises/probability-03-evidence.md" }, "minutes": 140, "contentHash": "885ece097a303a05c8c195774eb72f9abcd75ad850620788cbe97176199dae1b", "visualArtifact": { "image": "figures/probability-03.svg", "png": "figures/probability-03.png", "data": "figures/probability-03.json", "generatedAt": "2026-10-20T10:46:20.839431+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "885ece097a303a05c8c195774eb72f9abcd75ad850620788cbe97176199dae1b", "kind": "executed" }, "execution": { "contentHash": "885ece097a303a05c8c195774eb72f9abcd75ad850620788cbe97176199dae1b", "executedAt": "2026-10-20T10:48:50.098468+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "mean: [ 0.98358583 -1.0185838 ] covariance: [[0.98042357 0.7727629 ]\n [0.7727629 0.9861013 ]]\nclassical split R-hat: [1.0002148 1.000253 ] acceptance: [0.97749996 0.98625 0.98125 0.980625 ]\nmean: [ 0.98358583 -1.0185838 ] covariance: [[0.98042357 0.7727629 ]\n [0.7727629 0.9861013 ]]\nclassical split R-hat: [1.0002148 1.000253 ] acceptance: [0.97749996 0.98625 0.98125 0.980625 ]\nbad acceptance: 0.0 maximum energy error: 50229501952.0\nPASS: probability-03\n", "url": "executions/probability-03.json" }, "url": "lesson.html?lesson=probability-03" }, { "id": "probability-04", "title": "Variational inference and model checking", "status": "authored", "objective": "Derive the diagonal Gaussian variational optimum for a correlated target Separate optimization error from variational-family error", "prerequisites": [ "probability-03" ], "exercise": "Replace the target correlation by zero and rederive the optimal diagonal variances and minimum KL. Verify that the diagonal family can now represent the target exactly.", "evidence": "Analytic and optimized KL, variance comparison, Monte Carlo estimate and explained residual stress fixture. Keep the environment, observed outputs and your explanation of the figure.", "check": "The variational mean is correct and the optimization gradient is tiny. Why can posterior intervals still be too narrow?", "path": "phases/11-probability/04-variational-inference-and-model-checking", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/11-probability/04-variational-inference-and-model-checking/docs/en.md", "script": "exercises/probability-04.py", "notebook": "exercises/probability-04.ipynb", "markdown": "exercises/probability-04.md", "scriptSource": "phases/11-probability/04-variational-inference-and-model-checking/code/main.py", "notebookSource": "phases/11-probability/04-variational-inference-and-model-checking/notebooks/exercise.ipynb", "quizSource": "phases/11-probability/04-variational-inference-and-model-checking/quiz.json", "evidenceSource": "phases/11-probability/04-variational-inference-and-model-checking/outputs/evidence.md", "evidence": "exercises/probability-04-evidence.md" }, "minutes": 125, "contentHash": "6376255cbf8b00058c317b070f37c405c6d1a3d223f63852cccd9c0bcb42a727", "visualArtifact": { "image": "figures/probability-04.svg", "png": "figures/probability-04.png", "data": "figures/probability-04.json", "generatedAt": "2026-10-20T10:46:22.399640+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "6376255cbf8b00058c317b070f37c405c6d1a3d223f63852cccd9c0bcb42a727", "kind": "executed", "mechanismImage": "figures/probability-04-mechanism.svg", "mechanismPng": "figures/probability-04-mechanism.png" }, "execution": { "contentHash": "6376255cbf8b00058c317b070f37c405c6d1a3d223f63852cccd9c0bcb42a727", "executedAt": "2026-10-20T10:48:56.444623+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "VI mean: [ 0.99999994 -0.99999994] variances: [0.36000016 0.36000016] remaining KL: 0.5108256340026855\nsynthetic nominal-95% interval coverage: 0.2222222238779068\nstandardized residuals: [12. 5. 0. -3. -4. -3. 0. 5. 12.]\nVI mean: [ 0.99999994 -0.99999994] variances: [0.36000016 0.36000016] remaining KL: 0.5108256340026855\nsynthetic nominal-95% interval coverage: 0.2222222238779068\nstandardized residuals: [12. 5. 0. -3. -4. -3. 0. 5. 12.]\nMonte Carlo KL: 0.5119153261184692\nmodel self-coverage: 0.9493999481201172\nPASS: probability-04\n", "url": "executions/probability-04.json" }, "url": "lesson.html?lesson=probability-04" } ] }, { "id": "rl", "number": "12", "title": "Reinforcement learning", "part": "Specializations", "prerequisitePhaseIds": [ "networks" ], "hardware": "CPU; target accelerator execution requires separate evidence", "project": "An evaluated policy in a functional environment", "checkpoint": "Compare a trained policy to a baseline across several seeds and report variability.", "source": "https://github.com/FLAIROx/Jumanji", "description": "Build a functional environment, collect randomized trajectories, train policy-gradient and clipped-policy agents, and evaluate independent seeds.", "learningAdvice": "Keep terminal and timeout behavior explicit. Verify a small enumerated expectation before trusting a sampled policy gradient, then separate training from held-out evaluation.", "projectId": "policy-evaluation", "studyGuide": { "question": "Is the policy improving under a fixed evaluation protocol?", "why": "Environment semantics and data collection define the learning problem. Verify transitions, masks and behavior-policy probabilities before trusting a noisy return curve.", "readiness": { "prompt": "An episode stops because its time budget expires. Must its bootstrap value be zero?", "answer": "Not necessarily. Distinguish truncation from genuine terminal states and state whether the task is finite-horizon or continuing. A timeout alone does not imply that future value is zero.", "lessonId": "rl-01" }, "milestones": [ { "lessonIds": [ "rl-01", "rl-02" ], "title": "Test environment and rollout semantics", "task": "Check transitions independently, preserve terminal rewards and label time/environment axes. Store the behavior policy that collected each action." }, { "lessonIds": [ "rl-03", "rl-04" ], "title": "Compare updates and evaluations fairly", "task": "Audit the clipped objective for both advantage signs. Freeze the evaluation protocol, compare a simple baseline and distinguish training seeds from repeated episodes." } ], "transfer": { "prompt": "You evaluate one trained agent over many episodes and report a tiny uncertainty interval. Does this measure sensitivity to training randomness?", "approach": "No. Those episodes characterize that agent under the evaluation distribution. Train independent agents to assess training variation, retain agent-level means and state the experimental unit." }, "pitfall": { "symptom": "PPO ratios always equal one after every update.", "check": "Check that old log probabilities were stored at collection rather than recomputed from the current policy." }, "completion": "Use policy-evaluation. Keep independent transition checks, behavior-policy identity, changed-horizon cases and paired evaluations across declared independent training seeds.", "furtherWork": "Continuous actions, value-function estimation and larger benchmark environments need new implementation and evaluation labs; this tabular task is deliberately bounded." }, "projectNote": "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.", "lessons": [ { "id": "rl-01", "title": "Write a functional environment", "status": "authored", "objective": "Implement reset and step as pure functions with explicit state. Distinguish termination, time-limit truncation and an absorbing finished state.", "prerequisites": [], "exercise": "Begin at position \\(0\\), choose left for three actions with horizon \\(3\\), and then attempt one more action. Verify the event flags, total reward and absorbing state.", "evidence": "Keep independent transition enumeration, reward/terminal checks and the absorbing-state diagnosis. Explain the state and observation contract.", "check": "What should step emit when called on a finished state?", "path": "phases/12-rl/01-write-a-functional-environment", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/12-rl/01-write-a-functional-environment/docs/en.md", "script": "exercises/rl-01.py", "notebook": "exercises/rl-01.ipynb", "markdown": "exercises/rl-01.md", "scriptSource": "phases/12-rl/01-write-a-functional-environment/code/main.py", "notebookSource": "phases/12-rl/01-write-a-functional-environment/notebooks/exercise.ipynb", "quizSource": "phases/12-rl/01-write-a-functional-environment/quiz.json", "evidenceSource": "phases/12-rl/01-write-a-functional-environment/outputs/evidence.md", "evidence": "exercises/rl-01-evidence.md" }, "minutes": 75, "contentHash": "c955c6d35d9321a6fd91d24e9ab99063227cfc001399d2a64781de738d4dedd9", "visualArtifact": { "image": "figures/rl-01.svg", "png": "figures/rl-01.png", "data": "figures/rl-01.json", "generatedAt": "2026-10-20T10:46:23.313509+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "c955c6d35d9321a6fd91d24e9ab99063227cfc001399d2a64781de738d4dedd9", "kind": "executed" }, "execution": { "contentHash": "c955c6d35d9321a6fd91d24e9ab99063227cfc001399d2a64781de738d4dedd9", "executedAt": "2026-10-20T10:48:59.683282+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "positions: [1, 2, 3, 3, 3] rewards: [-0.009999999776482582, 1.0, 0.0, 0.0]\npositions: [1, 2, 3, 3, 3] rewards: [-0.009999999776482582, 1.0, 0.0, 0.0]\ndeadline goal: True False\nfraction starting at one: 0.5020000338554382\nPASS: rl-01\n", "url": "executions/rl-01.json" }, "url": "lesson.html?lesson=rl-01" }, { "id": "rl-02", "title": "Batch environments and scan rollouts", "status": "authored", "objective": "Build a scan over vmapped transitions with explicit key ownership. Reconstruct episode returns and reward-to-go from padded trajectories.", "prerequisites": [ "rl-01" ], "exercise": "Change the batch to \\(7\\) environments and the horizon to \\(5\\). Verify array shapes, one ending event per episode and zero rewards after finishing.", "evidence": "Keep reconstructed trajectories, randomized keys, padding masks and independent return calculations. Show how a changed key affects interaction without changing the environment rules.", "check": "Which axis should be summed to obtain one return per environment from a time-major reward array?", "path": "phases/12-rl/02-batch-environments-and-scan-rollouts", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/12-rl/02-batch-environments-and-scan-rollouts/docs/en.md", "script": "exercises/rl-02.py", "notebook": "exercises/rl-02.ipynb", "markdown": "exercises/rl-02.md", "scriptSource": "phases/12-rl/02-batch-environments-and-scan-rollouts/code/main.py", "notebookSource": "phases/12-rl/02-batch-environments-and-scan-rollouts/notebooks/exercise.ipynb", "quizSource": "phases/12-rl/02-batch-environments-and-scan-rollouts/quiz.json", "evidenceSource": "phases/12-rl/02-batch-environments-and-scan-rollouts/outputs/evidence.md", "evidence": "exercises/rl-02-evidence.md" }, "minutes": 85, "contentHash": "2b067eae4b74d07513fcdf4a0311cb087d4b88923073fdc1669c330e647038ed", "visualArtifact": { "image": "figures/rl-02.svg", "png": "figures/rl-02.png", "data": "figures/rl-02.json", "generatedAt": "2026-10-20T10:46:24.417916+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "2b067eae4b74d07513fcdf4a0311cb087d4b88923073fdc1669c330e647038ed", "kind": "executed", "mechanismImage": "figures/rl-02-mechanism.svg", "mechanismPng": "figures/rl-02-mechanism.png" }, "execution": { "contentHash": "2b067eae4b74d07513fcdf4a0311cb087d4b88923073fdc1669c330e647038ed", "executedAt": "2026-10-20T10:49:07.271650+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "return mean: 0.5848437547683716\nactive environments by transition: [64 64 58 53 47 44 39 29]\nreturn mean: 0.5848437547683716\nactive environments by transition: [64 64 58 53 47 44 39 29]\nreward-to-go: [0.98 0.99 1. 0. ]\nPASS: rl-02\n", "url": "executions/rl-02.json" }, "url": "lesson.html?lesson=rl-02" }, { "id": "rl-03", "title": "Policy gradients and PPO", "status": "authored", "objective": "Derive and implement the score-function policy gradient for complete episodes. Train with fresh stochastic rollouts and a fixed behavior policy per PPO batch.", "prerequisites": [ "rl-02" ], "exercise": "Run REINFORCE with a different training seed and retain the initial and final exact return. Verify that the final policy improves on the uniform baseline without requiring every intermediate point to improve.", "evidence": "Keep the exact small-state policy-gradient comparison, actual training curves and frozen PPO reference quantities. Diagnose changed clipping signs and distinguish Monte Carlo variation from an incorrect objective.", "check": "During multiple PPO epochs on one rollout batch, what must remain fixed?", "path": "phases/12-rl/03-policy-gradients-and-ppo", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/12-rl/03-policy-gradients-and-ppo/docs/en.md", "script": "exercises/rl-03.py", "notebook": "exercises/rl-03.ipynb", "markdown": "exercises/rl-03.md", "scriptSource": "phases/12-rl/03-policy-gradients-and-ppo/code/main.py", "notebookSource": "phases/12-rl/03-policy-gradients-and-ppo/notebooks/exercise.ipynb", "quizSource": "phases/12-rl/03-policy-gradients-and-ppo/quiz.json", "evidenceSource": "phases/12-rl/03-policy-gradients-and-ppo/outputs/evidence.md", "evidence": "exercises/rl-03-evidence.md" }, "minutes": 110, "contentHash": "d6bbf2419915a048da0fe938b65db4c4792e2070f4c62e6f1714251e0ef5c702", "visualArtifact": { "image": "figures/rl-03.svg", "png": "figures/rl-03.png", "data": "figures/rl-03.json", "generatedAt": "2026-10-20T10:46:25.846602+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "d6bbf2419915a048da0fe938b65db4c4792e2070f4c62e6f1714251e0ef5c702", "kind": "executed" }, "execution": { "contentHash": "d6bbf2419915a048da0fe938b65db4c4792e2070f4c62e6f1714251e0ef5c702", "executedAt": "2026-10-20T10:49:19.804218+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "REINFORCE initial/final exact return: 0.46687505 0.9521624\nPPO initial/final exact return: 0.46687505 0.97578394\nREINFORCE initial/final exact return: 0.46687505 0.9521624\nPPO initial/final exact return: 0.46687505 0.97578394\npositive / negative targets: [0.6 1. 1.2] [-0.8 -1. -1.4]\nsampled / exact gradient: [0.12067959 0.21948992 0.13326262] [0.12001554 0.21822323 0.13327113]\nchanged seed final exact return: 0.9501661\nPASS: rl-03\n", "url": "executions/rl-03.json" }, "url": "lesson.html?lesson=rl-03" }, { "id": "rl-04", "title": "Evaluate agents across seeds", "status": "authored", "objective": "Separate training seeds, evaluation streams and individual episodes. Report per-agent means and variability with an explicit aggregation unit.", "prerequisites": [ "rl-03" ], "exercise": "Compute the paired held-out improvement over the uniform baseline for each trained agent. Report all \\(5\\) differences and their mean, keeping the training run as the unit.", "evidence": "Keep multiple trained seeds, fixed held-out streams, exact versus sampled returns and failed trajectories. Explain uncertainty and demonstrate evaluation leaves the policy unchanged.", "check": "What does evaluating one policy on many episodes fail to measure?", "path": "phases/12-rl/04-evaluate-agents-across-seeds", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/12-rl/04-evaluate-agents-across-seeds/docs/en.md", "script": "exercises/rl-04.py", "notebook": "exercises/rl-04.ipynb", "markdown": "exercises/rl-04.md", "scriptSource": "phases/12-rl/04-evaluate-agents-across-seeds/code/main.py", "notebookSource": "phases/12-rl/04-evaluate-agents-across-seeds/notebooks/exercise.ipynb", "quizSource": "phases/12-rl/04-evaluate-agents-across-seeds/quiz.json", "evidenceSource": "phases/12-rl/04-evaluate-agents-across-seeds/outputs/evidence.md", "evidence": "exercises/rl-04-evidence.md" }, "minutes": 90, "contentHash": "60ee26b667e26e302f7f6417e510aff429e2068da88b725c7268c907aa9f2a03", "visualArtifact": { "image": "figures/rl-04.svg", "png": "figures/rl-04.png", "data": "figures/rl-04.json", "generatedAt": "2026-10-20T10:46:27.202039+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "60ee26b667e26e302f7f6417e510aff429e2068da88b725c7268c907aa9f2a03", "kind": "executed" }, "execution": { "contentHash": "60ee26b667e26e302f7f6417e510aff429e2068da88b725c7268c907aa9f2a03", "executedAt": "2026-10-20T10:49:26.500487+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "training seeds: [0, 7, 23, 41, 59]\nheld-out mean per agent: [0.9759717 0.97472161 0.97537106 0.96944821 0.97455072]\nexact mean per agent: [0.97578394 0.9749347 0.97562152 0.97113729 0.97404611]\nuniform baseline mean: 0.46628910303115845\nacross-agent mean / sample sd: 0.9740126609802247 0.0026129212275767476\ntraining seeds: [0, 7, 23, 41, 59]\nheld-out mean per agent: [0.9759717 0.97472161 0.97537106 0.96944821 0.97455072]\nexact mean per agent: [0.97578394 0.9749347 0.97562152 0.97113729 0.97404611]\nuniform baseline mean: 0.46628910303115845\nacross-agent mean / sample sd: 0.9740126609802247 0.0026129212275767476\nleft / random / right: -0.07999999821186066 0.46628910303115845 0.9850389361381531\npaired improvements: [0.5096826 0.50843251 0.50908196 0.50315911 0.50826162]\nmean improvement across trained agents: 0.5077235579490662\nsample sd / independence-assuming SE: 0.09999999999999998 0.05773502691896257\nPASS: rl-04\n", "url": "executions/rl-04.json" }, "url": "lesson.html?lesson=rl-04" } ] }, { "id": "kernels", "number": "13", "title": "Pallas kernels", "part": "Specializations", "prerequisitePhaseIds": [ "distributed" ], "hardware": "CPU interpretation/simulation reference; TPU target lab requires separately validated hardware", "project": "A checked Pallas kernel with an explicit target qualification plan", "checkpoint": "Compare against a reference across shapes and justify each optimization with evidence.", "source": "https://docs.jax.dev/en/latest/pallas/index.html", "description": "Write actual Pallas grids and buffered pipelines, verify boundaries and precision on CPU, and prepare guarded target measurements before claiming an optimization win.", "learningAdvice": "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.", "projectId": "kernel-audit", "studyGuide": { "question": "Does a custom kernel earn its complexity?", "why": "A lower-level implementation first owes you the same numerical contract. Use hostile shapes, independent references and explicit target limits before attempting a speed claim.", "readiness": { "prompt": "A matrix dimension is not divisible by the tile size. What must the last tile do?", "answer": "Mask or pad out-of-range loads and stores under a declared policy. Valid elements still need exactly one correct output; fewer launched programs do not excuse missing boundary elements.", "lessonId": "kernels-01" }, "milestones": [ { "lessonIds": [ "kernels-01", "kernels-02" ], "title": "Prove indexing and arithmetic", "task": "Map each grid coordinate to input/output slices, test rectangular and boundary shapes and reject unsupported layouts or dtypes before launch." }, { "lessonIds": [ "kernels-03", "kernels-04" ], "title": "Measure only a supported target", "task": "Check pipeline semantics and buffer assumptions, compare with a strong compiled baseline and make an end-to-end decision from the operation’s measured share." } ], "transfer": { "prompt": "Interpret mode passes all numerical tests. Can you use its timings to choose the fastest TPU tile?", "approach": "No. Interpretation supports bounded semantic checks. Compile and measure supported candidates on the named TPU with equivalent work and synchronization, including numerical verification there." }, "pitfall": { "symptom": "A kernel passes square inputs but corrupts rectangular ones.", "check": "Inspect operand-specific index maps and boundary masks; square dimensions can hide an axis swap." }, "completion": "Use kernel-audit. Retain independent numerical checks, rejected unsupported contracts and a target measurement plan. Keep hardware timings unmeasured until actually executed.", "furtherWork": "Actual TPU/GPU compilation, race diagnosis and target tuning remain unqualified; CPU interpretation cannot close these gaps." }, "projectNote": "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.", "lessons": [ { "id": "kernels-01", "title": "Pallas grids and BlockSpecs", "status": "authored", "objective": "Implement an actual Pallas Ref-writing kernel with grid and BlockSpecs. Derive block coordinates and distinguish them from element offsets.", "prerequisites": [], "exercise": "Test a singleton row and a changed nondivisible matrix with negative values. Verify full NumPy agreement and exact output shapes for two block choices.", "evidence": "Independent full-array comparisons, ownership map, tail fixtures and a floor-grid omission diagnosis. Keep the environment, observed outputs and your explanation of the figure.", "check": "For block shape (2,4), what does a BlockSpec index map returning (i,j) select?", "path": "phases/13-kernels/01-pallas-grids-and-blockspecs", "runtime": { "kind": "cpu" }, "hardware": "CPU Pallas interpretation; actual TPU lowering and performance unverified", "artifacts": { "source": "phases/13-kernels/01-pallas-grids-and-blockspecs/docs/en.md", "script": "exercises/kernels-01.py", "notebook": "exercises/kernels-01.ipynb", "markdown": "exercises/kernels-01.md", "scriptSource": "phases/13-kernels/01-pallas-grids-and-blockspecs/code/main.py", "notebookSource": "phases/13-kernels/01-pallas-grids-and-blockspecs/notebooks/exercise.ipynb", "quizSource": "phases/13-kernels/01-pallas-grids-and-blockspecs/quiz.json", "evidenceSource": "phases/13-kernels/01-pallas-grids-and-blockspecs/outputs/evidence.md", "evidence": "exercises/kernels-01-evidence.md" }, "minutes": 95, "contentHash": "0a36a139ed87e8fad68216085d916f9734fc00a60182c82f9e3da5a0ed3dff19", "visualArtifact": { "image": "figures/kernels-01.svg", "png": "figures/kernels-01.png", "data": "figures/kernels-01.json", "generatedAt": "2026-10-20T10:46:28.332597+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "0a36a139ed87e8fad68216085d916f9734fc00a60182c82f9e3da5a0ed3dff19", "kind": "executed", "mechanismImage": "figures/kernels-01-mechanism.svg", "mechanismPng": "figures/kernels-01-mechanism.png" }, "execution": { "contentHash": "0a36a139ed87e8fad68216085d916f9734fc00a60182c82f9e3da5a0ed3dff19", "executedAt": "2026-10-20T10:49:31.979818+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Shape and boundary: (5, 11) -3.75 7.050000190734863\nGrid: (3, 3) logical/padded elements: 55 72\nPadding fraction: 0.2361111111111111\nShape and boundary: (5, 11) -3.75 7.050000190734863\nGrid: (3, 3) logical/padded elements: 55 72\nPadding fraction: 0.2361111111111111\nChanged grid: (2, 3) same numerical result\nFloor-grid missing logical cells: 23\nSingleton and changed tails verified\nMismatched global shape rejected\nShape, block/program/padded counts: (5, 11) [(2, 4, 9, 72), (8, 128, 1, 1024), (16, 256, 1, 4096)]\nShape, block/program/padded counts: (33, 129) [(2, 4, 561, 4488), (8, 128, 10, 10240), (16, 256, 3, 12288)]\nPASS: kernels-01\n", "url": "executions/kernels-01.json" }, "url": "lesson.html?lesson=kernels-01" }, { "id": "kernels-02", "title": "A first TPU kernel", "status": "authored", "objective": "Implement a real TPU-targeted Pallas launch with a broadcast BlockSpec. Specify target layout and float32 compute/output-cast policies.", "prerequisites": [ "kernels-01" ], "exercise": "Change the activation shape to (17,257), freeze a NumPy random seed and compare two target-compatible block choices in interpretation. Keep the bias values distinct by column.", "evidence": "Broadcast/boundary checks, represented-input precision oracle and a tested no-fallback TPU guard; target execution remains separately unverified. Keep the environment, observed outputs and your explanation of the figure.", "check": "What does a successful CPU interpret run establish for this target-oriented kernel?", "path": "phases/13-kernels/02-a-first-tpu-kernel", "runtime": { "kind": "cpu" }, "hardware": "CPU Pallas interpretation; actual TPU lowering and performance unverified", "artifacts": { "source": "phases/13-kernels/02-a-first-tpu-kernel/docs/en.md", "script": "exercises/kernels-02.py", "notebook": "exercises/kernels-02.ipynb", "markdown": "exercises/kernels-02.md", "scriptSource": "phases/13-kernels/02-a-first-tpu-kernel/code/main.py", "notebookSource": "phases/13-kernels/02-a-first-tpu-kernel/notebooks/exercise.ipynb", "quizSource": "phases/13-kernels/02-a-first-tpu-kernel/quiz.json", "evidenceSource": "phases/13-kernels/02-a-first-tpu-kernel/outputs/evidence.md", "evidence": "exercises/kernels-02-evidence.md" }, "minutes": 110, "contentHash": "59bc73b55f193ea8dbea59a1a0b4aa9ebf1cf8688fa7788d85092e22c27f943c", "visualArtifact": { "image": "figures/kernels-02.svg", "png": "figures/kernels-02.png", "data": "figures/kernels-02.json", "generatedAt": "2026-10-20T10:46:29.396490+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "59bc73b55f193ea8dbea59a1a0b4aa9ebf1cf8688fa7788d85092e22c27f943c", "kind": "executed" }, "execution": { "contentHash": "59bc73b55f193ea8dbea59a1a0b4aa9ebf1cf8688fa7788d85092e22c27f943c", "executedAt": "2026-10-20T10:49:37.641695+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Interpretation max error: 0.0\nZero fraction and endpoint: 0.48578811369509045 2.4000000953674316\nExpected target refusal: Real TPU inputs/backend required; CPU fallback is disabled\nInterpretation max error: 0.0\nZero fraction and endpoint: 0.48578811369509045 2.4000000953674316\nExpected target refusal: Real TPU inputs/backend required; CPU fallback is disabled\nbfloat16 represented-input reference matched exactly\nBoundary fixture: [[0. 0. 0.]\n [2. 0. 0.]]\nChanged tails and two TPU-compatible blocks verified in interpretation\nMalformed bias axes rejected\nColumn/bias pairing survives a joint permutation\nPASS: kernels-02\n", "url": "executions/kernels-02.json" }, "url": "lesson.html?lesson=kernels-02" }, { "id": "kernels-03", "title": "Tiling, memory, and pipelining", "status": "authored", "objective": "Execute actual emit_pipeline semantics under explicit CPU TPU-layout simulation. Compare synchronous and buffered schedules against an independent oracle.", "prerequisites": [ "kernels-02" ], "exercise": "Use a (9,129) random fixture and verify synchronous, double-input-buffered and triple-input-buffered pipelines for both float32 and bfloat16 represented-input contracts. Keep output buffering at two slots.", "evidence": "Synchronous/buffered CPU simulation agreement, two/three input slots, unchanged output buffering, correctness across tiles and modeled footprint counts. Keep the environment, observed outputs and your explanation of the figure.", "check": "The three-buffer simulation is correct and has fewer grid programs after a tile change. Can you conclude the TPU version is faster?", "path": "phases/13-kernels/03-tiling-memory-and-pipelining", "runtime": { "kind": "cpu" }, "hardware": "CPU Pallas interpretation; actual TPU lowering and performance unverified", "artifacts": { "source": "phases/13-kernels/03-tiling-memory-and-pipelining/docs/en.md", "script": "exercises/kernels-03.py", "notebook": "exercises/kernels-03.ipynb", "markdown": "exercises/kernels-03.md", "scriptSource": "phases/13-kernels/03-tiling-memory-and-pipelining/code/main.py", "notebookSource": "phases/13-kernels/03-tiling-memory-and-pipelining/notebooks/exercise.ipynb", "quizSource": "phases/13-kernels/03-tiling-memory-and-pipelining/quiz.json", "evidenceSource": "phases/13-kernels/03-tiling-memory-and-pipelining/outputs/evidence.md", "evidence": "exercises/kernels-03-evidence.md" }, "minutes": 120, "contentHash": "ffafe87d6ffe3d1dddcd0ab058e87d4adff0aae938e053a9460b4bb5cbe3f3b6", "visualArtifact": { "image": "figures/kernels-03.svg", "png": "figures/kernels-03.png", "data": "figures/kernels-03.json", "generatedAt": "2026-10-20T10:46:31.177124+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "ffafe87d6ffe3d1dddcd0ab058e87d4adff0aae938e053a9460b4bb5cbe3f3b6", "kind": "executed" }, "execution": { "contentHash": "ffafe87d6ffe3d1dddcd0ab058e87d4adff0aae938e053a9460b4bb5cbe3f3b6", "executedAt": "2026-10-20T10:49:49.884306+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Actual backend: cpu simulated layout: TPU v5 lite\nSynchronous/buffered endpoints: -1.75 2.25\nprograms, padded elements, modeled data-buffer bytes: [(9, 9216, 24576), (6, 12288, 49152), (4, 16384, 98304)]\nActual backend: cpu simulated layout: TPU v5 lite\nSynchronous/buffered endpoints: -1.75 2.25\nprograms, padded elements, modeled data-buffer bytes: [(9, 9216, 24576), (6, 12288, 49152), (4, 16384, 98304)]\nThree-buffer schedule matches the two-buffer result\nSingle-tile pipeline checked\nChanged shape, dtype and schedules verified\nUnsupported pipeline tile rejected\nTile, two/three-input-slot modeled bytes: (8, 128) 24576 32768\nTile, two/three-input-slot modeled bytes: (16, 128) 49152 65536\nTile, two/three-input-slot modeled bytes: (16, 256) 98304 131072\nPASS: kernels-03\n", "url": "executions/kernels-03.json" }, "url": "lesson.html?lesson=kernels-03" }, { "id": "kernels-04", "title": "Iterate with correctness and performance evidence", "status": "authored", "objective": "Build a representative shape/dtype correctness matrix for an actual Pallas pipeline. Separate kernel implementation error from input/output precision differences.", "prerequisites": [ "kernels-03" ], "exercise": "Add a (15,255) fixture, compare two tile choices and two buffer counts, and retain the same frozen inputs for every candidate. Report errors and padded work separately.", "evidence": "Shape/dtype reference checks, separate precision gaps, guarded target benchmark and a bounded end-to-end speedup argument. Keep the environment, observed outputs and your explanation of the figure.", "check": "The kernel matches the represented-input oracle exactly but differs from the original float32 dataset in bfloat16. What should the report say?", "path": "phases/13-kernels/04-iterate-with-correctness-and-performance-evidence", "runtime": { "kind": "cpu" }, "hardware": "CPU Pallas interpretation; actual TPU lowering and performance unverified", "artifacts": { "source": "phases/13-kernels/04-iterate-with-correctness-and-performance-evidence/docs/en.md", "script": "exercises/kernels-04.py", "notebook": "exercises/kernels-04.ipynb", "markdown": "exercises/kernels-04.md", "scriptSource": "phases/13-kernels/04-iterate-with-correctness-and-performance-evidence/code/main.py", "notebookSource": "phases/13-kernels/04-iterate-with-correctness-and-performance-evidence/notebooks/exercise.ipynb", "quizSource": "phases/13-kernels/04-iterate-with-correctness-and-performance-evidence/quiz.json", "evidenceSource": "phases/13-kernels/04-iterate-with-correctness-and-performance-evidence/outputs/evidence.md", "evidence": "exercises/kernels-04-evidence.md" }, "minutes": 120, "contentHash": "502f1d93126e5e86a2e224f8d494da3faadd751606503d415c9d3c37d2f1b7d0", "visualArtifact": { "image": "figures/kernels-04.svg", "png": "figures/kernels-04.png", "data": "figures/kernels-04.json", "generatedAt": "2026-10-20T10:46:33.923095+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "502f1d93126e5e86a2e224f8d494da3faadd751606503d415c9d3c37d2f1b7d0", "kind": "executed" }, "execution": { "contentHash": "502f1d93126e5e86a2e224f8d494da3faadd751606503d415c9d3c37d2f1b7d0", "executedAt": "2026-10-20T10:50:03.077466+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Kernel errors versus represented-input oracle: [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]\nGaps versus original float32 data: [[0.0, 0.000862419605255127], [0.0, 0.025217056274414062], [0.0, 0.030788421630859375], [0.0, 0.0292205810546875]]\nExpected benchmark refusal: TPU benchmark requires real TPU inputs; interpretation timings are not accepted\nCPU correctness audit complete; no TPU latency or speedup measured\nKernel errors versus represented-input oracle: [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]\nGaps versus original float32 data: [[0.0, 0.000862419605255127], [0.0, 0.025217056274414062], [0.0, 0.030788421630859375], [0.0, 0.0292205810546875]]\nExpected benchmark refusal: TPU benchmark requires real TPU inputs; interpretation timings are not accepted\nCPU correctness audit complete; no TPU latency or speedup measured\nJoint row permutation preserves output identity\nEarly/final-only rounding: 1.0 1.0078125\nConfiguration, padded elements: (8, 128) 2 4096\nConfiguration, padded elements: (8, 128) 3 4096\nConfiguration, padded elements: (16, 256) 2 4096\nConfiguration, padded elements: (16, 256) 3 4096\nMixed and integer dtypes rejected\nModeled twofold-local speedup and zero-cost limit: 1.1764705882352942 1.4285714285714286\nPASS: kernels-04\n", "url": "executions/kernels-04.json" }, "url": "lesson.html?lesson=kernels-04" } ] }, { "id": "internals", "number": "14", "title": "Autodiff & JAX internals", "part": "Specializations", "prerequisitePhaseIds": [ "optimization" ], "hardware": "CPU", "project": "Audit derivatives and a tiny transformation", "checkpoint": "Explain tracing and verify the derivative rule numerically.", "source": "https://docs.jax.dev/en/latest/", "description": "Trace functions into jaxpr, check forward and reverse sensitivities, write exact custom derivatives and compile a small transformation with explicit supported primitives.", "learningAdvice": "Keep an independent derivative reference and test changed values, higher derivatives and unsupported cases before trusting a transformation.", "projectId": "derivative-audit", "studyGuide": { "question": "Which mathematical operation does the transformed program represent?", "why": "Use internal representations to explain and verify transformations, not to memorize printed variable names. Separate primal values, tangent directions, cotangents and compiled artifacts.", "readiness": { "prompt": "A function maps an input vector to an output vector. Does a JVP multiply by the Jacobian or its transpose?", "answer": "A JVP maps an input tangent through the Jacobian. A VJP maps an output cotangent through its transpose. Check the spaces and the adjoint inner-product identity.", "lessonId": "internals-02" }, "milestones": [ { "lessonIds": [ "internals-01", "internals-02" ], "title": "Read and check transformed programs", "task": "Follow one jaxpr equation, compare directional derivatives independently and verify the adjoint identity at changed inputs." }, { "lessonIds": [ "internals-03", "internals-04" ], "title": "Enforce a custom transformation boundary", "task": "Require tangent linearity and test higher-order promises. For the tiny interpreter, reject unsupported primitives rather than returning a plausible incomplete result." } ], "transfer": { "prompt": "A custom JVP returns the correct derivative when the tangent equals one but squares the tangent internally. Is it valid?", "approach": "No. Derivatives are linear maps in the direction at a fixed base point. Check multiple directions and linear combinations; agreement for one seed is insufficient." }, "pitfall": { "symptom": "A derivative test passes at one input but fails after a shift.", "check": "Inspect captured constants, detached values and custom rules. Compare both primal and derivative behavior at changed base points." }, "completion": "Use derivative-audit. Keep analytic/directional checks, the adjoint identity, supported higher-order behavior and explicit unsupported-program rejections.", "furtherWork": "The tiny interpreter is intentionally incomplete. Production extension APIs, effects and full compiler transformation coverage require further targeted work." }, "projectNote": "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.", "lessons": [ { "id": "internals-01", "title": "Read your first jaxpr", "status": "authored", "objective": "Inspect a small transformed function and identify its operations.", "prerequisites": [], "exercise": "Change the scalar offset from 1 to 2. Draw the new value flow for [1, 2, 3], calculate the sum and gradient by hand, then inspect both jaxprs. Explain which structural features remain and which values change.", "evidence": "Save your annotated forward and derivative traces, offset-2 hand calculation, captured/explicit input comparison, unrolled/scan comparison and traced-branch error with both repaired numeric tests.", "check": "A jaxpr has more equations after grad. What can you conclude from that observation alone?", "path": "phases/14-internals/01-read-your-first-jaxpr", "artifacts": { "source": "phases/14-internals/01-read-your-first-jaxpr/docs/en.md", "script": "exercises/internals-01.py", "notebook": "exercises/internals-01.ipynb", "markdown": "exercises/internals-01.md", "scriptSource": "phases/14-internals/01-read-your-first-jaxpr/code/main.py", "notebookSource": "phases/14-internals/01-read-your-first-jaxpr/notebooks/exercise.ipynb", "quizSource": "phases/14-internals/01-read-your-first-jaxpr/quiz.json", "evidenceSource": "phases/14-internals/01-read-your-first-jaxpr/outputs/evidence.md", "evidence": "exercises/internals-01-evidence.md" }, "runtime": { "kind": "cpu" }, "hardware": "CPU", "minutes": 90, "contentHash": "4605e7019f566f40efc6dba64e28ee6d18610e761aabe9e5eb3e84c46d3896c2", "visualArtifact": { "image": "figures/internals-01.svg", "png": "figures/internals-01.png", "data": "figures/internals-01.json", "generatedAt": "2026-10-20T10:46:35.063790+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "4605e7019f566f40efc6dba64e28ee6d18610e761aabe9e5eb3e84c46d3896c2", "kind": "executed" }, "execution": { "contentHash": "4605e7019f566f40efc6dba64e28ee6d18610e761aabe9e5eb3e84c46d3896c2", "executedAt": "2026-10-20T10:50:05.982947+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Forward program:\n{ lambda ; a:f32[3]. let\n b:f32[3] = add a 1.0:f32[]\n c:f32[3] = mul b b\n d:f32[] = reduce_sum[axes=(0,) out_sharding=None] c\n in (d,) }\nOperation: add output shapes: [(3,)]\nOperation: mul output shapes: [(3,)]\nOperation: reduce_sum output shapes: [()]\nValue / gradient: 29.0 [4. 6. 8.]\nForward program:\n{ lambda ; a:f32[3]. let\n b:f32[3] = add a 1.0:f32[]\n c:f32[3] = mul b b\n d:f32[] = reduce_sum[axes=(0,) out_sharding=None] c\n in (d,) }\nOperation: add output shapes: [(3,)]\nOperation: mul output shapes: [(3,)]\nOperation: reduce_sum output shapes: [()]\nValue / gradient: 29.0 [4. 6. 8.]\nDerivative program: { lambda ; a:f32[3]. let\n b:f32[3] = add a 1.0:f32[]\n c:f32[3] = mul b b\n _:f32[] = reduce_sum[axes=(0,) out_sharding=None] c\n d:f32[3] = broadcast_in_dim 1.0:f32[]\n e:f32[3] = mul b d\n f:f32[3] = mul d b\n g:f32[3] = add_any e f\n in (g,) }\nCaptured: { lambda a:f32[3]; b:f32[3]. let\n c:f32[3] = add b a\n d:f32[] = reduce_sum[axes=(0,) out_sharding=None] c\n in (d,) }\nExplicit: { lambda ; a:f32[3] b:f32[3]. let\n c:f32[3] = add a b\n d:f32[] = reduce_sum[axes=(0,) out_sharding=None] c\n in (d,) }\n{ lambda ; a:f32[3]. let\n b:f32[3] = add a 2.0:f32[]\n c:f32[3] = mul b b\n d:f32[] = reduce_sum[axes=(0,) out_sharding=None] c\n in (d,) }\nUnrolled: { lambda ; a:f32[]. let\n b:f32[] = mul 0.5:f32[] a\n c:f32[] = add b 1.0:f32[]\n d:f32[] = mul 0.5:f32[] c\n e:f32[] = add d 1.0:f32[]\n f:f32[] = mul 0.5:f32[] e\n g:f32[] = add f 1.0:f32[]\n h:f32[] = mul 0.5:f32[] g\n i:f32[] = add h 1.0:f32[]\n in (i,) }\nScanned: { lambda ; a:f32[]. let\n b:f32[] _:f32[4] = scan[\n _split_transpose=False\n jaxpr={ lambda ; c:f32[]. let\n d:f32[] = mul 0.5:f32[] c\n e:f32[] = add d 1.0:f32[]\n in (e, e) }\n length=4\n linear=(False,)\n num_carry=1\n num_consts=0\n reverse=False\n unroll=1\n ] a\n in (b,) }\nExpected failure: traced boolean used by Python if\nRepaired: { lambda ; a:f32[3]. let\n b:f32[1] = slice[limit_indices=(1,) start_indices=(0,) strides=None] a\n c:f32[] = squeeze[dimensions=(0,)] b\n d:bool[] = gt c 0.0:f32[]\n e:i32[] = convert_element_type[new_dtype=int32 weak_type=False] d\n f:f32[] = cond[\n branches=(\n { lambda ; g:f32[3]. let\n h:f32[] = reduce_sum[axes=(0,) out_sharding=None] g\n i:f32[] = neg h\n in (i,) }\n { lambda ; j:f32[3]. let\n k:f32[] = reduce_sum[axes=(0,) out_sharding=None] j\n in (k,) }\n )\n ] e a\n in (f,) }\nPASS: internals-01\n", "url": "executions/internals-01.json" }, "url": "lesson.html?lesson=internals-01" }, { "id": "internals-02", "title": "JVPs, VJPs, and higher-order derivatives", "status": "authored", "objective": "Compute and interpret JVP and VJP products with correct domains and shapes. Verify the adjoint dot-product identity against a hand-derived Jacobian.", "prerequisites": [ "internals-01" ], "exercise": "Repeat the Jacobian, JVP, VJP and adjoint checks at \\(x=(-0.2,0.8)\\) with \\(v=(-0.5,1.3)\\) and \\(u=(1.1,-0.4)\\). Explain which quantities belong to input and output space.", "evidence": "Hand-derived Jacobian/Hessian, adjoint scalar, changed-point checks, finite-difference HVP and a Gauss–Newton distinction. Keep the environment, observed outputs and your explanation of the figure.", "check": "A VJP is seeded with an output weighting u. What does the returned input cotangent represent?", "path": "phases/14-internals/02-jvps-vjps-and-higher-order-derivatives", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/14-internals/02-jvps-vjps-and-higher-order-derivatives/docs/en.md", "script": "exercises/internals-02.py", "notebook": "exercises/internals-02.ipynb", "markdown": "exercises/internals-02.md", "scriptSource": "phases/14-internals/02-jvps-vjps-and-higher-order-derivatives/code/main.py", "notebookSource": "phases/14-internals/02-jvps-vjps-and-higher-order-derivatives/notebooks/exercise.ipynb", "quizSource": "phases/14-internals/02-jvps-vjps-and-higher-order-derivatives/quiz.json", "evidenceSource": "phases/14-internals/02-jvps-vjps-and-higher-order-derivatives/outputs/evidence.md", "evidence": "exercises/internals-02-evidence.md" }, "minutes": 105, "contentHash": "b4c3b8720c5f59e80a50216fbd223c61ea013199a86108a3d73205efb94fc75a", "visualArtifact": { "image": "figures/internals-02.svg", "png": "figures/internals-02.png", "data": "figures/internals-02.json", "generatedAt": "2026-10-20T10:46:36.318127+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "b4c3b8720c5f59e80a50216fbd223c61ea013199a86108a3d73205efb94fc75a", "kind": "executed", "mechanismImage": "figures/internals-02-mechanism.svg", "mechanismPng": "figures/internals-02-mechanism.png" }, "execution": { "contentHash": "b4c3b8720c5f59e80a50216fbd223c61ea013199a86108a3d73205efb94fc75a", "executedAt": "2026-10-20T10:50:09.725259+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Jacobian: [[-0.7 0.4 ]\n [ 0.92106099 -1.4 ]]\nJVP: [-0.96 1.52527319] VJP: [ 1.49212199 -2.6 ] adjoint pair: 2.570546385606924 2.570546385606924\nHessian-vector product: [ 1.74991568 -3.38303348]\nJacobian: [[-0.7 0.4 ]\n [ 0.92106099 -1.4 ]]\nJVP: [-0.96 1.52527319] VJP: [ 1.49212199 -2.6 ] adjoint pair: 2.570546385606924 2.570546385606924\nHessian-vector product: [ 1.74991568 -3.38303348]\nForward columns and reverse rows agree\nExact HVP: [ 1.74991568 -3.38303348] Gauss-Newton product: [ 2.07686964 -2.51938247]\nChanged-point adjoint verified\nDirection linearity holds; the Jacobian changes with the base point\nBilinear Hessian symmetry and finite-difference product verified\nPASS: internals-02\n", "url": "executions/internals-02.json" }, "url": "lesson.html?lesson=internals-02" }, { "id": "internals-03", "title": "Custom derivative rules", "status": "authored", "objective": "Implement actual custom_jvp and custom_vjp rules with correct return and residual contracts. Verify seed linearity, stable extreme values and second derivatives.", "prerequisites": [ "internals-02" ], "exercise": "Check the custom log-sum-exp gradient on \\((-3,0.2,2,5)\\) against an independent stable softmax and central differences along two directions. Keep the output cotangent explicit.", "evidence": "Stable extreme inputs, analytic slopes/Hessians, seed linearity, independent finite differences and a deliberately wrong-rule diagnosis. Keep the environment, observed outputs and your explanation of the figure.", "check": "A custom derivative is linear in its tangent and returns finite values. What additional evidence is necessary before calling it exact?", "path": "phases/14-internals/03-custom-derivative-rules", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/14-internals/03-custom-derivative-rules/docs/en.md", "script": "exercises/internals-03.py", "notebook": "exercises/internals-03.ipynb", "markdown": "exercises/internals-03.md", "scriptSource": "phases/14-internals/03-custom-derivative-rules/code/main.py", "notebookSource": "phases/14-internals/03-custom-derivative-rules/notebooks/exercise.ipynb", "quizSource": "phases/14-internals/03-custom-derivative-rules/quiz.json", "evidenceSource": "phases/14-internals/03-custom-derivative-rules/outputs/evidence.md", "evidence": "exercises/internals-03-evidence.md" }, "minutes": 120, "contentHash": "16c3fd20cd983fed1236cd2f8b1589b5edf49f3be77e784652c254b20578c51a", "visualArtifact": { "image": "figures/internals-03.svg", "png": "figures/internals-03.png", "data": "figures/internals-03.json", "generatedAt": "2026-10-20T10:46:38.018377+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "16c3fd20cd983fed1236cd2f8b1589b5edf49f3be77e784652c254b20578c51a", "kind": "executed" }, "execution": { "contentHash": "16c3fd20cd983fed1236cd2f8b1589b5edf49f3be77e784652c254b20578c51a", "executedAt": "2026-10-20T10:50:14.579922+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Softplus first/second: [0.01798621 0.5 0.95257413] [0.01766271 0.25 0.04517666]\nLogsumexp gradient: [0.09085944 0.30166396 0.60747661]\nStable extreme values: [0.00000000e+00 6.93147181e-01 1.00000000e+03] naive positive extreme: inf\nSoftplus first/second: [0.01798621 0.5 0.95257413] [0.01766271 0.25 0.04517666]\nLogsumexp gradient: [0.09085944 0.30166396 0.60747661]\nStable extreme values: [0.00000000e+00 6.93147181e-01 1.00000000e+03] naive positive extreme: inf\nWrong/correct/finite-difference derivative: 0.29934383005622606 0.5986876601124521 0.5986876601138391\nExpected custom_vjp forward-mode boundary: can't apply forward-mode autodiff (jvp) to a custom_vjp function.\nOrdinary JAX common-shift JVP: 1.0\nChanged-size gradient and cotangent scaling verified\nStable shift invariance and zero common-direction curvature verified\nNonlinear seed rule rejected by its contract; repaired JVP is linear\nPASS: internals-03\n", "url": "executions/internals-03.json" }, "url": "lesson.html?lesson=internals-03" }, { "id": "internals-04", "title": "Lowering, compilation, and a tiny transformation", "status": "authored", "objective": "Interpret literals, captured constants and variables in a supported flat jaxpr. Implement actual forward derivative rules without calling jax.jvp inside the interpreter.", "prerequisites": [ "internals-03" ], "exercise": "Trace a new polynomial using only supported primitives, verify its interpreted derivative at two points, and compile that transformed computation independently.", "evidence": "Annotated primitive rules, constant/literal handling, independent expression oracles, changed compiled values and explicit unsupported/signature errors. Keep the environment, observed outputs and your explanation of the figure.", "check": "Why does the tiny interpreter raise an error for an unsupported primitive instead of calling that primitive directly?", "path": "phases/14-internals/04-lowering-compilation-and-a-tiny-transformation", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/14-internals/04-lowering-compilation-and-a-tiny-transformation/docs/en.md", "script": "exercises/internals-04.py", "notebook": "exercises/internals-04.ipynb", "markdown": "exercises/internals-04.md", "scriptSource": "phases/14-internals/04-lowering-compilation-and-a-tiny-transformation/code/main.py", "notebookSource": "phases/14-internals/04-lowering-compilation-and-a-tiny-transformation/notebooks/exercise.ipynb", "quizSource": "phases/14-internals/04-lowering-compilation-and-a-tiny-transformation/quiz.json", "evidenceSource": "phases/14-internals/04-lowering-compilation-and-a-tiny-transformation/outputs/evidence.md", "evidence": "exercises/internals-04-evidence.md" }, "minutes": 130, "contentHash": "2143e38783dec4194ec9bb3705c2aee898aaefa4b683badd7305f6fac2a8e76b", "visualArtifact": { "image": "figures/internals-04.svg", "png": "figures/internals-04.png", "data": "figures/internals-04.json", "generatedAt": "2026-10-20T10:46:39.082535+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "2143e38783dec4194ec9bb3705c2aee898aaefa4b683badd7305f6fac2a8e76b", "kind": "executed" }, "execution": { "contentHash": "2143e38783dec4194ec9bb3705c2aee898aaefa4b683badd7305f6fac2a8e76b", "executedAt": "2026-10-20T10:50:16.810777+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Supported primitives: ['sin', 'mul', 'add', 'reduce_sum']\nInterpreted primal and tangent: 2.3099803057106576 0.4391291800455618\nLowered IR characters: 1421\nCompiled transformed outputs match; changed shape explicitly rejected\nSupported primitives: ['sin', 'mul', 'add', 'reduce_sum']\nInterpreted primal and tangent: 2.3099803057106576 0.4391291800455618\nLowered IR characters: 1421\nCompiled transformed outputs match; changed shape explicitly rejected\nExpected unsupported primitive: unsupported primitive: exp\nChanged-value compiled tangent: 0.5640324983393596\nChanged polynomial transformed and compiled\nExplicit offset tangent contribution: -0.09999999999999998\nMultiple output contract and malformed tangent rejection verified\nPASS: internals-04\n", "url": "executions/internals-04.json" }, "url": "lesson.html?lesson=internals-04" } ] }, { "id": "deployment", "number": "15", "title": "Deployment, interoperability & edge AI", "part": "Specializations", "prerequisitePhaseIds": [ "recovery", "performance" ], "hardware": "CPU for core labs; optional framework environments and target devices for conversion and edge validation", "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.", "source": "https://docs.jax.dev/en/latest/", "description": "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.", "learningAdvice": "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.", "projectId": "deployment-audit", "additionalProjectIds": [ "engineering-release", "weight-conversion" ], "studyGuide": { "question": "Does the exported service preserve the model’s complete contract?", "why": "A deployment includes preprocessing, parameter layout, numeric policy and runtime behavior. Verify the same raw examples across each boundary before comparing memory or speed.", "readiness": { "prompt": "Weights match after conversion, but predictions differ. What besides weights must be compared?", "answer": "Compare input normalization, feature/channel order, operation semantics, non-trainable state, output labels, dtype and shape. Start with a hand-computed asymmetric example to expose layout mistakes.", "lessonId": "deployment-01" }, "milestones": [ { "title": "Choose the model and preserve its meaning", "lessonIds": [ "deployment-02", "deployment-01", "deployment-08" ], "task": "Adapt the model under a fixed held-out protocol before freezing its release. Compare framework layouts, preprocessing and state with independent predictions. Convert the actual PyTorch model to Flax and locate an injected operation mismatch numerically. Keep a candidate decision table, feature/parameter mapping table, and the first divergent operation from a deliberately broken conversion." }, { "title": "Freeze the artifact and numeric policy", "lessonIds": [ "deployment-03", "deployment-05" ], "task": "Export and load in a fresh process, reject malformed inputs, then track weight, activation, accumulator, bias and output formats. Separate clipping from rounding and retain held-out conversion checks. Include a fresh-interpreter export receipt, raw-request rejection cases, clipping counts and actual versus ideal packed byte counts." }, { "title": "Deliver to a stated runtime", "lessonIds": [ "deployment-07", "deployment-04", "deployment-06" ], "task": "Package and qualify the selected service, relate queue assumptions to measured service work, then examine the additional conversion and device budgets for an edge target. Leave actual device fields unmeasured until tested. Keep separate artifact, process, image and device receipts. Explain an interpolated percentile and a burst deadline failure before choosing a capacity intervention." } ], "transfer": { "prompt": "An INT8 model matches typical inputs but fails after a sensor’s range changes. How do you locate the regression?", "approach": "Verify raw preprocessing first; count values outside the frozen calibration range, inspect clipped channels and compare float-converted versus quantized outputs. Do not recalibrate on final test labels. Re-evaluate a revised policy on a separate held-out set." }, "pitfall": { "symptom": "A smaller model file produces no lower request latency.", "check": "Inspect load-time expansion, fallback operators, transfers and preprocessing on the named runtime. File size is not a kernel or latency measurement." }, "completion": "Use deployment-audit and engineering-release to assemble one release dossier: candidate selection and retention criteria; a named feature and parameter mapping; source/export and layerwise parity; full precision policy and calibration identity; raw-request validation; artifact and process receipts; measured request samples; and a queue/deadline analysis. Mark container and target-device evidence unmeasured unless you actually ran those environments. Another learner should be able to reproduce each claimed check from the recorded inputs and command.", "furtherWork": "Actual edge-device delegates, thermal/energy budgets, packed low-bit kernels and production serving loads remain unqualified. The revised sequence moves adaptation before export and containerization before service capacity planning." }, "projectNote": "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.", "lessons": [ { "id": "deployment-02", "title": "Adapt a pretrained model and choose a post-training objective", "status": "authored", "objective": "Pretrain and reload an artifact with recorded data and checkpoint hashes Compare supervised adaptation and teacher distillation from the same starting weights", "prerequisites": [], "exercise": "Adapt for zero steps and verify that predictions exactly match the reloaded checkpoint. Then run a short adaptation and show that the saved base is unchanged while the adapted parameters differ.", "evidence": "Source/checkpoint hashes, independent loss and gradient, disjoint evaluation results and teacher-failure diagnosis. Keep the environment, observed outputs and your explanation of the figure.", "check": "Why should the two adaptation branches be compared using a common held-out criterion rather than their own training losses?", "path": "phases/15-deployment/02-adapt-a-pretrained-model-and-choose-a-post-training-objective", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/15-deployment/02-adapt-a-pretrained-model-and-choose-a-post-training-objective/docs/en.md", "script": "exercises/deployment-02.py", "notebook": "exercises/deployment-02.ipynb", "markdown": "exercises/deployment-02.md", "scriptSource": "phases/15-deployment/02-adapt-a-pretrained-model-and-choose-a-post-training-objective/code/main.py", "notebookSource": "phases/15-deployment/02-adapt-a-pretrained-model-and-choose-a-post-training-objective/notebooks/exercise.ipynb", "quizSource": "phases/15-deployment/02-adapt-a-pretrained-model-and-choose-a-post-training-objective/quiz.json", "evidenceSource": "phases/15-deployment/02-adapt-a-pretrained-model-and-choose-a-post-training-objective/outputs/evidence.md", "evidence": "exercises/deployment-02-evidence.md" }, "minutes": 130, "contentHash": "6db9af3ac8d9c5f8b63bdbeff90da292a562996376bbb32f458317ee9d879ee8", "visualArtifact": { "image": "figures/deployment-02.svg", "png": "figures/deployment-02.png", "data": "figures/deployment-02.json", "generatedAt": "2026-10-20T10:46:41.298818+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "6db9af3ac8d9c5f8b63bdbeff90da292a562996376bbb32f458317ee9d879ee8", "kind": "executed" }, "execution": { "contentHash": "6db9af3ac8d9c5f8b63bdbeff90da292a562996376bbb32f458317ee9d879ee8", "executedAt": "2026-10-20T10:50:24.147400+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "{\n \"pretrained\": {\n \"held_cross_entropy\": 0.9780182838439941,\n \"held_brier\": 0.17108705639839172,\n \"source_cross_entropy\": 0.5451400876045227\n },\n \"supervised\": {\n \"held_cross_entropy\": 0.5554870367050171,\n \"held_brier\": 0.012093935161828995,\n \"source_cross_entropy\": 0.7904530763626099\n },\n \"teacher\": {\n \"held_cross_entropy\": 0.5240296125411987,\n \"held_brier\": 0.00022982103109825402,\n \"source_cross_entropy\": 1.0046310424804688\n }\n}\ncheckpoint source hash: f5fb50c38195b46766bee2edaa81a7ad28386523baa850c5294eb71b47f1f84c\n{\n \"pretrained\": {\n \"held_cross_entropy\": 0.9780182838439941,\n \"held_brier\": 0.17108705639839172,\n \"source_cross_entropy\": 0.5451400876045227\n },\n \"supervised\": {\n \"held_cross_entropy\": 0.5554870367050171,\n \"held_brier\": 0.012093935161828995,\n \"source_cross_entropy\": 0.7904530763626099\n },\n \"teacher\": {\n \"held_cross_entropy\": 0.5240296125411987,\n \"held_brier\": 0.00022982103109825402,\n \"source_cross_entropy\": 1.0046310424804688\n }\n}\ncheckpoint source hash: f5fb50c38195b46766bee2edaa81a7ad28386523baa850c5294eb71b47f1f84c\nwrong-teacher held loss: 1.3151594400405884\nsource retention losses: [0.5451400876045227, 0.7904530763626099, 1.0046310424804688]\nsupervised source loss increase 0.24531298875808716 passes declared gate True\nteacher source loss increase 0.45949095487594604 passes declared gate False\nCandidate selected under the declared fixture budget: supervised\nPASS: deployment-02\n", "url": "executions/deployment-02.json" }, "url": "lesson.html?lesson=deployment-02" }, { "id": "deployment-01", "title": "Move models between JAX, Keras, TensorFlow and PyTorch", "status": "authored", "objective": "Verify parameter layouts and inference semantics before switching frameworks.", "prerequisites": [ "deployment-02" ], "exercise": "Add a batch with four rows and preserve the same parameters. Verify each row independently using a feature-by-feature sum.", "evidence": "Keep the layout map, two changed-input comparisons, runtime versions and optional backend-lab output. Identify which model state and preprocessing are not covered by a dense-layer match.", "check": "The Dense predictions match after transposing weights. What has been established?", "path": "phases/15-deployment/01-keras-and-pytorch-bridges-to-explicit-jax", "runtime": { "kind": "cpu" }, "hardware": "CPU; optional framework and edge-device labs have separate requirements", "artifacts": { "source": "phases/15-deployment/01-keras-and-pytorch-bridges-to-explicit-jax/docs/en.md", "script": "exercises/deployment-01.py", "notebook": "exercises/deployment-01.ipynb", "markdown": "exercises/deployment-01.md", "scriptSource": "phases/15-deployment/01-keras-and-pytorch-bridges-to-explicit-jax/code/main.py", "notebookSource": "phases/15-deployment/01-keras-and-pytorch-bridges-to-explicit-jax/notebooks/exercise.ipynb", "quizSource": "phases/15-deployment/01-keras-and-pytorch-bridges-to-explicit-jax/quiz.json", "evidenceSource": "phases/15-deployment/01-keras-and-pytorch-bridges-to-explicit-jax/outputs/evidence.md", "evidence": "exercises/deployment-01-evidence.md" }, "minutes": 110, "contentHash": "996b7ae3906cffdb498add59edb1b30dc571162f999c431c441f51793de81f2c", "visualArtifact": { "image": "figures/deployment-01.svg", "png": "figures/deployment-01.png", "data": "figures/deployment-01.json", "generatedAt": "2026-10-20T10:46:42.378114+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "996b7ae3906cffdb498add59edb1b30dc571162f999c431c441f51793de81f2c", "kind": "executed" }, "execution": { "contentHash": "996b7ae3906cffdb498add59edb1b30dc571162f999c431c441f51793de81f2c", "executedAt": "2026-10-20T10:50:26.277718+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Matched dense outputs: [[ 3.1 -0.45]\n [-4.9 4.55]]\nMatched dense outputs: [[ 3.1 -0.45]\n [-4.9 4.55]]\nEqual shapes, unequal outputs\nFeature order failed, then matched after aligning kernel rows.\nChanged batch verified\nTwo-layer reference scores: [ 6.5 -4.25]\nPASS: deployment-01\n", "url": "executions/deployment-01.json" }, "url": "lesson.html?lesson=deployment-01" }, { "id": "deployment-08", "title": "Convert PyTorch weights to Flax and locate numerical errors", "status": "authored", "objective": "Convert a real PyTorch state dictionary to Flax NNX and validate intermediate outputs, gradients and reloaded artifacts with explicit numerical tolerances.", "prerequisites": [ "deployment-01" ], "exercise": "Compare a near-zero reference against two candidate errors. Verify that the absolute tolerance accepts rounding noise but rejects a meaningful discrepancy.", "evidence": "Keep source checkpoint and architecture identity, complete parameter mapping, per-layer max-absolute/relative-L2 errors, elementwise tolerances, input-gradient comparisons, wrong-epsilon diagnosis and a fresh-process converted-artifact receipt.", "check": "The first linear layer agrees, but LayerNorm is the first mismatch. What should you inspect next?", "path": "phases/15-deployment/08-pytorch-to-flax-numerical-parity", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/15-deployment/08-pytorch-to-flax-numerical-parity/docs/en.md", "script": "exercises/deployment-08.py", "notebook": "exercises/deployment-08.ipynb", "markdown": "exercises/deployment-08.md", "scriptSource": "phases/15-deployment/08-pytorch-to-flax-numerical-parity/code/main.py", "notebookSource": "phases/15-deployment/08-pytorch-to-flax-numerical-parity/notebooks/exercise.ipynb", "quizSource": "phases/15-deployment/08-pytorch-to-flax-numerical-parity/quiz.json", "evidenceSource": "phases/15-deployment/08-pytorch-to-flax-numerical-parity/outputs/evidence.md", "evidence": "exercises/deployment-08-evidence.md" }, "minutes": 130, "contentHash": "ca01511bddec81ea2245eb5085cd368dc5d5c02708cf83929268b92134ae94f5", "visualArtifact": { "image": "figures/deployment-08.svg", "png": "figures/deployment-08.png", "data": "figures/deployment-08.json", "generatedAt": "2026-10-20T10:46:48.017044+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "ca01511bddec81ea2245eb5085cd368dc5d5c02708cf83929268b92134ae94f5", "kind": "executed", "mechanismImage": "figures/deployment-08-mechanism.svg", "mechanismPng": "figures/deployment-08-mechanism.png" }, "execution": { "contentHash": "ca01511bddec81ea2245eb5085cd368dc5d5c02708cf83929268b92134ae94f5", "executedAt": "2026-10-20T10:50:37.934658+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Maximum absolute error per layer: {'hidden': 2.384185791015625e-07, 'norm': 3.5762786865234375e-07, 'activation': 3.5762786865234375e-07, 'output': 1.1920928955078125e-07}\nWrong epsilon: first mismatch is norm; intermediate outputs and input gradients verified on CPU.\nMaximum absolute error per layer: {'hidden': 2.384185791015625e-07, 'norm': 3.5762786865234375e-07, 'activation': 3.5762786865234375e-07, 'output': 1.1920928955078125e-07}\nWrong epsilon: first mismatch is norm; intermediate outputs and input gradients verified on CPU.\nPer-layer wrong-epsilon report: {'hidden': {'max_abs': 2.384185791015625e-07, 'relative_l2': 3.5662515782829134e-08, 'passed': True}, 'norm': {'max_abs': 0.046508073806762695, 'relative_l2': 0.018277932316979384, 'passed': False}, 'activation': {'max_abs': 0.05144989490509033, 'relative_l2': 0.022399576976497533, 'passed': False}, 'output': {'max_abs': 0.011958837509155273, 'relative_l2': 0.013684121610437203, 'passed': False}}\nNear-zero, excessive near-zero, large-value: True False True\nNear-zero acceptance and rejection verified.\nWrong source layout rejected\nNonuniform cotangent input-gradient parity: {'max_abs': 2.086162567138672e-07, 'relative_l2': 4.910744924550249e-07, 'passed': True}\nPASS: deployment-08\n", "url": "executions/deployment-08.json" }, "url": "lesson.html?lesson=deployment-08" }, { "id": "deployment-03", "title": "Export a computation and verify its serving contract", "status": "authored", "objective": "Round-trip an exported inference function and reject incompatible requests.", "prerequisites": [ "deployment-08" ], "exercise": "Verify a new input [-2, 0, 3] against an independent NumPy expression.", "evidence": "Keep the export/reload code, version/platform manifest, known-output and changed-input checks, rejected-shape result, and the planned serving request boundary.", "check": "What does a successful local export round trip prove?", "path": "phases/15-deployment/03-export-and-serve-a-trained-computation", "runtime": { "kind": "cpu" }, "hardware": "CPU; optional framework and edge-device labs have separate requirements", "artifacts": { "source": "phases/15-deployment/03-export-and-serve-a-trained-computation/docs/en.md", "script": "exercises/deployment-03.py", "notebook": "exercises/deployment-03.ipynb", "markdown": "exercises/deployment-03.md", "scriptSource": "phases/15-deployment/03-export-and-serve-a-trained-computation/code/main.py", "notebookSource": "phases/15-deployment/03-export-and-serve-a-trained-computation/notebooks/exercise.ipynb", "quizSource": "phases/15-deployment/03-export-and-serve-a-trained-computation/quiz.json", "evidenceSource": "phases/15-deployment/03-export-and-serve-a-trained-computation/outputs/evidence.md", "evidence": "exercises/deployment-03-evidence.md" }, "minutes": 110, "contentHash": "2bb69b19747a5a934d1051e96f9936cb9b6418e6a8a670ac3781c7f73f427b3f", "visualArtifact": { "image": "figures/deployment-03.svg", "png": "figures/deployment-03.png", "data": "figures/deployment-03.json", "generatedAt": "2026-10-20T10:46:49.818939+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "2bb69b19747a5a934d1051e96f9936cb9b6418e6a8a670ac3781c7f73f427b3f", "kind": "executed", "mechanismImage": "figures/deployment-03-mechanism.svg", "mechanismPng": "figures/deployment-03-mechanism.png" }, "execution": { "contentHash": "2bb69b19747a5a934d1051e96f9936cb9b6418e6a8a670ac3781c7f73f427b3f", "executedAt": "2026-10-20T10:50:40.025501+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Verified serialized bytes: 1172\nVerified serialized bytes: 1172\nRejected incompatible shape: ValueError\nFresh interpreter loaded bytes and matched the known scores.\nChanged exported request verified\nInvalid requests rejected before the one valid inference call.\nPASS: deployment-03\n", "url": "executions/deployment-03.json" }, "url": "lesson.html?lesson=deployment-03" }, { "id": "deployment-05", "title": "Choose weight, activation and accumulation precision", "status": "authored", "objective": "Compare floating-point policies, weight-only quantization and calibrated integer activations.", "exercise": "Add a zero-valued output channel and verify that its quantization scale is finite and its reconstruction is exactly zero.", "evidence": "Keep a precision table listing weight, activation, accumulator and output dtypes; max error and clipping counts on changed inputs; calibration provenance; and storage estimates that include scales and packing.", "check": "What must a deployment policy specify beyond “INT8 weights”?", "path": "phases/15-deployment/05-weight-activation-and-accumulation-precision", "runtime": { "kind": "cpu" }, "hardware": "CPU; optional framework and edge-device labs have separate requirements", "artifacts": { "source": "phases/15-deployment/05-weight-activation-and-accumulation-precision/docs/en.md", "script": "exercises/deployment-05.py", "notebook": "exercises/deployment-05.ipynb", "markdown": "exercises/deployment-05.md", "scriptSource": "phases/15-deployment/05-weight-activation-and-accumulation-precision/code/main.py", "notebookSource": "phases/15-deployment/05-weight-activation-and-accumulation-precision/notebooks/exercise.ipynb", "quizSource": "phases/15-deployment/05-weight-activation-and-accumulation-precision/quiz.json", "evidenceSource": "phases/15-deployment/05-weight-activation-and-accumulation-precision/outputs/evidence.md", "evidence": "exercises/deployment-05-evidence.md" }, "prerequisites": [ "deployment-03" ], "minutes": 110, "contentHash": "c229344c4a2c52d7669643c92b4b25cb997cf90ecb36600645fff53e8c8cbe68", "visualArtifact": { "image": "figures/deployment-05.svg", "png": "figures/deployment-05.png", "data": "figures/deployment-05.json", "generatedAt": "2026-10-20T10:46:51.082107+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "c229344c4a2c52d7669643c92b4b25cb997cf90ecb36600645fff53e8c8cbe68", "kind": "executed", "mechanismImage": "figures/deployment-05-mechanism.svg", "mechanismPng": "figures/deployment-05-mechanism.png" }, "execution": { "contentHash": "c229344c4a2c52d7669643c92b4b25cb997cf90ecb36600645fff53e8c8cbe68", "executedAt": "2026-10-20T10:50:41.243511+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": " max absolute error 0.0\n max absolute error 0.0010547637939453125\n max absolute error 0.006171703338623047\nW8A32 simulated output error 0.008031368255615234\nW4A32 simulated output error 0.14571428298950195\n max absolute error 0.0\n max absolute error 0.0010547637939453125\n max absolute error 0.006171703338623047\nW8A32 simulated output error 0.008031368255615234\nW4A32 simulated output error 0.14571428298950195\nAffine accumulators: [[5, 15]] output codes: [[2, 7]]\nCorrect / missing zero point / clipped: [[0.625, 0.9375]] [[-0.125, 0.5625]] [[0.52734375, 0.52734375]]\nClipped values: 1 reconstructed: [0.19685039 0.8976378 1. ]\nAffine accumulators: [[5, 15]] output codes: [[2, 7]]\nCorrect / missing zero point / clipped: [[0.625, 0.9375]] [[-0.125, 0.5625]] [[0.52734375, 0.52734375]]\nSmall-channel max error; shared / per-channel: 0.03 7.874053e-05\nZero channel handled\nW8A8 fixture error 0.00742650032043457\nCode-origin invariance and bias-only zero input verified.\nWeight-plus-scale subtotals: 24 14 11\nPASS: deployment-05\n", "url": "executions/deployment-05.json" }, "url": "lesson.html?lesson=deployment-05" }, { "id": "deployment-07", "title": "Containerize a model service and verify its boundary", "status": "authored", "objective": "Package an inference service and verify model identity, request validation and readiness across process and container boundaries.", "prerequisites": [ "deployment-05" ], "exercise": "Reject an artifact whose bytes changed after its expected digest was recorded.", "evidence": "Keep the artifact digest, valid and rejected requests, fresh-process output and, when Docker is available, the actual image identity, runtime user and container qualification receipt.", "check": "The image built successfully. What should happen before serving traffic?", "path": "phases/15-deployment/07-containerize-a-model-service", "runtime": { "kind": "cpu" }, "hardware": "CPU; Docker lab requires a local engine", "artifacts": { "source": "phases/15-deployment/07-containerize-a-model-service/docs/en.md", "script": "exercises/deployment-07.py", "notebook": "exercises/deployment-07.ipynb", "markdown": "exercises/deployment-07.md", "scriptSource": "phases/15-deployment/07-containerize-a-model-service/code/main.py", "notebookSource": "phases/15-deployment/07-containerize-a-model-service/notebooks/exercise.ipynb", "quizSource": "phases/15-deployment/07-containerize-a-model-service/quiz.json", "evidenceSource": "phases/15-deployment/07-containerize-a-model-service/outputs/evidence.md", "evidence": "exercises/deployment-07-evidence.md" }, "minutes": 110, "contentHash": "93a1f10568b65730bcac178c53015434c5c5d86e436c7219598fee4a97ea04ba", "visualArtifact": { "image": "figures/deployment-07.svg", "png": "figures/deployment-07.png", "data": "figures/deployment-07.json", "generatedAt": "2026-10-20T10:46:52.305883+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "93a1f10568b65730bcac178c53015434c5c5d86e436c7219598fee4a97ea04ba", "kind": "executed", "mechanismImage": "figures/deployment-07-mechanism.svg", "mechanismPng": "figures/deployment-07-mechanism.png" }, "execution": { "contentHash": "93a1f10568b65730bcac178c53015434c5c5d86e436c7219598fee4a97ea04ba", "executedAt": "2026-10-20T10:50:42.914381+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Fresh-process valid requests: 2 rejected boundaries: 3\nThis companion tests the process/artifact contract. Run the Docker lab for container evidence.\nFresh-process valid requests: 2 rejected boundaries: 3\nThis companion tests the process/artifact contract. Run the Docker lab for container evidence.\nIndependent three-input predictions: 0.5, 2.5, 5.0\nMatching digest did not bypass model validation.\nChanged artifact requires a new recorded digest.\nNonfinite input rejected before output.\nNew artifact loads, but the old behavioral expectation no longer passes.\nPASS: deployment-07\n", "url": "executions/deployment-07.json" }, "url": "lesson.html?lesson=deployment-07" }, { "id": "deployment-04", "title": "Inference capacity, batching, and autoscaling", "status": "authored", "objective": "Measure completed inference after warming each batch shape Report batch latency and example throughput with correct units", "prerequisites": [ "deployment-07" ], "exercise": "Simulate a burst of four requests arriving at time zero with a two-millisecond serial service duration. Derive all finish times and response times before executing.", "evidence": "Warm timing samples, latency/throughput units, independent queue timeline and labeled simulation/replica assumptions. Keep the environment, observed outputs and your explanation of the figure.", "check": "A batch of eight finishes in four milliseconds. Which statement follows for a fully occupied serial worker?", "path": "phases/15-deployment/04-inference-capacity-batching-and-autoscaling", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/15-deployment/04-inference-capacity-batching-and-autoscaling/docs/en.md", "script": "exercises/deployment-04.py", "notebook": "exercises/deployment-04.ipynb", "markdown": "exercises/deployment-04.md", "scriptSource": "phases/15-deployment/04-inference-capacity-batching-and-autoscaling/code/main.py", "notebookSource": "phases/15-deployment/04-inference-capacity-batching-and-autoscaling/notebooks/exercise.ipynb", "quizSource": "phases/15-deployment/04-inference-capacity-batching-and-autoscaling/quiz.json", "evidenceSource": "phases/15-deployment/04-inference-capacity-batching-and-autoscaling/outputs/evidence.md", "evidence": "exercises/deployment-04-evidence.md" }, "minutes": 130, "contentHash": "caf11b7c04e394c5eda2ab681d76adb0d11b51fbf9641b519177c99b2d175b3a", "visualArtifact": { "image": "figures/deployment-04.svg", "png": "figures/deployment-04.png", "data": "figures/deployment-04.json", "generatedAt": "2026-10-20T10:46:53.337823+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "caf11b7c04e394c5eda2ab681d76adb0d11b51fbf9641b519177c99b2d175b3a", "kind": "executed" }, "execution": { "contentHash": "caf11b7c04e394c5eda2ab681d76adb0d11b51fbf9641b519177c99b2d175b3a", "executedAt": "2026-10-20T10:50:44.185786+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "{'batch': 1, 'first_call_ms': 30.380667187273502, 'p50_ms': 0.009396113455295563, 'p95_ms': 0.02765634562820195, 'steady_examples_per_second': 106426.98225790464}\n{'batch': 8, 'first_call_ms': 19.219750072807074, 'p50_ms': 0.011957716196775436, 'p95_ms': 0.018120976164937012, 'steady_examples_per_second': 669024.0735231122}\n{'batch': 32, 'first_call_ms': 18.149458803236485, 'p50_ms': 0.01568719744682312, 'p95_ms': 0.03619375638663767, 'steady_examples_per_second': 2039879.979102351}\nmeasured batch-one service proxy (ms): 0.009396113455295563\nsimulated last stable/overloaded response (ms): 0.009396113455295563 0.23114439100027084\n{'batch': 1, 'first_call_ms': 18.95762514322996, 'p50_ms': 0.009458279237151146, 'p95_ms': 0.020428886637091616, 'steady_examples_per_second': 105727.47694655736}\n{'batch': 8, 'first_call_ms': 17.656249925494194, 'p50_ms': 0.012186821550130844, 'p95_ms': 0.01899753697216509, 'steady_examples_per_second': 656446.7992816477}\n{'batch': 32, 'first_call_ms': 17.30937510728836, 'p50_ms': 0.016249949112534523, 'p95_ms': 0.02756859175860881, 'steady_examples_per_second': 1969236.9359677904}\nmeasured batch-one service proxy (ms): 0.009458279237151146\nsimulated last stable/overloaded response (ms): 0.009458279237151146 0.2326736692339182\nhypothetical planned replicas: 3\nAnalytic sample p95 / max (ms): 5.95000000000007 100.0\nRequests beyond the declared deadline: 2\nPASS: deployment-04\n", "url": "executions/deployment-04.json" }, "url": "lesson.html?lesson=deployment-04" }, { "id": "deployment-06", "title": "Deploy at the edge: conversion, budgets and device checks", "status": "authored", "objective": "Define an edge runtime contract and measure an end-to-end batch-one request.", "exercise": "Send a second raw input [0, 255, 128] through the full request, and verify its independent normalized NumPy result.", "evidence": "Keep the CPU request report, preprocessing equivalence checks, chosen runtime/operator/precision matrix and an unfilled target-device report until actual device measurements exist. For the optional converter, retain the artifact, quantization metadata and held-out comparison.", "check": "A LiteRT file converts successfully. What must happen before claiming NPU latency?", "path": "phases/15-deployment/06-edge-ai-conversion-and-device-validation", "runtime": { "kind": "cpu" }, "hardware": "CPU; optional framework and edge-device labs have separate requirements", "artifacts": { "source": "phases/15-deployment/06-edge-ai-conversion-and-device-validation/docs/en.md", "script": "exercises/deployment-06.py", "notebook": "exercises/deployment-06.ipynb", "markdown": "exercises/deployment-06.md", "scriptSource": "phases/15-deployment/06-edge-ai-conversion-and-device-validation/code/main.py", "notebookSource": "phases/15-deployment/06-edge-ai-conversion-and-device-validation/notebooks/exercise.ipynb", "quizSource": "phases/15-deployment/06-edge-ai-conversion-and-device-validation/quiz.json", "evidenceSource": "phases/15-deployment/06-edge-ai-conversion-and-device-validation/outputs/evidence.md", "evidence": "exercises/deployment-06-evidence.md" }, "prerequisites": [ "deployment-04" ], "minutes": 110, "contentHash": "c543e431adadcd9167421861846fa832f48eb681f9d3d667ca20bec9ea696264", "visualArtifact": { "image": "figures/deployment-06.svg", "png": "figures/deployment-06.png", "data": "figures/deployment-06.json", "generatedAt": "2026-10-20T10:46:54.179503+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "c543e431adadcd9167421861846fa832f48eb681f9d3d667ca20bec9ea696264", "kind": "executed", "mechanismImage": "figures/deployment-06-mechanism.svg", "mechanismPng": "figures/deployment-06-mechanism.png" }, "execution": { "contentHash": "c543e431adadcd9167421861846fa832f48eb681f9d3d667ca20bec9ea696264", "executedAt": "2026-10-20T10:50:45.293154+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "{\n \"runtime\": \"JAX CPU instructional proxy\",\n \"jax\": \"0.9.2\",\n \"first_request_ms\": 11.448000092059374,\n \"warm_samples_ms\": [\n 0.15420792624354362,\n 0.059125013649463654,\n 0.04437519237399101,\n 0.03974977880716324,\n 0.037208665162324905,\n 0.03679096698760986,\n 0.033999793231487274,\n 0.03370875492691994,\n 0.032040756195783615,\n 0.032833777368068695,\n 0.03237510100007057,\n 0.030708033591508865,\n 0.030125025659799576,\n 0.029750168323516846,\n 0.03154110163450241,\n 0.029833056032657623,\n 0.029833056032657623,\n 0.029582995921373367,\n 0.028749927878379822,\n 0.032541342079639435,\n 0.029374845325946808,\n 0.029250048100948334,\n 0.02966681495308876,\n 0.0462094321846962,\n 0.029208138585090637,\n 0.029792077839374542,\n 0.02904096618294716,\n 0.028999987989664078,\n 0.02895807847380638,\n 0.02816738560795784\n ],\n \"p50_ms\": 0.03041652962565422,\n \"p95_ms\": 0.053313001990318264,\n \"boundary\": \"JSON decode + normalize + transfer + infer + wait + encode\",\n \"edge_device_validated\": false\n}\n{\n \"runtime\": \"JAX CPU instructional proxy\",\n \"jax\": \"0.9.2\",\n \"first_request_ms\": 9.076584130525589,\n \"warm_samples_ms\": [\n 0.13345805928111076,\n 0.06820773705840111,\n 0.06900029256939888,\n 0.042415689677000046,\n 0.035208649933338165,\n 0.032833777368068695,\n 0.031082890927791595,\n 0.03233412280678749,\n 0.02950010821223259,\n 0.02941582351922989,\n 0.03008311614394188,\n 0.028708018362522125,\n 0.028458889573812485,\n 0.028749927878379822,\n 0.028332695364952087,\n 0.028125010430812836,\n 0.028582755476236343,\n 0.02837507054209709,\n 0.028625130653381348,\n 0.028499867767095566,\n 0.029124785214662552,\n 0.028209295123815536,\n 0.02845795825123787,\n 0.027582980692386627,\n 0.03120768815279007,\n 0.028708018362522125,\n 0.027917325496673584,\n 0.027833040803670883,\n 0.028125010430812836,\n 0.028541311621665955\n ],\n \"p50_ms\": 0.028708018362522125,\n \"p95_ms\": 0.06864364258944988,\n \"boundary\": \"JSON decode + normalize + transfer + infer + wait + encode\",\n \"edge_device_validated\": false\n}\nPreprocessing mismatch max error 380.5019836425781\nSeven malformed or out-of-domain requests rejected; zero input returns the bias.\nChanged end-to-end request verified\nNot device-validated; missing: device, runtime_version, delegate, quality_metric, p95_ms, peak_memory_bytes, cold_start_ms, fallback_operators\nIllustrative request before/after/lower-bound (ms): 24.0 20.0 16.0\nPASS: deployment-06\n", "url": "executions/deployment-06.json" }, "url": "lesson.html?lesson=deployment-06" } ] }, { "id": "operations", "number": "16", "title": "Workload operations", "part": "Specializations", "prerequisitePhaseIds": [ "distributed" ], "hardware": "Executed local CPU processes; accelerator and cluster transfer require separate validation", "project": "An observable, restartable local workload", "checkpoint": "Diagnose a failed or stalled job, restore it correctly, and justify its resource plan from recorded measurements.", "source": "https://docs.jax.dev/en/latest/", "description": "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.", "learningAdvice": "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.", "projectId": "workload-operations", "additionalProjectIds": [ "engineering-release" ], "studyGuide": { "question": "Can you detect, explain and safely recover a failed model release?", "why": "Operational evidence connects work done to user-visible behavior. Separate worker health, model quality and release ownership so an alert leads to a specific investigation rather than automatic retraining.", "readiness": { "prompt": "A worker process is alive but has emitted no progress for a long time. Is it healthy?", "answer": "Liveness only establishes that the process exists or responds. Check readiness, last completed work, expected duration and checkpoint freshness before choosing a bounded timeout or recovery action.", "lessonId": "operations-02" }, "milestones": [ { "lessonIds": [ "operations-01", "operations-02", "operations-03" ], "title": "Observe and recover real work", "task": "Launch a bounded worker, retain structured events, rehearse a failure and verify the next restored update against uninterrupted work." }, { "lessonIds": [ "operations-04", "operations-05" ], "title": "Control change and capacity", "task": "Reject invalid artifacts without changing selection. Revalidate rollback targets and calculate demand, reserve and cost with consistent measured units." }, { "lessonIds": [ "operations-06", "operations-07", "operations-08" ], "title": "Connect MLOps, LLMOps and ModelOps", "task": "Use data and slice-quality gates, trace a versioned text application and bind review to the complete release bundle. Keep evaluator replay separate from live-model quality." } ], "transfer": { "prompt": "An input-distribution alert fires while delayed labels have not arrived. Should the system automatically retrain and promote?", "approach": "No conclusion about quality follows yet. Check data integrity, affected slices and service health; gather outcome evidence under a declared response policy. Any candidate must pass independent release gates, and selection must remain unchanged after failure." }, "pitfall": { "symptom": "A release passes average error while a rare critical slice fails.", "check": "Require per-slice counts and quality thresholds. Do not average slice means equally or treat a missing critical slice as a passing zero." }, "completion": "Use workload-operations and engineering-release. Keep a real failure/recovery timeline, failed-gate invariance, a rollback drill, run/model identity and named review responsibilities.", "furtherWork": "Authenticated review, multi-writer promotion, staged cloud rollout, SLO alerting and live-model abuse evaluation still require separate infrastructure and application-specific labs." }, "projectNote": "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.", "lessons": [ { "id": "operations-01", "title": "Accelerator jobs and runtime lifecycle", "status": "authored", "objective": "Launch and reap a real bounded subprocess with explicit runtime requirements. Distinguish successful completion, worker failure and supervisor timeout.", "prerequisites": [], "exercise": "Launch a changed job with \\(3\\) updates and batch size \\(5\\). Verify the observed progress steps and completion event rather than relying only on its exit code.", "evidence": "Keep actual child exit codes, bounded timeout/termination output and proof that children were reaped. Distinguish launch failure, failed preflight, slow work and a stalled job.", "check": "What must a timeout handler do beyond recording a timeout status?", "path": "phases/16-operations/01-accelerator-jobs-and-runtime-lifecycle", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/16-operations/01-accelerator-jobs-and-runtime-lifecycle/docs/en.md", "script": "exercises/operations-01.py", "notebook": "exercises/operations-01.ipynb", "markdown": "exercises/operations-01.md", "scriptSource": "phases/16-operations/01-accelerator-jobs-and-runtime-lifecycle/code/main.py", "notebookSource": "phases/16-operations/01-accelerator-jobs-and-runtime-lifecycle/notebooks/exercise.ipynb", "quizSource": "phases/16-operations/01-accelerator-jobs-and-runtime-lifecycle/quiz.json", "evidenceSource": "phases/16-operations/01-accelerator-jobs-and-runtime-lifecycle/outputs/evidence.md", "evidence": "exercises/operations-01-evidence.md" }, "minutes": 90, "contentHash": "92e64dddbfd55d0bb45eb5e333d18b463aba0cb4772ce7eebfae42bddb28c297", "visualArtifact": { "image": "figures/operations-01.svg", "png": "figures/operations-01.png", "data": "figures/operations-01.json", "generatedAt": "2026-10-20T10:47:02.024462+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "92e64dddbfd55d0bb45eb5e333d18b463aba0cb4772ce7eebfae42bddb28c297", "kind": "executed" }, "execution": { "contentHash": "92e64dddbfd55d0bb45eb5e333d18b463aba0cb4772ce7eebfae42bddb28c297", "executedAt": "2026-10-20T10:51:15.896016+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "statuses: completed failed timed_out\nactual exit codes: 0 23 -15\nstatuses: completed failed timed_out\nactual exit codes: 0 23 -15\ntimed-out worker events: ['started', 'ready', 'stall_injected']\nPASS: operations-01\n", "url": "executions/operations-01.json" }, "url": "lesson.html?lesson=operations-01" }, { "id": "operations-02", "title": "Health, logs, and workload observability", "status": "authored", "objective": "Build useful summaries from structured events emitted by a real worker. Define throughput with explicit numerator, denominator and synchronization boundary.", "prerequisites": [ "operations-01" ], "exercise": "Compute the fraction of supervisor wall time spent inside synchronized updates, and state precisely what that fraction excludes.", "evidence": "Keep raw structured events, observed work counts, synchronized timing units and checkpoint lag. Derive job and update throughput with different boundaries.", "check": "What does update duty fraction establish in this lab?", "path": "phases/16-operations/02-health-logs-and-workload-observability", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/16-operations/02-health-logs-and-workload-observability/docs/en.md", "script": "exercises/operations-02.py", "notebook": "exercises/operations-02.ipynb", "markdown": "exercises/operations-02.md", "scriptSource": "phases/16-operations/02-health-logs-and-workload-observability/code/main.py", "notebookSource": "phases/16-operations/02-health-logs-and-workload-observability/notebooks/exercise.ipynb", "quizSource": "phases/16-operations/02-health-logs-and-workload-observability/quiz.json", "evidenceSource": "phases/16-operations/02-health-logs-and-workload-observability/outputs/evidence.md", "evidence": "exercises/operations-02-evidence.md" }, "minutes": 85, "contentHash": "dd61186655e00859690d650c18acd842b02a696530e9711d97dcc136d41f25ab", "visualArtifact": { "image": "figures/operations-02.svg", "png": "figures/operations-02.png", "data": "figures/operations-02.json", "generatedAt": "2026-10-20T10:47:04.316701+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "dd61186655e00859690d650c18acd842b02a696530e9711d97dcc136d41f25ab", "kind": "executed" }, "execution": { "contentHash": "dd61186655e00859690d650c18acd842b02a696530e9711d97dcc136d41f25ab", "executedAt": "2026-10-20T10:51:22.258741+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "healthy: {'status': 'completed', 'updates': 12, 'examples': 96, 'last_step': 12, 'checkpoint_step': 12, 'uncheckpointed_updates': 0, 'update_s': 0.0003859163261950016, 'update_examples_per_s': 248758.58698834025, 'job_examples_per_s': 176.72946936655072, 'update_duty_fraction': 0.0007104457036284513}\nincident: {'status': 'failed', 'updates': 5, 'examples': 40, 'last_step': 5, 'checkpoint_step': 4, 'uncheckpointed_updates': 1, 'update_s': 0.00017587607726454735, 'update_examples_per_s': 227432.86421969283, 'job_examples_per_s': 69.44720752337886, 'update_duty_fraction': 0.0003053525609047209}\nhealthy: {'status': 'completed', 'updates': 12, 'examples': 96, 'last_step': 12, 'checkpoint_step': 12, 'uncheckpointed_updates': 0, 'update_s': 0.0003807488828897476, 'update_examples_per_s': 252134.68591527935, 'job_examples_per_s': 174.31857639191733, 'update_duty_fraction': 0.0006913708669599339}\nincident: {'status': 'failed', 'updates': 5, 'examples': 40, 'last_step': 5, 'checkpoint_step': 4, 'uncheckpointed_updates': 1, 'update_s': 0.00018524983897805214, 'update_examples_per_s': 215924.61413550316, 'job_examples_per_s': 70.63820657474201, 'update_duty_fraction': 0.0003271429098417335}\nexamples / synchronized seconds: 96 0.0003807488828897476\nmeasured update duty fraction, not hardware utilization: 0.0006913708669599339\nPASS: operations-02\n", "url": "executions/operations-02.json" }, "url": "lesson.html?lesson=operations-02" }, { "id": "operations-03", "title": "Failure recovery and operational runbooks", "status": "authored", "objective": "Rehearse a real worker failure and resume from an atomic complete-state checkpoint. Verify the next update and every later state against an uninterrupted reference.", "prerequisites": [ "operations-02" ], "exercise": "Repeat the recovery drill with seed \\(9\\), checkpoint cadence \\(3\\) and failure after update \\(5\\). Verify restoration at step \\(3\\) and equality with a matching uninterrupted run.", "evidence": "Keep the injected failure timeline, accepted checkpoint identity and full-state comparisons for every resumed update. Record the incompatible and corrupted restore rejections.", "check": "Which evidence most directly supports a correct resume?", "path": "phases/16-operations/03-failure-recovery-and-operational-runbooks", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/16-operations/03-failure-recovery-and-operational-runbooks/docs/en.md", "script": "exercises/operations-03.py", "notebook": "exercises/operations-03.ipynb", "markdown": "exercises/operations-03.md", "scriptSource": "phases/16-operations/03-failure-recovery-and-operational-runbooks/code/main.py", "notebookSource": "phases/16-operations/03-failure-recovery-and-operational-runbooks/notebooks/exercise.ipynb", "quizSource": "phases/16-operations/03-failure-recovery-and-operational-runbooks/quiz.json", "evidenceSource": "phases/16-operations/03-failure-recovery-and-operational-runbooks/outputs/evidence.md", "evidence": "exercises/operations-03-evidence.md" }, "minutes": 110, "contentHash": "adb65ce1ba3c1a6a5cb60827246e91ab261edabf44974dca0f9868aa5d2b86ed", "visualArtifact": { "image": "figures/operations-03.svg", "png": "figures/operations-03.png", "data": "figures/operations-03.json", "generatedAt": "2026-10-20T10:47:07.645578+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "adb65ce1ba3c1a6a5cb60827246e91ab261edabf44974dca0f9868aa5d2b86ed", "kind": "executed" }, "execution": { "contentHash": "adb65ce1ba3c1a6a5cb60827246e91ab261edabf44974dca0f9868aa5d2b86ed", "executedAt": "2026-10-20T10:51:36.597024+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "restored step: 2 replayed updates: [3, 4, 5, 6, 7, 8]\nall next-update state hashes match: True\nrestored step: 2 replayed updates: [3, 4, 5, 6, 7, 8]\nall next-update state hashes match: True\nPASS: operations-03\n", "url": "executions/operations-03.json" }, "url": "lesson.html?lesson=operations-03" }, { "id": "operations-04", "title": "Changes, rollback, and artifact provenance", "status": "authored", "objective": "Separate content identity, validation evidence and active selection. Reject corrupt or failed artifacts before a pointer change.", "prerequisites": [ "operations-03" ], "exercise": "Publish two compatible validated artifacts that differ in an explicit release note. Activate both in sequence, then rollback and verify the exact selected ID and content.", "evidence": "Keep source/configuration/data hashes, actual held-out candidate predictions, the rejected degraded release and unchanged active pointer, followed by verified rollback.", "check": "What does a matching artifact digest establish?", "path": "phases/16-operations/04-changes-rollback-and-artifact-provenance", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/16-operations/04-changes-rollback-and-artifact-provenance/docs/en.md", "script": "exercises/operations-04.py", "notebook": "exercises/operations-04.ipynb", "markdown": "exercises/operations-04.md", "scriptSource": "phases/16-operations/04-changes-rollback-and-artifact-provenance/code/main.py", "notebookSource": "phases/16-operations/04-changes-rollback-and-artifact-provenance/notebooks/exercise.ipynb", "quizSource": "phases/16-operations/04-changes-rollback-and-artifact-provenance/quiz.json", "evidenceSource": "phases/16-operations/04-changes-rollback-and-artifact-provenance/outputs/evidence.md", "evidence": "exercises/operations-04-evidence.md" }, "minutes": 95, "contentHash": "e4af015be92bf0b2b6ccb5343222cd5de15b9dbdef7fbc0df0a5e892c717717c", "visualArtifact": { "image": "figures/operations-04.svg", "png": "figures/operations-04.png", "data": "figures/operations-04.json", "generatedAt": "2026-10-20T10:47:09.407572+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "e4af015be92bf0b2b6ccb5343222cd5de15b9dbdef7fbc0df0a5e892c717717c", "kind": "executed", "mechanismImage": "figures/operations-04-mechanism.svg", "mechanismPng": "figures/operations-04-mechanism.png" }, "execution": { "contentHash": "e4af015be92bf0b2b6ccb5343222cd5de15b9dbdef7fbc0df0a5e892c717717c", "executedAt": "2026-10-20T10:51:39.679728+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "first / second content IDs: 810b956e3376228ff3c1f03f469feb3ab26653b488ba51271e82277672590b80 aed1378a7842598237e59c5d189b13d486068b6284d9c1ca3d4e865866ba6990\nfailed gate kept selection: True rollback restored first: True\nfirst / second content IDs: 810b956e3376228ff3c1f03f469feb3ab26653b488ba51271e82277672590b80 aed1378a7842598237e59c5d189b13d486068b6284d9c1ca3d4e865866ba6990\nfailed gate kept selection: True rollback restored first: True\nPASS: operations-04\n", "url": "executions/operations-04.json" }, "url": "lesson.html?lesson=operations-04" }, { "id": "operations-05", "title": "Capacity, utilization, and operating cost", "status": "authored", "objective": "Measure end-to-end local service time at more than one workload length. Compute offered load, worker reserve and illustrative cost with consistent units.", "prerequisites": [ "operations-04" ], "exercise": "Using the measured service-time input, calculate the minimum worker count for \\(1200\\) arrivals per hour with \\(40\\)% reserve, then independently verify that the chosen count satisfies the load bound.", "evidence": "Keep measured job/update samples and independent unit calculations. Label demand, reserve and price assumptions as hypothetical before deriving a resource plan.", "check": "Which claim is supported by the measured update duty fraction?", "path": "phases/16-operations/05-capacity-utilization-and-operating-cost", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/16-operations/05-capacity-utilization-and-operating-cost/docs/en.md", "script": "exercises/operations-05.py", "notebook": "exercises/operations-05.ipynb", "markdown": "exercises/operations-05.md", "scriptSource": "phases/16-operations/05-capacity-utilization-and-operating-cost/code/main.py", "notebookSource": "phases/16-operations/05-capacity-utilization-and-operating-cost/notebooks/exercise.ipynb", "quizSource": "phases/16-operations/05-capacity-utilization-and-operating-cost/quiz.json", "evidenceSource": "phases/16-operations/05-capacity-utilization-and-operating-cost/outputs/evidence.md", "evidence": "exercises/operations-05-evidence.md" }, "minutes": 90, "contentHash": "bd985f582fbaea197458e4dfca2d3965c26d712bb05b83aca5893f9c7b19969e", "visualArtifact": { "image": "figures/operations-05.svg", "png": "figures/operations-05.png", "data": "figures/operations-05.json", "generatedAt": "2026-10-20T10:47:12.783181+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "bd985f582fbaea197458e4dfca2d3965c26d712bb05b83aca5893f9c7b19969e", "kind": "executed" }, "execution": { "contentHash": "bd985f582fbaea197458e4dfca2d3965c26d712bb05b83aca5893f9c7b19969e", "executedAt": "2026-10-20T10:51:47.444738+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "measured local job seconds: [0.543132500257343, 0.5599589156918228, 0.5557856247760355]\nillustrative capacity plan: {'offered_worker_hours_per_hour': 0.09332648594863713, 'load_fraction': 0.046663242974318564, 'minimum_workers_with_reserve': 1, 'hourly_budget': 3.0, 'hypothetical_cost_per_job': 0.0002333162148715928, 'within_reserve': True}\nassumptions: 600 arrivals/hour, 2 workers, hypothetical 1.50 per worker-hour, 25% reserve\nmeasured local job seconds: [0.6144686662591994, 0.5638747080229223, 0.5770753328688443]\nillustrative capacity plan: {'offered_worker_hours_per_hour': 0.10241144437653323, 'load_fraction': 0.05120572218826661, 'minimum_workers_with_reserve': 1, 'hourly_budget': 3.0, 'hypothetical_cost_per_job': 0.00025602861094133306, 'within_reserve': True}\nassumptions: 600 arrivals/hour, 2 workers, hypothetical 1.50 per worker-hour, 25% reserve\nnominal load at two demand scenarios: 0.05120572218826661 0.10241144437653323\nrequired workers under changed hypothetical reserve: 1\nPASS: operations-05\n", "url": "executions/operations-05.json" }, "url": "lesson.html?lesson=operations-05" }, { "id": "operations-06", "title": "MLOps: data contracts, CI gates and monitoring", "status": "authored", "objective": "Reject data leakage and incomplete evaluation, compute count-weighted slice metrics and distinguish drift investigation from model promotion.", "prerequisites": [ "operations-05" ], "exercise": "Show that averaging slice means without counts changes the overall error.", "evidence": "Keep the leaking-group rejection, data fingerprint, slice counts and MSEs, aggregate/slice gate disagreement, changed-population calculation and a written response to drift without labels.", "check": "An input drift alert fired before new labels arrived. What follows?", "path": "phases/16-operations/06-mlops-data-contracts-and-release-pipelines", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/16-operations/06-mlops-data-contracts-and-release-pipelines/docs/en.md", "script": "exercises/operations-06.py", "notebook": "exercises/operations-06.ipynb", "markdown": "exercises/operations-06.md", "scriptSource": "phases/16-operations/06-mlops-data-contracts-and-release-pipelines/code/main.py", "notebookSource": "phases/16-operations/06-mlops-data-contracts-and-release-pipelines/notebooks/exercise.ipynb", "quizSource": "phases/16-operations/06-mlops-data-contracts-and-release-pipelines/quiz.json", "evidenceSource": "phases/16-operations/06-mlops-data-contracts-and-release-pipelines/outputs/evidence.md", "evidence": "exercises/operations-06-evidence.md" }, "minutes": 70, "contentHash": "28078c290e10a6e46e86fb9fd10cee8c44f9b8df69443bbdb51da4f712a9af8c", "visualArtifact": { "image": "figures/operations-06.svg", "png": "figures/operations-06.png", "data": "figures/operations-06.json", "generatedAt": "2026-10-20T10:47:13.651979+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "28078c290e10a6e46e86fb9fd10cee8c44f9b8df69443bbdb51da4f712a9af8c", "kind": "executed" }, "execution": { "contentHash": "28078c290e10a6e46e86fb9fd10cee8c44f9b8df69443bbdb51da4f712a9af8c", "executedAt": "2026-10-20T10:51:48.399230+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Overall MSE: 0.409 high-slice MSE: 4.0 release: False\nTwo-bin input total variation: 0.4 (requires investigation; not an automatic retraining order)\nOverall MSE: 0.409 high-slice MSE: 4.0 release: False\nTwo-bin input total variation: 0.4 (requires investigation; not an automatic retraining order)\nAggregate-only acceptance would miss the high-slice failure.\nUnweighted slice average: 2.005\nNew population MSE: 3.601\nPASS: operations-06\n", "url": "executions/operations-06.json" }, "url": "lesson.html?lesson=operations-06" }, { "id": "operations-07", "title": "LLMOps: version prompts, evaluate behavior and trace requests", "status": "authored", "objective": "Evaluate a versioned text application with a complete case inventory, slice-level answer contracts and inspectable MLflow traces.", "prerequisites": [ "operations-06" ], "exercise": "Add a response with an invented citation and require the evaluator to fail it.", "evidence": "Keep all case IDs and per-slice replay scores, missing/duplicate-response rejections, the retained failed case, trace identities and application-version fields. Explain what requires separate live-model evaluation.", "check": "The controlled replay suite passes every case. What is established?", "path": "phases/16-operations/07-llmops-prompts-evaluation-and-traces", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/16-operations/07-llmops-prompts-evaluation-and-traces/docs/en.md", "script": "exercises/operations-07.py", "notebook": "exercises/operations-07.ipynb", "markdown": "exercises/operations-07.md", "scriptSource": "phases/16-operations/07-llmops-prompts-evaluation-and-traces/code/main.py", "notebookSource": "phases/16-operations/07-llmops-prompts-evaluation-and-traces/notebooks/exercise.ipynb", "quizSource": "phases/16-operations/07-llmops-prompts-evaluation-and-traces/quiz.json", "evidenceSource": "phases/16-operations/07-llmops-prompts-evaluation-and-traces/outputs/evidence.md", "evidence": "exercises/operations-07-evidence.md" }, "minutes": 75, "contentHash": "43e1f34e84f93a58fe344a5357aee212ede59116f582807c39c167b4ae0f84a7", "visualArtifact": { "image": "figures/operations-07.svg", "png": "figures/operations-07.png", "data": "figures/operations-07.json", "generatedAt": "2026-10-20T10:47:17.120280+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "43e1f34e84f93a58fe344a5357aee212ede59116f582807c39c167b4ae0f84a7", "kind": "executed" }, "execution": { "contentHash": "43e1f34e84f93a58fe344a5357aee212ede59116f582807c39c167b4ae0f84a7", "executedAt": "2026-10-20T10:51:53.818799+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Observed replay scores: [{'overall': 0.9, 'answerable': 1.0, 'unanswerable': 0.0}, {'overall': 1.0, 'answerable': 1.0, 'unanswerable': 1.0}]\nTwo real MLflow traces; no LLM API was called and no general model quality was evaluated.\nObserved replay scores: [{'overall': 0.9, 'answerable': 1.0, 'unanswerable': 0.0}, {'overall': 1.0, 'answerable': 1.0, 'unanswerable': 1.0}]\nTwo real MLflow traces; no LLM API was called and no general model quality was evaluated.\nRelease decisions: [False, True]\nComplete inventory: {'count': 10, 'passed': 9, 'fraction': 0.9} ; missing, duplicate and unknown responses rejected.\nUnrecognized citations rejected.\nAn abstention flag alone does not satisfy the output contract.\nFailed answerable work remains in denominator: {'count': 10, 'passed': 9, 'fraction': 0.9}\nPASS: operations-07\n", "url": "executions/operations-07.json" }, "url": "lesson.html?lesson=operations-07" }, { "id": "operations-08", "title": "ModelOps: own, approve, roll back and retire a model", "status": "authored", "objective": "Bind a release decision to immutable evidence, reject changed or expired approvals and preserve active selection after failure.", "prerequisites": [ "operations-07" ], "exercise": "Change the evaluation hash while keeping the model unchanged, and prove the old approval no longer applies.", "evidence": "Keep the reviewed bundle hash, owner and target, expiry boundary, failed-transition pointer equality and changed-owner rejection. Explain why the local reviewer field is not authenticated approval.", "check": "An approval record contains a reviewer name and a matching hash. Does that authenticate the reviewer?", "path": "phases/16-operations/08-modelops-ownership-approval-and-retirement", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/16-operations/08-modelops-ownership-approval-and-retirement/docs/en.md", "script": "exercises/operations-08.py", "notebook": "exercises/operations-08.ipynb", "markdown": "exercises/operations-08.md", "scriptSource": "phases/16-operations/08-modelops-ownership-approval-and-retirement/code/main.py", "notebookSource": "phases/16-operations/08-modelops-ownership-approval-and-retirement/notebooks/exercise.ipynb", "quizSource": "phases/16-operations/08-modelops-ownership-approval-and-retirement/quiz.json", "evidenceSource": "phases/16-operations/08-modelops-ownership-approval-and-retirement/outputs/evidence.md", "evidence": "exercises/operations-08-evidence.md" }, "minutes": 65, "contentHash": "5a242aea00f1dd2f24dbfcb2bc3b6cb306416e79defc37b721a8536ad40ab193", "visualArtifact": { "image": "figures/operations-08.svg", "png": "figures/operations-08.png", "data": "figures/operations-08.json", "generatedAt": "2026-10-20T10:47:18.557826+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "5a242aea00f1dd2f24dbfcb2bc3b6cb306416e79defc37b721a8536ad40ab193", "kind": "executed" }, "execution": { "contentHash": "5a242aea00f1dd2f24dbfcb2bc3b6cb306416e79defc37b721a8536ad40ab193", "executedAt": "2026-10-20T10:51:54.787441+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Accepted original / changed artifact / expired approval / wrong target: [True, False, False, False]\nAll failed gates preserved the selected version. This fixture does not authenticate reviewers.\nAccepted original / changed artifact / expired approval / wrong target: [True, False, False, False]\nAccepted original / changed artifact / expired approval / wrong target: [True, False, False, False]\nAll failed gates preserved the selected version. This fixture does not authenticate reviewers.\nExpiry is exclusive: 199 accepted, 200 rejected.\nChanged evaluation invalidates the old approval.\nOwnership change requires renewed review.\nPASS: operations-08\n", "url": "executions/operations-08.json" }, "url": "lesson.html?lesson=operations-08" } ] }, { "id": "pretraining", "number": "17", "title": "Self-supervised pretraining: masked and contrastive learning", "part": "Training systems", "prerequisitePhaseIds": [ "transformers" ], "hardware": "Executed CPU synthetic objectives; large-scale and human-feedback evaluation remain separate", "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.", "source": "https://docs.jax.dev/en/latest/", "description": "Train small masked-token, hidden-patch and paired-view models. Audit information leakage, reduction counts and pair identity before interpreting a falling objective.", "learningAdvice": "Predict the mask, denominator or preference direction before training. Keep independent objective checks and changed-condition evidence.", "projectId": "training-methods", "projectStageIds": [ "1", "2", "3" ], "studyGuide": { "question": "What learning signal can unlabelled data provide?", "why": "Train small masked-token, hidden-patch and paired-view models. Audit information leakage, reduction counts and pair identity before interpreting a falling objective.", "readiness": { "prompt": "If an image reconstruction is perfect before training, what boundary should you inspect?", "answer": "Check whether hidden pixels entered the encoder or were copied into the prediction. A valid objective can still be defeated by leaked inputs.", "lessonId": "pretraining-02" }, "milestones": [ { "title": "Masked language modeling: predict hidden tokens", "task": "Trace clean targets, corrupted IDs, attention axes and selected counts. Reproduce unequal-mask loss with a host calculation, derive the full logit gradient, and explain the final probability matrix.", "lessonIds": [ "pretraining-01" ] }, { "title": "Masked image modeling: reconstruct missing patches", "task": "Verify every patch on a rectangular RGB input, intervene on hidden pixels, and explain an ambiguous target that no deterministic decoder can reconstruct perfectly. Compare actual hidden targets with predictions.", "lessonIds": [ "pretraining-02" ] }, { "title": "Contrastive learning: views, positives and negatives", "task": "Derive a two-pair loss, change temperature without changing ranking, compare encoder gradients with finite differences, and diagnose broken identity or duplicate false negatives from the similarity matrix.", "lessonIds": [ "pretraining-03" ] } ], "transfer": { "prompt": "A batch contains duplicate semantic pairs. Should each duplicate automatically be a negative?", "approach": "No. Decide positives by source or semantic identity and use a compatible objective; the duplicate-pair exercise shows how a diagonal-only loss penalizes duplicates." }, "pitfall": { "symptom": "A low masked loss does not transfer to the downstream task.", "check": "Check target leakage, mask distribution, encoder-versus-projection features and a held-out downstream protocol." }, "completion": "Complete stages 1–3 of training-methods with independent loss references, real training curves, changed masks or pairs and explicit synthetic-data limits. Add the new lesson evidence: unequal-count gradient oracle, rectangular patch/intervention checks, irreducible ambiguity example and temperature/ranking comparison. Explain each saved figure without copying its caption.", "furtherWork": "Full BERT/ViT masked autoencoders, natural corpora, augmentation studies and downstream probe benchmarks remain extensions. The current labs execute small attention, linear reconstruction and embedding models." }, "projectNote": "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.", "additionalProjectIds": [ "embedding-lineage" ], "lessons": [ { "id": "pretraining-01", "title": "Masked language modeling: predict hidden tokens", "status": "authored", "objective": "Separate the clean targets from the corrupted input. Normalize over selected target positions.", "prerequisites": [], "exercise": "Select the first position instead of the second and evaluate the trained model without updating it.", "evidence": "Keep corrupted inputs, selected-target counts, the log-vocabulary baseline, independent masked-loss checks, trained parameters and changed-mask results.", "check": "Why are visible input tokens useful even when their output positions are excluded from the loss?", "path": "phases/17-pretraining/01-mlm", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/17-pretraining/01-mlm/docs/en.md", "script": "exercises/pretraining-01.py", "notebook": "exercises/pretraining-01.ipynb", "markdown": "exercises/pretraining-01.md", "scriptSource": "phases/17-pretraining/01-mlm/code/main.py", "notebookSource": "phases/17-pretraining/01-mlm/notebooks/exercise.ipynb", "quizSource": "phases/17-pretraining/01-mlm/quiz.json", "evidenceSource": "phases/17-pretraining/01-mlm/outputs/evidence.md", "evidence": "exercises/pretraining-01-evidence.md" }, "minutes": 105, "contentHash": "e44d399501edc9669f44321bf1190105f65df97b7153323d41a087617b8458b3", "visualArtifact": { "image": "figures/pretraining-01.svg", "png": "figures/pretraining-01.png", "data": "figures/pretraining-01.json", "generatedAt": "2026-10-20T10:47:20.818351+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "e44d399501edc9669f44321bf1190105f65df97b7153323d41a087617b8458b3", "kind": "executed", "mechanismImage": "figures/pretraining-01-mechanism.svg", "mechanismPng": "figures/pretraining-01-mechanism.png" }, "execution": { "contentHash": "e44d399501edc9669f44321bf1190105f65df97b7153323d41a087617b8458b3", "executedAt": "2026-10-20T10:51:59.021149+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "MLM initial/final/changed-mask: 1.0986123085021973 0.010633035562932491 0.010457751341164112\nMLM initial/final/changed-mask: 1.0986123085021973 0.010633035562932491 0.010457751341164112\nUnselected output logits have zero direct loss gradient.\nToken-weighted / sequence-weighted: 0.6577723026275635 0.51289606\nChanged position loss: 0.010457751341164112\nEmpty mask rejected\nIndependent masked softmax gradient agrees.\nPASS: pretraining-01\n", "url": "executions/pretraining-01.json" }, "url": "lesson.html?lesson=pretraining-01" }, { "id": "pretraining-02", "title": "Masked image modeling: reconstruct missing patches", "status": "authored", "objective": "Preserve patch order and pixel order. Keep hidden pixels out of the encoder.", "prerequisites": [ "pretraining-01" ], "exercise": "Verify the bottom-right patch against a direct slice and explain why a whole-array reshape can hide a spatial permutation.", "evidence": "Keep patch-order slices, visible-only feature shapes, hidden-patch counts, independent MSE and held-amplitude reconstruction results.", "check": "Why must the encoder input be inspected separately from reconstruction loss?", "path": "phases/17-pretraining/02-mim", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/17-pretraining/02-mim/docs/en.md", "script": "exercises/pretraining-02.py", "notebook": "exercises/pretraining-02.ipynb", "markdown": "exercises/pretraining-02.md", "scriptSource": "phases/17-pretraining/02-mim/code/main.py", "notebookSource": "phases/17-pretraining/02-mim/notebooks/exercise.ipynb", "quizSource": "phases/17-pretraining/02-mim/quiz.json", "evidenceSource": "phases/17-pretraining/02-mim/outputs/evidence.md", "evidence": "exercises/pretraining-02-evidence.md" }, "minutes": 105, "contentHash": "546d7b69e5cbe151e99542665025a7bbfdd507e951bcde81b5c981400f065eab", "visualArtifact": { "image": "figures/pretraining-02.svg", "png": "figures/pretraining-02.png", "data": "figures/pretraining-02.json", "generatedAt": "2026-10-20T10:47:23.102205+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "546d7b69e5cbe151e99542665025a7bbfdd507e951bcde81b5c981400f065eab", "kind": "executed", "mechanismImage": "figures/pretraining-02-mechanism.svg", "mechanismPng": "figures/pretraining-02-mechanism.png" }, "execution": { "contentHash": "546d7b69e5cbe151e99542665025a7bbfdd507e951bcde81b5c981400f065eab", "executedAt": "2026-10-20T10:52:02.405713+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Masked pixel MSE initial/final/held amplitudes: 0.5538085699081421 0.0016928859986364841 0.0011766081443056464\nMasked pixel MSE initial/final/held amplitudes: 0.5538085699081421 0.0016928859986364841 0.0011766081443056464\nHidden MSE unchanged; full-image MSE is large.\nAmbiguous target MSE at predictions 0, 1, 2: 2, 1, 2\nBottom-right patch order verified.\nThe zero-loss shortcut receives hidden pixels; the real features contain only eight visible pixels and a bias.\nAll rectangular RGB patches agree; hidden-pixel intervention preserves visible features.\nPASS: pretraining-02\n", "url": "executions/pretraining-02.json" }, "url": "lesson.html?lesson=pretraining-02" }, { "id": "pretraining-03", "title": "Contrastive learning: views, positives and negatives", "status": "authored", "objective": "Define the positive pair before augmenting. Normalize embeddings and introduce temperature.", "prerequisites": [ "pretraining-02" ], "exercise": "Permute only the right-hand embeddings after training and demonstrate the effect of broken pair identity.", "evidence": "Keep source pairing, temperature, stable two-direction loss checks, collapse and duplicate baselines, training outputs and broken-pair diagnosis.", "check": "What does successful retrieval on the training pairs establish?", "path": "phases/17-pretraining/03-contrastive", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/17-pretraining/03-contrastive/docs/en.md", "script": "exercises/pretraining-03.py", "notebook": "exercises/pretraining-03.ipynb", "markdown": "exercises/pretraining-03.md", "scriptSource": "phases/17-pretraining/03-contrastive/code/main.py", "notebookSource": "phases/17-pretraining/03-contrastive/notebooks/exercise.ipynb", "quizSource": "phases/17-pretraining/03-contrastive/quiz.json", "evidenceSource": "phases/17-pretraining/03-contrastive/outputs/evidence.md", "evidence": "exercises/pretraining-03-evidence.md" }, "minutes": 105, "contentHash": "f831dc2fce1ef37c4dd173016bbc86e0579d06d729ff6ff7c5989e9cffadba28", "visualArtifact": { "image": "figures/pretraining-03.svg", "png": "figures/pretraining-03.png", "data": "figures/pretraining-03.json", "generatedAt": "2026-10-20T10:47:24.904609+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "f831dc2fce1ef37c4dd173016bbc86e0579d06d729ff6ff7c5989e9cffadba28", "kind": "executed" }, "execution": { "contentHash": "f831dc2fce1ef37c4dd173016bbc86e0579d06d729ff6ff7c5989e9cffadba28", "executedAt": "2026-10-20T10:52:06.547124+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Paired contrastive initial/final: 0.564157247543335 0.013556718826293945 ; all four nearest pairs correct\nPaired contrastive initial/final: 0.564157247543335 0.013556718826293945 ; all four nearest pairs correct\nCollapsed baseline: 1.3862943649291992\nTemperature/loss: 0.2 0.006715348921716213\nTemperature/loss: 1.0 0.31326162815093994\nTemperature/loss: 2.0 0.4740769863128662\nIncorrect pair mapping loss: 7.493563652038574\nDuplicated diagonal-only penalty: 0.6931472215801477\nStep / finite difference / autodiff: 0.001 -1.3803242444992065 -1.3803226947784424\nStep / finite difference / autodiff: 0.0005 -1.3802646398544312 -1.3803226947784424\nPASS: pretraining-03\n", "url": "executions/pretraining-03.json" }, "url": "lesson.html?lesson=pretraining-03" } ] }, { "id": "posttraining", "number": "18", "title": "Post-training: SFT, LoRA, reward models and RLHF", "part": "Training systems", "prerequisitePhaseIds": [ "pretraining" ], "hardware": "Executed CPU synthetic objectives; large-scale and human-feedback evaluation remain separate", "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.", "source": "https://docs.jax.dev/en/latest/", "description": "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.", "learningAdvice": "Predict the mask, denominator or preference direction before training. Keep independent objective checks and changed-condition evidence.", "projectId": "posttraining-audit", "projectStageIds": [ "1", "2", "3", "4" ], "studyGuide": { "question": "Which parameters and signals change during adaptation?", "why": "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.", "readiness": { "prompt": "Should response-only loss remove prompt tokens from the model context?", "answer": "No. Prompt inputs remain visible under the causal attention contract; the shifted response mask controls which targets receive loss.", "lessonId": "posttraining-01" }, "milestones": [ { "title": "Supervised fine-tuning with response-only token loss", "task": "Write an input/next-target/role table, audit ragged likelihood sums, identify supervised positions in the gradient map, and demonstrate a context distinction the transition model cannot represent.", "lessonIds": [ "posttraining-01" ] }, { "title": "LoRA: adapt, save and merge low-rank updates", "task": "Derive and check both adapter-factor gradients, explain asymmetric initialization, verify nonunit merge scaling, and calculate a rank-one capacity limit before interpreting the adaptation curve.", "lessonIds": [ "posttraining-02" ] }, { "title": "Learn a reward model from pairwise preferences", "task": "Derive preference likelihood and its linear gradient, inspect contradictory labels, distinguish offsets from scale, and explain the measured comparison margins before handing off a reward model.", "lessonIds": [ "posttraining-03" ] }, { "title": "RLHF mechanics: a frozen reward and PPO policy updates", "task": "Identify current, behavior, reference and reward states; work both clipping signs; verify baseline cancellation by enumeration; and compare the trained bandit with its regularized optimum and probability plot.", "lessonIds": [ "posttraining-04" ] }, { "title": "Direct preference optimization and reference-corrected margins", "task": "Work through reference correction and its gradient, reproduce a falling chosen-probability counterexample, and explain why positive margins do not by themselves prove improved response behavior.", "lessonIds": [ "posttraining-05" ] } ], "transfer": { "prompt": "Reward rises, but independently reviewed outputs get worse. What has been established?", "approach": "Only improvement against the learned reward. Investigate reward exploitation, preference coverage and independent task metrics; do not label the result human alignment." }, "pitfall": { "symptom": "A PPO ratio is always one during multiple optimizer epochs.", "check": "Check whether old log probabilities are incorrectly recomputed from the changing policy." }, "completion": "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. Add the new lesson evidence: ragged likelihood oracle, rank-capacity counterexample, conflicting preferences, exact bandit optimum and a DPO ratio improvement with lower chosen probability. State which model integrations remain absent.", "furtherWork": "The PPO lab is a terminal synthetic preference bandit, not full language-model RLHF. Human preference collection, sequence PPO/value models, large-model PEFT and held-out human evaluation remain separate work. For the learned-reward PPO lesson, complete rl-01 through rl-04 first. Response-only SFT and dense LoRA do not require the RL phase." }, "projectNote": "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.", "additionalProjectIds": [ "training-methods" ], "lessons": [ { "id": "posttraining-01", "title": "Supervised fine-tuning with response-only token loss", "status": "authored", "objective": "Align predictions with their target tokens. Separate the attention mask from the loss mask.", "prerequisites": [], "exercise": "Score only the first response token in each sequence and compare the resulting count and objective.", "evidence": "Keep shifted target/role masks, valid response-token counts, gradient support and the changed target-selection result.", "check": "Should prompt inputs disappear when using response-only loss?", "path": "phases/18-posttraining/01-sft", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/18-posttraining/01-sft/docs/en.md", "script": "exercises/posttraining-01.py", "notebook": "exercises/posttraining-01.ipynb", "markdown": "exercises/posttraining-01.md", "scriptSource": "phases/18-posttraining/01-sft/code/main.py", "notebookSource": "phases/18-posttraining/01-sft/notebooks/exercise.ipynb", "quizSource": "phases/18-posttraining/01-sft/quiz.json", "evidenceSource": "phases/18-posttraining/01-sft/outputs/evidence.md", "evidence": "exercises/posttraining-01-evidence.md" }, "minutes": 105, "contentHash": "d48f77f40513e77744330ab3407f8ff07c0aa48d1694fb60099b1fec0d05d3ee", "visualArtifact": { "image": "figures/posttraining-01.svg", "png": "figures/posttraining-01.png", "data": "figures/posttraining-01.json", "generatedAt": "2026-10-20T10:47:26.920905+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "d48f77f40513e77744330ab3407f8ff07c0aa48d1694fb60099b1fec0d05d3ee", "kind": "executed" }, "execution": { "contentHash": "d48f77f40513e77744330ab3407f8ff07c0aa48d1694fb60099b1fec0d05d3ee", "executedAt": "2026-10-20T10:52:09.960042+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Response-only token NLL initial/final: 1.6094379425048828 0.07985548675060272\nResponse-only token NLL initial/final: 1.6094379425048828 0.07985548675060272\nPrompt rows learn to predict first answers; the final end-token row has no target.\nSupervised targets per sequence: [2 3]\nTwo first-response targets: 0.07894014567136765\nNo-target SFT batch rejected\nIdentical final tokens force identical table predictions despite different contexts.\nPASS: posttraining-01\n", "url": "executions/posttraining-01.json" }, "url": "lesson.html?lesson=posttraining-01" }, { "id": "posttraining-02", "title": "LoRA: adapt, save and merge low-rank updates", "status": "authored", "objective": "Separate the frozen base from trainable adapters. Understand initialization and the first gradient.", "prerequisites": [ "posttraining-01" ], "exercise": "Save both adapter factors and their rank/scaling metadata, reload them and check merged versus unmerged predictions.", "evidence": "Keep base-training evidence, unchanged base arrays, adapter factors/rank/alpha, initial-gradient checks and reloaded/new-input merge parity.", "check": "Does LoRA require updating the base weights during adapter training?", "path": "phases/18-posttraining/02-lora", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/18-posttraining/02-lora/docs/en.md", "script": "exercises/posttraining-02.py", "notebook": "exercises/posttraining-02.ipynb", "markdown": "exercises/posttraining-02.md", "scriptSource": "phases/18-posttraining/02-lora/code/main.py", "notebookSource": "phases/18-posttraining/02-lora/notebooks/exercise.ipynb", "quizSource": "phases/18-posttraining/02-lora/quiz.json", "evidenceSource": "phases/18-posttraining/02-lora/outputs/evidence.md", "evidence": "exercises/posttraining-02-evidence.md" }, "minutes": 105, "contentHash": "5a084f0f528c1207e76c4deeaaa121102dcf48d223e7cb5781e9ce18431f807d", "visualArtifact": { "image": "figures/posttraining-02.svg", "png": "figures/posttraining-02.png", "data": "figures/posttraining-02.json", "generatedAt": "2026-10-20T10:47:28.955567+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "5a084f0f528c1207e76c4deeaaa121102dcf48d223e7cb5781e9ce18431f807d", "kind": "executed", "mechanismImage": "figures/posttraining-02-mechanism.svg", "mechanismPng": "figures/posttraining-02-mechanism.png" }, "execution": { "contentHash": "5a084f0f528c1207e76c4deeaaa121102dcf48d223e7cb5781e9ce18431f807d", "executedAt": "2026-10-20T10:52:13.228766+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "LoRA initial/final: 0.043869297951459885 2.466217665642034e-07 ; trainable: 10 base: 24\nLoRA initial/final: 0.043869297951459885 2.466217665642034e-07 ; trainable: 10 base: 24\nBoth zero factors produce zero first gradients.\nBoth LoRA factor gradients match the matrix-calculus oracle at alpha/r = 1.5.\nAdapter round trip and merge checked.\nNew-input merge passes; accidental double merge changes outputs.\nNonunit scaling passes; best rank-one identity-input MSE is 0.25.\nPASS: posttraining-02\n", "url": "executions/posttraining-02.json" }, "url": "lesson.html?lesson=posttraining-02" }, { "id": "posttraining-03", "title": "Learn a reward model from pairwise preferences", "status": "authored", "objective": "Represent a preference as a comparison. Fit score differences with a stable likelihood.", "prerequisites": [ "posttraining-02" ], "exercise": "Calculate the initial loss for tied reward scores independently and verify the trained ordering.", "evidence": "Keep synthetic label provenance, stable preference-loss and gradient references, offset/scale tests and the learned ordering.", "check": "What do pairwise preferences identify directly?", "path": "phases/18-posttraining/03-reward", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/18-posttraining/03-reward/docs/en.md", "script": "exercises/posttraining-03.py", "notebook": "exercises/posttraining-03.ipynb", "markdown": "exercises/posttraining-03.md", "scriptSource": "phases/18-posttraining/03-reward/code/main.py", "notebookSource": "phases/18-posttraining/03-reward/notebooks/exercise.ipynb", "quizSource": "phases/18-posttraining/03-reward/quiz.json", "evidenceSource": "phases/18-posttraining/03-reward/outputs/evidence.md", "evidence": "exercises/posttraining-03-evidence.md" }, "minutes": 105, "contentHash": "1f49dabb4b7fd72624bfb2aaf2c1b85347f1283ee4498dbb44b6602a861bbbd7", "visualArtifact": { "image": "figures/posttraining-03.svg", "png": "figures/posttraining-03.png", "data": "figures/posttraining-03.json", "generatedAt": "2026-10-20T10:47:30.613999+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "1f49dabb4b7fd72624bfb2aaf2c1b85347f1283ee4498dbb44b6602a861bbbd7", "kind": "executed" }, "execution": { "contentHash": "1f49dabb4b7fd72624bfb2aaf2c1b85347f1283ee4498dbb44b6602a861bbbd7", "executedAt": "2026-10-20T10:52:15.398245+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Synthetic preference NLL initial/final: 0.6931471824645996 0.08669456839561462 ; rewards: [ 1.9811294 0.327278 -2.3084073]\nSynthetic preference NLL initial/final: 0.6931471824645996 0.08669456839561462 ; rewards: [ 1.9811294 0.327278 -2.3084073]\nForward/reversed preference NLL: 0.08596369624137878 2.94565486907959\nBalanced conflicting labels prefer a zero margin.\nTied baseline and preference directions verified.\nRanking unchanged; confidence and downstream reward scale change.\nIndependent reward gradients agree in both label directions.\nPASS: posttraining-03\n", "url": "executions/posttraining-03.json" }, "url": "lesson.html?lesson=posttraining-03" }, { "id": "posttraining-04", "title": "RLHF mechanics: a frozen reward and PPO policy updates", "status": "authored", "objective": "Keep three identities separate. Collect actions before optimizing them.", "prerequisites": [ "posttraining-03" ], "exercise": "Check the KL at the fixed reference and compare it with the final policy’s exact KL.", "evidence": "Keep reward/reference identity, collected old probabilities, signed clipping checks, exact reward/KL history and detached behavior information.", "check": "Which policy is refreshed for the next rollout batch?", "path": "phases/18-posttraining/04-rlhf", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/18-posttraining/04-rlhf/docs/en.md", "script": "exercises/posttraining-04.py", "notebook": "exercises/posttraining-04.ipynb", "markdown": "exercises/posttraining-04.md", "scriptSource": "phases/18-posttraining/04-rlhf/code/main.py", "notebookSource": "phases/18-posttraining/04-rlhf/notebooks/exercise.ipynb", "quizSource": "phases/18-posttraining/04-rlhf/quiz.json", "evidenceSource": "phases/18-posttraining/04-rlhf/outputs/evidence.md", "evidence": "exercises/posttraining-04-evidence.md" }, "minutes": 105, "contentHash": "f95b5b63c839a7fb666de4e7d7d05613900eb9c24c3e422b20aed70502400f42", "visualArtifact": { "image": "figures/posttraining-04.svg", "png": "figures/posttraining-04.png", "data": "figures/posttraining-04.json", "generatedAt": "2026-10-20T10:47:32.867260+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "f95b5b63c839a7fb666de4e7d7d05613900eb9c24c3e422b20aed70502400f42", "kind": "executed", "mechanismImage": "figures/posttraining-04-mechanism.svg", "mechanismPng": "figures/posttraining-04-mechanism.png" }, "execution": { "contentHash": "f95b5b63c839a7fb666de4e7d7d05613900eb9c24c3e422b20aed70502400f42", "executedAt": "2026-10-20T10:52:18.834855+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Reward initial/final and final KL: 0.28191882371902466 1.9185447692871094 0.7824773192405701\nSynthetic preference bandit with actual sampled PPO updates; synthetic preference target or language-generation claim.\nReward initial/final and final KL: 0.28191882371902466 1.9185447692871094 0.7824773192405701\nSynthetic preference bandit with actual sampled PPO updates; synthetic preference target or language-generation claim.\nClipped surrogate terms: [ 2.4 -1.6]\nEnumerated baseline score gradient matches exact expected-reward gradient.\nReference/final KL: 0.0 0.7824773192405701\nBehavior log probabilities and advantages are detached.\nRegularized objective reference/trained/optimum: 0.28191882371902466 1.7620493173599243 1.7987911701202393\nPASS: posttraining-04\n", "url": "executions/posttraining-04.json" }, "url": "lesson.html?lesson=posttraining-04" }, { "id": "posttraining-05", "title": "Direct preference optimization and reference-corrected margins", "status": "authored", "objective": "Compare chosen and rejected response log probabilities. Subtract the reference preference.", "prerequisites": [ "posttraining-04" ], "exercise": "Swap chosen and rejected IDs after training and verify that the loss rises.", "evidence": "Keep reference identity, response log-probability conventions, log(2) initialization, independent margin loss, reversed labels and frozen-reference checks.", "check": "Does this DPO loop require a separately trained reward model and fresh policy rollouts?", "path": "phases/18-posttraining/05-dpo", "runtime": { "kind": "cpu" }, "hardware": "CPU", "artifacts": { "source": "phases/18-posttraining/05-dpo/docs/en.md", "script": "exercises/posttraining-05.py", "notebook": "exercises/posttraining-05.ipynb", "markdown": "exercises/posttraining-05.md", "scriptSource": "phases/18-posttraining/05-dpo/code/main.py", "notebookSource": "phases/18-posttraining/05-dpo/notebooks/exercise.ipynb", "quizSource": "phases/18-posttraining/05-dpo/quiz.json", "evidenceSource": "phases/18-posttraining/05-dpo/outputs/evidence.md", "evidence": "exercises/posttraining-05-evidence.md" }, "minutes": 105, "contentHash": "247b12d5db6bb478abfe606e2a31ed9fe445bee53ee43b39281b48654fdf44cc", "visualArtifact": { "image": "figures/posttraining-05.svg", "png": "figures/posttraining-05.png", "data": "figures/posttraining-05.json", "generatedAt": "2026-10-20T10:47:34.370265+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "247b12d5db6bb478abfe606e2a31ed9fe445bee53ee43b39281b48654fdf44cc", "kind": "executed", "mechanismImage": "figures/posttraining-05-mechanism.svg", "mechanismPng": "figures/posttraining-05-mechanism.png" }, "execution": { "contentHash": "247b12d5db6bb478abfe606e2a31ed9fe445bee53ee43b39281b48654fdf44cc", "executedAt": "2026-10-20T10:52:21.092144+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "DPO initial/final: 0.6931471824645996 0.27227783203125 ; reference-corrected margins: [3.0890286 6.17805719 3.08902884]\nDPO initial/final: 0.6931471824645996 0.27227783203125 ; reference-corrected margins: [3.0890286 6.17805719 3.08902884]\nCorrected/unadjusted initial loss: 0.6931471824645996 0.6540467739105225\nLoss before/after: 0.6931471824645996 0.5418724417686462 ; chosen probability: 0.4 -> 0.3\nReversed-label loss: 1.5064733028411865\nReference gradient is zero.\nOne-pair logit gradient agrees with independent margin derivative.\nPASS: posttraining-05\n", "url": "executions/posttraining-05.json" }, "url": "lesson.html?lesson=posttraining-05" } ] }, { "id": "tpu", "number": "19", "title": "TPU Setup: Provisioning & Runtime Verification", "part": "Start here", "prerequisitePhaseIds": [ "welcome" ], "hardware": "CPU topology/contract verifier; Cloud TPU VM (v5litepod-4 / v6e-4) or Colab TPU for target runs", "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.", "source": "https://cloud.google.com/tpu/docs/run-calculation-jax", "description": "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.", "learningAdvice": "Complete this short 2-lesson setup course right after 00: Setup & first steps, then continue directly to 20: TPU Workflows and 21: TPU Systems.", "projectId": "tpu-runtime-audit", "projectStageIds": [ "1", "2", "3", "4" ], "projectNote": "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.", "studyGuide": { "question": "Can you provision a Cloud TPU VM, verify that JAX actually executes on its chips, and prove clean teardown?", "why": "Every accelerator experiment starts with knowing how many hosts and GiB of HBM your slice provides, verifying that libtpu is active inside an isolated .venv, and deleting billable VMs when finished.", "readiness": { "prompt": "You created a Cloud TPU VM and gcloud says READY. What still has to happen before you trust a JAX run on it?", "answer": "Bootstrap ~/jax-tpu-lab/.venv with jax[tpu], verify inside Python that jax.default_backend() == 'tpu' and jax.device_count() matches your slice, and run a synchronized pmap reduction receipt.", "lessonId": "tpu-01" }, "milestones": [ { "lessonIds": [ "tpu-01" ], "title": "Select a single-host TPU topology and provision the VM", "task": "Compare v5litepod-1, v5litepod-4, v5litepod-8, and v6e-4 HBM totals, distinguish single-host VMs from multi-host Pod slices, and run gcloud create/describe/delete." }, { "lessonIds": [ "tpu-04" ], "title": "Bootstrap .venv and verify the runtime contract receipt", "task": "Install jax[tpu] in ~/jax-tpu-lab/.venv, run verify_runtime_contract with a synchronized pmap sum of squares, and copy back the SHA-256 receipt." } ], "transfer": { "prompt": "A script runs on a Cloud TPU VM without raising an error, but throughput is identical to a laptop CPU. What should your receipt check first?", "approach": "Inspect observed_backend and device_count in the JSON receipt; if libtpu was missing or JAX_PLATFORMS=tpu was not set, JAX silently fell back to CPU." }, "pitfall": { "symptom": "You closed your SSH terminal after finishing a run, assuming the Cloud TPU VM stopped billing.", "check": "Run gcloud compute tpus tpu-vm delete and confirm gcloud compute tpus tpu-vm list returns 0 items." }, "completion": "Retain the slice HBM table from tpu-01, the verified JSON runtime receipt from tpu-04, and the empty gcloud compute tpus tpu-vm list confirmation.", "furtherWork": "Continue immediately to Phase 20 (TPU Workflows: Launching Jobs & Checkpointed Experiments) to stage code and run resumable training loops." }, "additionalProjectIds": [ "workload-operations" ], "lessons": [ { "id": "tpu-01", "title": "Select a TPU slice and provision your Cloud TPU VM", "status": "authored", "objective": "Choose a single-host TPU topology, create a Cloud TPU VM with gcloud (or connect to a Colab TPU), verify SSH access, and rehearse clean VM deletion.", "prerequisites": [], "exercise": "Compute the chip count and total HBM GiB for v5litepod-4 and v6e-4, verify both are single-host slices (hosts == 1), and list all single-host slices with at least 64 GiB of HBM.", "evidence": "Keep the slice HBM capacity table, single-host vs Pod classification, and gcloud create/describe/delete command log.", "check": "Which command sequence proves that your Cloud TPU VM is READY when you start and no longer billing when you finish?", "path": "phases/19-tpu/01-get-a-tpu-running-and-verify-the-runtime", "runtime": { "kind": "cpu" }, "hardware": "CPU topology planner; Cloud TPU VM or Colab TPU for target provisioning", "artifacts": { "source": "phases/19-tpu/01-get-a-tpu-running-and-verify-the-runtime/docs/en.md", "script": "exercises/tpu-01.py", "notebook": "exercises/tpu-01.ipynb", "markdown": "exercises/tpu-01.md", "scriptSource": "phases/19-tpu/01-get-a-tpu-running-and-verify-the-runtime/code/main.py", "notebookSource": "phases/19-tpu/01-get-a-tpu-running-and-verify-the-runtime/notebooks/exercise.ipynb", "quizSource": "phases/19-tpu/01-get-a-tpu-running-and-verify-the-runtime/quiz.json", "evidenceSource": "phases/19-tpu/01-get-a-tpu-running-and-verify-the-runtime/outputs/evidence.md", "evidence": "exercises/tpu-01-evidence.md" }, "minutes": 25, "contentHash": "a67833264e71f5170b34b3db66b014a849e8aa6e5840de5def2a404bee7aae74", "visualArtifact": { "image": "figures/tpu-01.svg", "png": "figures/tpu-01.png", "data": "figures/tpu-01.json", "generatedAt": "2026-10-20T10:47:35.530166+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "a67833264e71f5170b34b3db66b014a849e8aa6e5840de5def2a404bee7aae74", "kind": "executed", "mechanismImage": "figures/tpu-01-mechanism.svg", "mechanismPng": "figures/tpu-01-mechanism.png" }, "execution": { "contentHash": "a67833264e71f5170b34b3db66b014a849e8aa6e5840de5def2a404bee7aae74", "executedAt": "2026-10-20T10:52:22.243997+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Verified runtime receipt: {\"contract_ok\": true, \"device_count\": 1, \"device_kinds\": [\"cpu\"], \"expected_backend\": \"cpu\", \"expected_sum_of_squares\": 140.0, \"jax\": \"0.9.2\", \"local_device_count\": 1, \"observed_backend\": \"cpu\", \"pmap_sum_of_squares\": 140.0, \"python\": \"3.14.3\", \"receipt_sha256\": \"4ddc0cbb323fe3bd\"}\nSingle-host TPU slice HBM totals (GiB): {'v5litepod-1': 16, 'v5litepod-4': 64, 'v5litepod-8': 128, 'v6e-4': 128}\nVerified runtime receipt: {\"contract_ok\": true, \"device_count\": 1, \"device_kinds\": [\"cpu\"], \"expected_backend\": \"cpu\", \"expected_sum_of_squares\": 140.0, \"jax\": \"0.9.2\", \"local_device_count\": 1, \"observed_backend\": \"cpu\", \"pmap_sum_of_squares\": 140.0, \"python\": \"3.14.3\", \"receipt_sha256\": \"4ddc0cbb323fe3bd\"}\nSingle-host TPU slice HBM totals (GiB): {'v5litepod-1': 16, 'v5litepod-4': 64, 'v5litepod-8': 128, 'v6e-4': 128}\nMismatch contract_ok: False observed: cpu devices: 1\nRequires --worker=all and jax.distributed.initialize(): {'v5litepod-1': False, 'v5litepod-4': False, 'v5litepod-8': False, 'v5litepod-16': True}\nChecked receipt sha256: 4ddc0cbb323fe3bd v5litepod-4 chips: 4 HBM GiB: 64\nSlices fitting 48 GiB: ['v5litepod-4', 'v5litepod-8', 'v6e-4']\nExpected 4-device pmap sum of squares: 10416.0\nPASS: tpu-01\n", "url": "executions/tpu-01.json" }, "url": "lesson.html?lesson=tpu-01" }, { "id": "tpu-04", "title": "Bootstrap .venv and verify a TPU runtime receipt", "status": "authored", "objective": "Create ~/jax-tpu-lab/.venv on the TPU VM, install jax[tpu], run a synchronized pmap runtime contract check, and download the signed JSON receipt before deleting the VM.", "prerequisites": [ "tpu-01" ], "exercise": "Call verify_runtime_contract for the active backend with min_devices=1, verify contract_ok is True, and compute expected_pmap_sum_of_squares(4, 8) for a 4-chip v5litepod-4 slice.", "evidence": "Keep the JSON runtime receipt (observed_backend, device_count, pmap_sum_of_squares, receipt_sha256) and the 1-chip vs 4-chip vs 8-chip pmap reference values.", "check": "Why must verify_runtime_contract check both jax.default_backend() and a synchronized pmap reduction rather than only checking that import jax succeeds?", "path": "phases/19-tpu/02-bootstrap-venv-and-verify-tpu-runtime-receipt", "runtime": { "kind": "cpu" }, "hardware": "CPU contract verifier; Cloud TPU VM or Colab TPU for target receipt", "artifacts": { "source": "phases/19-tpu/02-bootstrap-venv-and-verify-tpu-runtime-receipt/docs/en.md", "script": "exercises/tpu-04.py", "notebook": "exercises/tpu-04.ipynb", "markdown": "exercises/tpu-04.md", "scriptSource": "phases/19-tpu/02-bootstrap-venv-and-verify-tpu-runtime-receipt/code/main.py", "notebookSource": "phases/19-tpu/02-bootstrap-venv-and-verify-tpu-runtime-receipt/notebooks/exercise.ipynb", "quizSource": "phases/19-tpu/02-bootstrap-venv-and-verify-tpu-runtime-receipt/quiz.json", "evidenceSource": "phases/19-tpu/02-bootstrap-venv-and-verify-tpu-runtime-receipt/outputs/evidence.md", "evidence": "exercises/tpu-04-evidence.md" }, "minutes": 20, "contentHash": "fbf7f69ea6d4eb812a7d2bfd81046c806a9092bca95d7ae4737ed297ced95f15", "visualArtifact": { "image": "figures/tpu-04.svg", "png": "figures/tpu-04.png", "data": "figures/tpu-04.json", "generatedAt": "2026-10-20T10:47:36.709793+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "fbf7f69ea6d4eb812a7d2bfd81046c806a9092bca95d7ae4737ed297ced95f15", "kind": "analytic" }, "execution": { "contentHash": "fbf7f69ea6d4eb812a7d2bfd81046c806a9092bca95d7ae4737ed297ced95f15", "executedAt": "2026-10-20T10:52:23.450063+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Runtime receipt: {\"backend\": \"cpu\", \"expected_total\": 140.0, \"jax\": \"0.9.2\", \"local_device_count\": 1, \"observed_total\": 140.0, \"per_device_sums\": [140.0], \"python\": \"3.14.3\", \"receipt_sha256\": \"8c6bd0cbe29edc78\", \"verified\": true}\nExpected pmap sums by slice: {'1-device (CPU / v5e-1)': 140.0, '4-device (v5litepod-4 / v6e-4)': 10416.0, '8-device (v5litepod-8 / v6e-8)': 85344.0}\nRuntime receipt: {\"backend\": \"cpu\", \"expected_total\": 140.0, \"jax\": \"0.9.2\", \"local_device_count\": 1, \"observed_total\": 140.0, \"per_device_sums\": [140.0], \"python\": \"3.14.3\", \"receipt_sha256\": \"8c6bd0cbe29edc78\", \"verified\": true}\nExpected pmap sums by slice: {'1-device (CPU / v5e-1)': 140.0, '4-device (v5litepod-4 / v6e-4)': 10416.0, '8-device (v5litepod-8 / v6e-8)': 85344.0}\nPer-chip sums on a 4-chip slice: [140.0, 1100.0, 3084.0, 6092.0] Total: 10416.0\nTeardown verifier passed for empty and non-empty gcloud list outputs.\n4-chip expected sum: 10416.0 8-chip expected sum: 85344.0 receipt hash: 8c6bd0cbe29edc78\nSimulated 4-chip TPU check: True Local CPU check: False\nOvernight idle cost avoided (USD): 74.4\nPASS: tpu-04\n", "url": "executions/tpu-04.json" }, "url": "lesson.html?lesson=tpu-04" } ] }, { "id": "tpu-workflows", "number": "20", "title": "TPU Workflows: Launching Jobs & Checkpointed Experiments", "part": "Start here", "prerequisitePhaseIds": [ "tpu" ], "hardware": "CPU launcher rehearsal; Cloud TPU VM (single-host or --worker=all Pod slice) for target execution", "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.", "source": "https://cloud.google.com/tpu/docs/run-calculation-jax", "description": "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.", "learningAdvice": "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.", "projectId": "tpu-workflow-audit", "projectStageIds": [ "1", "2", "3", "4" ], "projectNote": "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.", "studyGuide": { "question": "Can you launch a JAX job over SSH, separate XLA warmup from steady-state timing, and resume cleanly after interruption?", "why": "Cloud TPU VMs can be preempted and multi-host Pod slices require lockstep execution across all workers. A reliable workflow saves atomic checkpoints, logs JSONL events, and copies artifacts back before teardown.", "readiness": { "prompt": "Step 1 of your TPU job takes 14 seconds while steps 2–50 take 11 milliseconds each. How should your launcher record timing?", "answer": "Run and synchronize one warmup call before the timed loop and log it as a separate warmup event so XLA compilation does not pollute steady-state step latency.", "lessonId": "tpu-02" }, "milestones": [ { "lessonIds": [ "tpu-02" ], "title": "Launch resumable TPU jobs with warmup separation", "task": "Run run_tpu_ready_job, log warmup separately from timed steps, and verify that a checkpoint-resumed run matches an uninterrupted control run bit-for-bit." }, { "lessonIds": [ "tpu-05" ], "title": "Stage course scripts and launch multi-host Pod slices", "task": "Build scp-in, ssh-run, and scp-out commands for any course script, use --worker=all on multi-host slices, and gate host artifact writes on jax.process_index() == 0." } ], "transfer": { "prompt": "You move a script from a single-host v5litepod-4 VM to a 4-host v5litepod-16 Pod slice, and it hangs at startup. What is the first thing to check?", "approach": "Check that you launched the script across all hosts simultaneously using gcloud compute tpus tpu-vm ssh --worker=all and called jax.distributed.initialize() before querying jax.devices()." }, "pitfall": { "symptom": "A resumed run loads the saved model weights `w`, but its next loss diverges from the uninterrupted control run.", "check": "Verify that the checkpoint also restored optimizer momentum `v` and step counter `step` so the deterministic batch schedule continues at the exact next step." }, "completion": "Retain the resumed vs control events.jsonl and state hashes from tpu-02, plus the single-host and --worker=all command workflows from tpu-05.", "furtherWork": "Continue to Phase 21 (TPU Systems: Generations, Memory, Precision & XProf) to budget HBM, test BF16/INT8 precision, and capture XProf traces." }, "additionalProjectIds": [ "workload-operations" ], "lessons": [ { "id": "tpu-02", "title": "Stage and launch checkpointed TPU jobs", "status": "authored", "objective": "Stage launch.py to a TPU VM, separate first-step XLA compilation warmup from steady-state step timing, and verify that resuming from an atomic checkpoint matches an uninterrupted run.", "prerequisites": [], "exercise": "Run run_tpu_ready_job for 4 steps (save_every=2), resume to 8 steps, and verify that the final state_hash at step 8 equals an uninterrupted 8-step run.", "evidence": "Keep events.jsonl (started, warmup, step, checkpoint, restored, completed), checkpoint-latest.json, and matching resumed vs control state hashes.", "check": "Why does run_tpu_ready_job execute train_step once and discard its outputs before starting the timed while step < steps loop?", "path": "phases/20-tpu-workflows/01-stage-and-launch-checkpointed-tpu-jobs", "runtime": { "kind": "cpu" }, "hardware": "CPU launcher rehearsal; Cloud TPU VM for target checkpointed runs", "artifacts": { "source": "phases/20-tpu-workflows/01-stage-and-launch-checkpointed-tpu-jobs/docs/en.md", "script": "exercises/tpu-02.py", "notebook": "exercises/tpu-02.ipynb", "markdown": "exercises/tpu-02.md", "scriptSource": "phases/20-tpu-workflows/01-stage-and-launch-checkpointed-tpu-jobs/code/main.py", "notebookSource": "phases/20-tpu-workflows/01-stage-and-launch-checkpointed-tpu-jobs/notebooks/exercise.ipynb", "quizSource": "phases/20-tpu-workflows/01-stage-and-launch-checkpointed-tpu-jobs/quiz.json", "evidenceSource": "phases/20-tpu-workflows/01-stage-and-launch-checkpointed-tpu-jobs/outputs/evidence.md", "evidence": "exercises/tpu-02-evidence.md" }, "minutes": 25, "contentHash": "90a369489019c4bccc04bade800129bfb8c90c341621d2a9a4160ef2a6d2b1d4", "visualArtifact": { "image": "figures/tpu-02.svg", "png": "figures/tpu-02.png", "data": "figures/tpu-02.json", "generatedAt": "2026-10-20T10:54:22.790744+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "90a369489019c4bccc04bade800129bfb8c90c341621d2a9a4160ef2a6d2b1d4", "kind": "executed", "mechanismImage": "figures/tpu-02-mechanism.svg", "mechanismPng": "figures/tpu-02-mechanism.png" }, "execution": { "contentHash": "90a369489019c4bccc04bade800129bfb8c90c341621d2a9a4160ef2a6d2b1d4", "executedAt": "2026-10-20T10:58:53.178365+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Resumed state_hash matches uninterrupted control: ecba6a13e41beb0003af0371f8422efa3b08538c4420efb07a9bc61b0a10cdc1\nLosses across 6 steps: [3.6712, 3.78913, 1.91372, 2.09553, 1.81174, 1.16636]\nResumed state_hash matches uninterrupted control: ecba6a13e41beb0003af0371f8422efa3b08538c4420efb07a9bc61b0a10cdc1\nLosses across 6 steps: [3.6712, 3.78913, 1.91372, 2.09553, 1.81174, 1.16636]\nWarmup ms: 79.34 Median steady step ms: 0.0432\nResume rejection: checkpoint state hash mismatch\nStep 8 resumed vs control hash: eba1c0bcfb80184d9f5108c1d6539369e21f95276e5d6806234237f41eccbed9 eba1c0bcfb80184d9f5108c1d6539369e21f95276e5d6806234237f41eccbed9\nRecorded event sequence: ['started', 'warmup', 'step', 'step', 'checkpoint', 'completed', 'restored', 'warmup', 'step', 'step', 'checkpoint', 'completed']\nLoaded step: 4 Leftover tmp files: []\nPASS: tpu-02\n", "url": "executions/tpu-02.json" }, "url": "lesson.html?lesson=tpu-02" }, { "id": "tpu-05", "title": "Run course scripts and multi-host TPU Pod slices", "status": "authored", "objective": "Stage any course script in public/exercises/ to a TPU VM, construct single-host and --worker=all multi-host SSH launch commands, and verify jax.distributed.initialize() requirements.", "prerequisites": [ "tpu-02" ], "exercise": "Generate the workflow for distributed-02.py on a 4-host v5litepod-16 slice (hosts=4) and verify that --worker=all is present in the SSH command and requires_distributed_init is True.", "evidence": "Keep the generated scp_in_cmd, ssh_run_cmd, and scp_out_cmd workflow for single-host and multi-host slices plus the process-0 checkpoint writer check.", "check": "What happens if you launch a JAX script on only worker 0 of a 4-host v5litepod-16 slice instead of passing --worker=all?", "path": "phases/20-tpu-workflows/02-run-course-scripts-and-multihost-tpu-slices", "runtime": { "kind": "cpu" }, "hardware": "CPU workflow planner; Cloud TPU VM or Pod slice for target runs", "artifacts": { "source": "phases/20-tpu-workflows/02-run-course-scripts-and-multihost-tpu-slices/docs/en.md", "script": "exercises/tpu-05.py", "notebook": "exercises/tpu-05.ipynb", "markdown": "exercises/tpu-05.md", "scriptSource": "phases/20-tpu-workflows/02-run-course-scripts-and-multihost-tpu-slices/code/main.py", "notebookSource": "phases/20-tpu-workflows/02-run-course-scripts-and-multihost-tpu-slices/notebooks/exercise.ipynb", "quizSource": "phases/20-tpu-workflows/02-run-course-scripts-and-multihost-tpu-slices/quiz.json", "evidenceSource": "phases/20-tpu-workflows/02-run-course-scripts-and-multihost-tpu-slices/outputs/evidence.md", "evidence": "exercises/tpu-05-evidence.md" }, "minutes": 20, "contentHash": "09cbfe2aa061d9ab3329236f15e12a6bbbdb04609765c10fc42a27e3e277a585", "visualArtifact": { "image": "figures/tpu-05.svg", "png": "figures/tpu-05.png", "data": "figures/tpu-05.json", "generatedAt": "2026-10-20T10:58:50.448372+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "09cbfe2aa061d9ab3329236f15e12a6bbbdb04609765c10fc42a27e3e277a585", "kind": "executed" }, "execution": { "contentHash": "09cbfe2aa061d9ab3329236f15e12a6bbbdb04609765c10fc42a27e3e277a585", "executedAt": "2026-10-20T10:58:54.082998+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "CPU route: python public/exercises/distributed-02.py\nCPU-only route remains excluded from TPU Pod execution\nCorrectly rejected CPU-only lesson for TPU Pod execution\nSlice catalog global chips: {'v5litepod-4': 4, 'v5litepod-8': 8, 'v5litepod-16': 16, 'v6e-16': 16}\nCPU route: python public/exercises/distributed-02.py\nv5litepod-16 global chips: 16 Idle chips without --worker=all: 12\nTotal HBM GiB by slice: {'v5litepod-4': 64, 'v5litepod-8': 128, 'v5litepod-16': 256, 'v6e-16': 512}\nCPU route: python public/exercises/distributed-02.py\nRejected unsafe route: {}\nRejected unsafe route: {'hosts': 4, 'execution_class': 'cpu-only'}\nMulti-host Pod slices: ['v5litepod-16', 'v6e-16']\nPer-host batch: 32 Global batch: 128\nPASS: tpu-05\n", "url": "executions/tpu-05.json" }, "url": "lesson.html?lesson=tpu-05" } ] }, { "id": "tpu-systems", "number": "21", "title": "TPU Systems: Generations, Memory, Precision & XProf", "part": "Start here", "prerequisitePhaseIds": [ "tpu-workflows" ], "hardware": "CPU numerical & profiler harness; Cloud TPU VM for target XProf traces", "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.", "source": "https://openxla.org/xprof", "description": "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.", "learningAdvice": "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.", "projectId": "tpu-systems-audit", "projectStageIds": [ "1", "2", "3", "4" ], "projectNote": "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.", "studyGuide": { "question": "How do you size TPU HBM capacity, bound BF16/INT8 precision error, and capture a clean XProf trace?", "why": "Provisioning a TPU slice without budgeting HBM wastes time on out-of-memory crashes, and profiling without warming up measures compiler overhead instead of MXU execution.", "readiness": { "prompt": "A matrix multiplication has arithmetic intensity 170 FLOP/byte on TPU v5e (ridge point 240.5 FLOP/byte). Is it compute-bound or HBM bandwidth-bound?", "answer": "HBM bandwidth-bound, because its arithmetic intensity (170 FLOP/byte) is below the TPU v5e ridge point (240.5 FLOP/byte). Increasing batch size or sequence length raises operational intensity toward the MXU ceiling.", "lessonId": "tpu-03" }, "milestones": [ { "lessonIds": [ "tpu-03" ], "title": "Compare TPU generations and budget HBM capacity", "task": "Compute TPU v5e, v5p, and v6e ridge points, calculate FP32 vs mixed-BF16 training state and KV-cache GiB, and determine the minimum chip count at an 80% HBM budget." }, { "lessonIds": [ "tpu-06" ], "title": "Audit BF16/INT8 precision and capture a warmed XProf trace", "task": "Compare BF16 and symmetric INT8 matmul errors against FP32, verify FP32 accumulation inside jax.lax.dot_general, and capture a warmed *.xplane.pb trace." } ], "transfer": { "prompt": "An XProf trace on your TPU VM shows a 12-second host compile block followed by five 9-millisecond step blocks. How do you get a clean steady-state profile?", "approach": "Run at least one warmup call and call jax.block_until_ready() before entering with jax.profiler.trace(trace_dir): so only steady-state execution is recorded." }, "pitfall": { "symptom": "Accumulating a long bfloat16 dot product directly in bfloat16 produces much larger drift than expected.", "check": "Pass preferred_element_type=jnp.float32 to jax.lax.dot_general so the TPU MXU accumulates products in FP32 before any output cast." }, "completion": "Retain the TPU ridge-point and HBM budget tables from tpu-03, plus the BF16/INT8 error summary and *.xplane.pb trace receipt from tpu-06.", "furtherWork": "Apply the 3-part Cloud TPU mini-course across Phases 07, 08, 09, 13, 15, 16, 17, and 18 by staging each lesson script in exercises/ to your TPU VM." }, "additionalProjectIds": [ "workload-operations" ], "lessons": [ { "id": "tpu-03", "title": "TPU generations and HBM memory budgeting", "status": "authored", "objective": "Compare TPU v5e, v5p, and v6e specifications and roofline ridge points, budget FP32 vs mixed-BF16 training state and KV-cache HBM, and size the minimum chip count for a model.", "prerequisites": [], "exercise": "Compute estimate_memory_gib for a 1.5B parameter model with 4096 KV tokens and verify that its mixed-precision training state fits comfortably inside a single-host v5litepod-4 (64 GiB total HBM).", "evidence": "Keep the TPU ridge-point table, 7B and 1.5B training/KV-cache HBM budgets, and minimum chip count calculation at an 80% HBM ceiling.", "check": "Why does mixed-precision Adam training still require ~12 bytes per parameter in HBM even when forward and backward matmuls use 2-byte bfloat16?", "path": "phases/21-tpu-systems/01-tpu-generations-and-hbm-memory-budgeting", "runtime": { "kind": "cpu" }, "hardware": "CPU capacity & roofline planner; Cloud TPU VM for target HBM verification", "artifacts": { "source": "phases/21-tpu-systems/01-tpu-generations-and-hbm-memory-budgeting/docs/en.md", "script": "exercises/tpu-03.py", "notebook": "exercises/tpu-03.ipynb", "markdown": "exercises/tpu-03.md", "scriptSource": "phases/21-tpu-systems/01-tpu-generations-and-hbm-memory-budgeting/code/main.py", "notebookSource": "phases/21-tpu-systems/01-tpu-generations-and-hbm-memory-budgeting/notebooks/exercise.ipynb", "quizSource": "phases/21-tpu-systems/01-tpu-generations-and-hbm-memory-budgeting/quiz.json", "evidenceSource": "phases/21-tpu-systems/01-tpu-generations-and-hbm-memory-budgeting/outputs/evidence.md", "evidence": "exercises/tpu-03-evidence.md" }, "minutes": 25, "contentHash": "e6dd2b245a8dfadb2626f5ecca9e89c0a274b8b8c32b51ec9ce69dffaefa00f9", "visualArtifact": { "image": "figures/tpu-03.svg", "png": "figures/tpu-03.png", "data": "figures/tpu-03.json", "generatedAt": "2026-10-20T10:58:50.515621+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "e6dd2b245a8dfadb2626f5ecca9e89c0a274b8b8c32b51ec9ce69dffaefa00f9", "kind": "analytic", "mechanismImage": "figures/tpu-03-mechanism.svg", "mechanismPng": "figures/tpu-03-mechanism.png" }, "execution": { "contentHash": "e6dd2b245a8dfadb2626f5ecca9e89c0a274b8b8c32b51ec9ce69dffaefa00f9", "executedAt": "2026-10-20T10:58:55.101094+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "TPU ridge points (FLOP/byte): {'v5e': 240.5, 'v5p': 166.0, 'v6e': 573.8}\n7B memory budget (GiB): {'fp32_train_state_gib': 104.308, 'mixed_bf16_train_state_gib': 117.347, 'kv_fp32_gib': 2.0, 'kv_bf16_gib': 1.0, 'kv_int8_gib': 0.5}\n1.5B memory budget (GiB): {'fp32_train_state_gib': 22.352, 'mixed_bf16_train_state_gib': 25.146, 'kv_fp32_gib': 1.0, 'kv_bf16_gib': 0.5, 'kv_int8_gib': 0.25}\nTPU ridge points (FLOP/byte): {'v5e': 240.5, 'v5p': 166.0, 'v6e': 573.8}\n7B memory budget (GiB): {'fp32_train_state_gib': 104.308, 'mixed_bf16_train_state_gib': 117.347, 'kv_fp32_gib': 2.0, 'kv_bf16_gib': 1.0, 'kv_int8_gib': 0.5}\n1.5B memory budget (GiB): {'fp32_train_state_gib': 22.352, 'mixed_bf16_train_state_gib': 25.146, 'kv_fp32_gib': 1.0, 'kv_bf16_gib': 0.5, 'kv_int8_gib': 0.25}\nKV cache GiB (FP32 / BF16 / INT8): 2.0 1.0 0.5\n7B mixed train state GiB: 117.347 Minimum chips at 80% usable HBM: 10\n1.5B state plus BF16 KV cache (GiB): 25.646\nArithmetic intensity: 341.3333333333333 FLOPs/byte\nv6e-4 80% usable GiB: 102.4 Fits 7B state estimate: False\nPASS: tpu-03\n", "url": "executions/tpu-03.json" }, "url": "lesson.html?lesson=tpu-03" }, { "id": "tpu-06", "title": "bfloat16, int8 precision, and XProf profiling on TPU", "status": "authored", "objective": "Measure bfloat16 and symmetric int8 matrix multiplication error against float32, verify float32 accumulator precision, and capture a warmed XProf (*.xplane.pb) trace.", "prerequisites": [ "tpu-03" ], "exercise": "Run capture_warmed_xprof_trace for 4 steps of size 96 on the active backend, verify trace_file_count >= 1, and assert that bf16_max_abs_err is smaller than int8_max_abs_err.", "evidence": "Keep the BF16 vs INT8 error summary, FP32-vs-BF16 accumulator drift comparison, and captured *.xplane.pb trace receipt.", "check": "Why should you run and synchronize at least one warmup step before entering with jax.profiler.trace(trace_dir): on a TPU VM?", "path": "phases/21-tpu-systems/02-bfloat16-int8-precision-and-xprof-profiling", "runtime": { "kind": "cpu" }, "hardware": "CPU numerical & profiler harness; Cloud TPU VM for target XProf traces", "artifacts": { "source": "phases/21-tpu-systems/02-bfloat16-int8-precision-and-xprof-profiling/docs/en.md", "script": "exercises/tpu-06.py", "notebook": "exercises/tpu-06.ipynb", "markdown": "exercises/tpu-06.md", "scriptSource": "phases/21-tpu-systems/02-bfloat16-int8-precision-and-xprof-profiling/code/main.py", "notebookSource": "phases/21-tpu-systems/02-bfloat16-int8-precision-and-xprof-profiling/notebooks/exercise.ipynb", "quizSource": "phases/21-tpu-systems/02-bfloat16-int8-precision-and-xprof-profiling/quiz.json", "evidenceSource": "phases/21-tpu-systems/02-bfloat16-int8-precision-and-xprof-profiling/outputs/evidence.md", "evidence": "exercises/tpu-06-evidence.md" }, "minutes": 25, "contentHash": "59d9ebf07ea1f8d61e29d1453c803ce4ba21902ab870f37fcb1976de8ffdc26f", "visualArtifact": { "image": "figures/tpu-06.svg", "png": "figures/tpu-06.png", "data": "figures/tpu-06.json", "generatedAt": "2026-10-20T10:58:51.060890+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "59d9ebf07ea1f8d61e29d1453c803ce4ba21902ab870f37fcb1976de8ffdc26f", "kind": "executed" }, "execution": { "contentHash": "59d9ebf07ea1f8d61e29d1453c803ce4ba21902ab870f37fcb1976de8ffdc26f", "executedAt": "2026-10-20T10:58:58.362237+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Precision comparison: [{'policy': 'bf16 (fp32 accum)', 'max_abs_err': 0.03718, 'mean_abs_err': 0.00756, 'weight_bytes': 16384, 'scale_bytes': 256}, {'policy': 'bf16 (bf16 accum)', 'max_abs_err': 0.05545, 'mean_abs_err': 0.00912, 'weight_bytes': 16384, 'scale_bytes': 256}, {'policy': 'int8 per-col (fp32 accum)', 'max_abs_err': 0.12516, 'mean_abs_err': 0.02322, 'weight_bytes': 16384, 'scale_bytes': 256}]\nCaptured *.xplane.pb trace count: 4 steady_ms: 0.0845\nPrecision comparison: [{'policy': 'bf16 (fp32 accum)', 'max_abs_err': 0.03718, 'mean_abs_err': 0.00756, 'weight_bytes': 16384, 'scale_bytes': 256}, {'policy': 'bf16 (bf16 accum)', 'max_abs_err': 0.05545, 'mean_abs_err': 0.00912, 'weight_bytes': 16384, 'scale_bytes': 256}, {'policy': 'int8 per-col (fp32 accum)', 'max_abs_err': 0.12516, 'mean_abs_err': 0.02322, 'weight_bytes': 16384, 'scale_bytes': 256}]\nRetained trace directory: /Users/tahabsn/Documents/ChatGPT/jaxpathways/tpu-xprof-trace\nCaptured *.xplane.pb trace count: 4 steady_ms: 0.0663\nBF16-accum / FP32-accum max error ratio: 1.49\nWarmup ms: 16.804 Traced steady step ms: 0.0663\nSeed 42 mean errors: {'bf16 (fp32 accum)': 0.0075, 'bf16 (bf16 accum)': 0.00902, 'int8 per-col (fp32 accum)': 0.02182} xplane files: 4\nMean error (per-channel vs per-tensor): 0.01389 0.0979\nxplane.pb bytes: 259893\nPASS: tpu-06\n", "url": "executions/tpu-06.json" }, "url": "lesson.html?lesson=tpu-06" } ] } ] }