Colab notebooks for the AI GDE Marathon: JAX on TPU Tutorial. Apache 2.0.
Two notebooks. Open one in Colab, pick Runtime → Change runtime type → v5e-1 TPU, then Runtime → Run all.
| Notebook | What it teaches | Token? | Status |
|---|---|---|---|
 01_jax_tpu_mechanics |
Buffer donation, jit retracing and static shapes, cached decode, int8 KV — measured on synthetic arrays |
No | In progress |
02_e2b_repack_vs_stock_on_jax |
Gemma 4 E2B in pure JAX: Google's W4A16 export against an exact repack of the QAT weights — weight error, next-token KL, speed and memory | No | Ready |
Notebook 02 downloads about 26 GB and takes 30 to 40 minutes on v5e-1. To keep your outputs, use File → Save a copy in Drive.
The measurements come from single-chip TPU rigs in ~/tpu-jax and
~/gemma4-dev. Colab's TPU runtimes are v5e-1 and v6e-1 — the same silicon
those rigs used — so the findings are re-measured live in the notebook rather
than quoted at the reader. These notebooks target v5e-1 (one TPU v5 lite,
16.91 GB HBM, JAX 0.7.2 as of 2026-09-20); the source rigs measured on v6e-1,
and citations say so. See docs/OUTLINE.md for the
section-by-section source map.
Notebooks are build artifacts. Edit src/*.py, never notebooks/*.ipynb.
Sources are jupytext py:percent files, so
a prose change is a one-line diff instead of a re-encoded JSON blob. The build
prepends the Apache-2.0 cell and the "Open in Colab" badge, and sets
metadata.accelerator = "TPU" so the reader lands on a TPU runtime instead of
silently benchmarking a CPU.
The .ipynb files are committed, because that is what Colab serves from
GitHub. make check fails if they have drifted from their sources.
make setup # pip install -r requirements-dev.txt: jupytext, nbformat, google-colab-cli
make build # src/*.py -> notebooks/*.ipynb
make check # fail if committed notebooks are stale (CI runs this)
make verify # run every notebook on a real Colab TPU; fail on any errormake verify is the gate. Rendered markdown will not tell you that cell 14
raises, so every notebook is executed on the accelerator it claims before it
ships, via the Colab CLI:
colab --auth adc new -s verify-01 --tpu v5e1
colab --auth adc exec -s verify-01 -f notebooks/01_jax_tpu_mechanics.ipynb
colab --auth adc stop -s verify-01
colab exec writes the executed copy back as *_output.ipynb;
tools/verify_on_tpu.py scans it for output_type == "error" and unexpected
stderr, then files it under runs/<date>/ as evidence.
--auth adc reuses your gcloud Application Default Credentials. They must be a
user credential, since a Colab runtime belongs to a Google user and bills
against that user's compute units, and they must carry the colaboratory
scope, which a plain gcloud auth application-default login does not grant.
--auth oauth2 uses the CLI's own login instead. Sign in once in a terminal,
colab --auth oauth2 sessions, and paste back the code it asks for; the token
is saved under ~/.config/colab-cli/ and refreshes itself, so later runs need
no prompt:
python3 tools/verify_on_tpu.py --auth oauth2 02_e2b_repack_vs_stock_on_jax
make verify-v6e1 targets the 32 GB chip, but v6e-1 needs a Colab tier this
account does not have — it returns Backend rejected accelerator 'V6E1'.
These runs provision real hardware and spend compute units. make sessions
lists anything still up; colab stop -s <name> releases it.
google-colab-cli— headless. Provisions--tpu v5e1|v6e1, executes a local.ipynb, tears down. This is the build and CI half, and the only one used bymake verify.colab-mcp— not a headless builder. It proxies over a websocket to a notebook already open in your browser, exposing one tool (open_colab_browser_connection) until the page hands over the editing tools. Good for live iteration on one notebook against a real TPU; wrong tool for building two deterministically.
make setup installs the CLI with plain pip.
src/ py:percent sources — the thing you edit
notebooks/ generated .ipynb — committed, served to Colab from GitHub
tools/ build_notebooks.py, verify_on_tpu.py
docs/ OUTLINE.md: section-by-section source map and known gaps
runs/ executed notebooks kept as evidence that a version ran clean
