TPU 101 with JAX: A Medium Series
A six-part, beginner-friendly course that takes you from “I’ve never touched a TPU” to training models across multiple TPU chips in JAX.
Every code snippet in this series was run on real TPU hardware, and the output blocks come from actual runs — not invented examples.
The series
Part 1 — Why TPUs, Why JAX, Why Now
The GPU vs. TPU mental model. What kind of chip a TPU actually is, how it differs from a GPU at the architectural level, and most importantly, which chip wins for which workload. No code; pure orientation.
Part 2 — Getting a TPU and Running Your First JAX Program
Provision a Spot TPU VM on Google Cloud. Install gcloud on a Mac, including the real errors you may hit, or skip the local setup and use Cloud Shell. SSH into the TPU VM, install JAX, and get a TpuDevice(...) printout from your first script.
Part 3 — JAX for PyTorch People
The seven core JAX concepts: immutability, jit, grad, vmap, PRNG keys, pytrees, and the functional training step. Each concept is mapped back to its PyTorch equivalent so the transition feels practical, not abstract.
Part 4 — Inside the TPU, and a Head-to-Head With a GPU
What @jit actually produces, including how to dump and read the compiled program. Why bf16 matters, with measured numbers. Then the same benchmark on a Colab T4 GPU, so the performance story is concrete rather than theoretical.
Part 5 — Your First Real Model: MNIST in JAX on TPU
An end-to-end training run with real data from disk, Optax’s Adam optimizer, and a jit-compiled training step. The model reaches 97.59% test accuracy in 3.8 seconds on a v6e-1.
Part 6 — Scaling Out, and Shutting Down
Multi-chip training on a v6e-4. Data parallelism with pmap, tensor parallelism with shard_map, and a live demonstration of where tensor parallelism can actually slow you down before it starts winning. The series ends with shutdown steps and a cost audit.
Who this is for
You know Python. You have trained models in PyTorch or TensorFlow. You have never touched JAX or TPUs.
Get Roya’s stories in your inbox
Join Medium for free to get updates from this writer.
You want to understand what the shift to JAX + TPU actually means for your workflow: when it helps, when it hurts, and how to get started.
No prior JAX or TPU experience is assumed.
What it costs to run
Everything in this series runs on Spot-priced TPUs. Reproducing the full tutorial yourself should cost roughly $1.50–$3.00 in total Google Cloud spend, depending on region, availability, and how long you leave resources running.
Reading time
Reading only: about 2 hours.
Reading + running every script yourself: about 3–4 hours.
