TPU 101 — Part 1: Why TPUs, Why JAX, Why Now

A beginner-friendly series on Google TPUs and JAX. This part: what TPUs actually are, why they’re different from GPUs, and which one is the right tool for which job. No code yet, just the mental model.

If you’ve ever trained a model, you’ve probably used a GPU. Maybe you clicked “GPU” in a Colab dropdown. Maybe you wrote .to("cuda"). Either way, the chip mostly stayed invisible; you handed it work, and it gave you results.

Lately, there’s been a quiet shift. The researchers at Google, DeepMind, and a growing chunk of the AI world are running their biggest models on a different kind of chip entirely: a TPU, short for Tensor Processing Unit, a specialized processor Google designed to accelerate machine-learning workloads. And the library they’re reaching for to drive those chips isn’t PyTorch. It’s JAX.

This series is a hands-on tour of that world. Over six parts, we’ll go from “I’ve never touched a TPU” to training a model across multiple TPU chips in JAX. But before any code, you deserve an honest answer to a question that usually gets glossed over:

What actually is a TPU, and why would I use one?

That’s what this part is about. No code. Just the right mental model so everything in the next five parts has a place to land.

The restaurant kitchen analogy

I’m going to use one analogy throughout this whole article, so let’s set it up.

Imagine two restaurants. Both serve food. Both have a kitchen. But they’re designed very differently.

Restaurant A is a diner. The kitchen has 30 cooks, each with their own small station. One cook grills burgers, another plates salads, another flips pancakes, and another makes milkshakes. If you order a weird custom dish, that’s fine; one of the cooks will figure it out. The kitchen handles the variety well. Orders can be big or small, simple or complicated; the cooks improvise.

Restaurant B is a pasta factory. There’s one enormous machine in the middle of the room. You pour flour and water in one end; perfect, identical pasta comes out the other. It’s spectacularly fast at making pasta. It is useless at making anything that isn’t pasta. The menu is short on purpose.

A GPU is the diner. A TPU is the pasta factory.

That’s the whole idea. The rest of this article is just filling in what the “flour and water” is, and why “pasta” turns out to be the overwhelming majority of what a deep-learning model actually asks a chip to do.

What a GPU actually is, in one paragraph

Before we compare, let’s demystify the GPU itself.

The “G” in GPU stands for “graphics.” GPUs were originally designed to render video game frames, millions of tiny pixels, each needing a similar but slightly different calculation, all at once. So NVIDIA built chips with thousands of small, independent calculators, called cores. Think of each core as one of those diner cooks; small, fast, capable of doing whatever you ask. They sit in clusters called streaming multiprocessors (SMs), which is just a fancy name for “a group of cooks that share a prep station.”

When deep learning took off around 2012, it turned out that matrix multiplication, the math at the heart of every neural network, could be split up and distributed across those thousands of cores. GPUs have become much more useful than their designers ever intended. Then NVIDIA added Tensor Cores, specialized accelerators that perform matrix math much faster than regular cores, and GPUs became the default chip for AI.

The important thing: a GPU is thousands of small, general-purpose calculators, plus some specialized ones, all sharing memory. Keeping all of them busy is the game.

What a TPU is, in one paragraph

A TPU was designed from scratch in 2015 for one purpose: to run neural networks fast. Google didn’t try to make a better GPU. They asked a different question — what do neural networks actually do? The answer, almost embarrassingly, is: they do huge matrix multiplications, over and over, with some other minor math sprinkled in between.

So Google built a chip that does exactly that. The centerpiece is a big dedicated matrix-multiplying unit called the MXU. Not thousands of small calculators, but a small number of very large, specialized ones (usually one or two MXUs per core, sometimes more on newer generations). You feed them two matrices; they multiply them. That’s the pasta machine.

The important thing: a TPU is built around large matrix engines, plus some helper hardware for everything else. Feeding those engines efficiently is the game.

Here’s what both look like side by side:

Now we can take a closer look at a TPU.

Inside a TPU chip

Here’s a simplified picture. Real TPUs also contain scalar processors, data-movement engines, and compiler-managed scheduling logic, but as a first mental model, it’s fine to think of a TPU as having four main pieces. I’ll give you both the technical name and what each one actually does, in plain terms.

The MXU (Matrix Multiply Unit). The pasta machine. This is a giant grid of multipliers, typically 128×128 or 256×256, arranged so that a matrix flows through it like an assembly line. Most of a TPU’s computing power is dedicated to matrix operations like these, and the MXU does exactly one thing: multiply two matrices together. That’s all you use it for, but it’s what nearly every line of a neural network eventually boils down to, so that’s fine.

The VPU (Vector Processing Unit). The side-dish station. Not everything in a neural network is a matrix multiplication. You also need to apply activation functions (like ReLU), normalize things, add bias terms, and compute loss. The VPU handles all of that. It’s much smaller than the MXU because there’s less work to do, but it’s there because if you didn’t have it, you’d have to ship intermediate results off the chip to do them, and that would be slow.

HBM (High-Bandwidth Memory). The pantry. This is the chip’s main memory — 16 to 32 gigabytes of it. It holds your model’s weights and activations. It’s fast, but “fast” is relative; reading from HBM is still much slower than using data that’s already in smaller, closer, super-fast memory right next to the MXU (called VMEM — think of it as the chef’s countertop). One of the TPU compiler’s main jobs is keeping the MXU’s countertop stocked so the machine never has to wait on the pantry.

ICI (Inter-Chip Interconnect). The conveyor belts between kitchens. When you have multiple TPU chips working together, they need to talk to each other fast. ICI is the dedicated high-speed wiring that connects neighboring chips. Most modern TPU systems arrange these connections in a torus — a grid where chips connect to their neighbors and the edges wrap around — because it’s a topology that works well for the communication patterns training needs. On NVIDIA GPUs, the equivalent is called NVLink. We won’t need this until Part 6, when we train across four chips at once.

Keep this picture in mind: large matrix engines, small helper engines, a big pantry, and wires out to neighboring chips.

Two things worth mentioning while we’re here, since beginners often trip on them:

• TPUs don’t work alone. Every TPU chip is paired with a host CPU that runs your Python, prepares batches, and orchestrates work. When you “SSH into a TPU,” you’re actually SSHing into that host CPU; the TPU chips are attached to it. We’ll see this in Part 2.
• TPUs come in generations. You’ll hear about v2, v3, v4, v5e, v5p, v6e. They differ in speed, memory, and interconnectivity. For this series, we’ll use v6e (“Trillium”), Google’s current-generation chip. The mental model in this article applies to all of them; the specific numbers don’t.

Side-by-side: GPU vs TPU, row by row

Okay — here’s the comparison table you came for. But I’m going to walk through every row, not just drop it on you. The table is a summary; the explanations below are where the actual learning happens.

Now let’s unpack each row.

Matrix engine: many small vs one big

A GPU spreads its matrix-multiplication hardware across the chip, many small Tensor Cores, each working on a slice of the problem. It’s like having 30 cooks each rolling out one ball of dough. Good news: they’re flexible, and if you hand them unusually-sized jobs, they’ll cope. Bad news: there’s overhead in handing work out, tracking who’s doing what, and collecting results.

A TPU concentrates its matrix-multiplication power into a small number of very large engines that eat whole matrices in one go. In the analogy, it’s a pasta machine instead of thirty cooks, or really, a couple of pasta machines, but the principle is the same: each one is enormous and handles whole sheets of dough rather than parceling out little balls. Good news: very little coordination overhead, and incredibly fast when matrices are big. Bad news: if you hand it a tiny job, most of the machine sits idle. A 100×100 matrix is a waste; a 10,000×10,000 matrix is heaven.

Practical implication: TPUs love big, uniform, predictable workloads. GPUs are more forgiving of variety.

The “everything else” engine

Neural networks aren’t only matrix multiplications. There’s also relu(x), x + bias, mean(x), and a long tail of little operations. Both chips need something to do those.

On a GPU, the same thousands of general cores that predate AI handle this. They’re good at it; that’s literally what they were built for.

On a TPU, there’s a dedicated smaller unit, the VPU. It’s less flexible, but it’s designed to hand data seamlessly between itself and the MXU. There’s no round-trip to the pantry.

Practical implication: for normal deep-learning workloads, both work fine. If you’ve written a weird custom operation (some exotic loss function, maybe), the GPU will usually handle it with less fuss; the TPU might force you to rephrase it as a sequence of standard ops.

Fast on-chip memory

Every chip has two kinds of memory: the big-but-slow main memory (the pantry), and the small-but-fast memory right next to the calculators (the countertop). The fast one is where you want your data when you’re actually doing math.

Get Roya’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

GPUs have a little of this fast memory attached to each cluster of cores, like each group of cooks having their own prep space. TPUs have one larger shared pool, one giant countertop next to the pasta machine.

Practical implication: the TPU arrangement matches the TPU’s philosophy, one big engine, one big countertop stocked for it. The GPU’s arrangement matches its philosophy: many cooks, each with their own prep space.

Chip-to-chip wiring

When you need more compute than one chip can provide (anything approaching LLM-scale), you need chips to talk to each other. Both GPUs and TPUs have proprietary high-speed wiring.

The interesting architectural choice: many TPU systems connect chips in a torus. Imagine chips arranged in a grid where every chip connects to its four neighbors (up, down, left, right), and the edges wrap around so the left-edge chips also talk to the right-edge chips. This topology works well for the communication patterns deep-learning training needs. The exact shape varies by TPU generation (v4 uses a 3D torus, smaller configurations use simpler meshes or rings), but the idea of a dense, regular grid of chip-to-chip links is consistent. NVIDIA’s NVLink connects chips in different arrangements depending on the server.

Practical implication: Both work well for multi-chip training. TPUs were designed from the start to scale to thousands of chips, and Google’s internal training clusters can hit scales that are genuinely hard to match elsewhere.

bfloat16: the quietly enormous thing

This one deserves more than a line. “bfloat16” is a 16-bit number format. You’ve probably heard of fp32 (standard 32-bit floats) and fp16 (16-bit floats). bfloat16 is a third option.

Why does it matter? Two reasons:

1. Smaller numbers mean faster math and less memory. A bf16 matrix is half the size of an fp32 one. That means it moves through memory twice as fast and fits twice as much on the chip.
2. bf16 is “safer” than fp16. Older fp16 has a narrow range of representable numbers , it can overflow or underflow during training, and people had to invent workarounds (like “loss scaling”) to deal with that. bf16 keeps the full range of fp32, at the cost of some precision in the smaller digits. For neural-network training, that tradeoff is almost always a win.

TPUs were designed around bf16 from day one. Their matrix engine takes bf16 inputs natively and accumulates results into fp32 for numerical stability. If you keep everything in fp32 on a TPU, you’re voluntarily running at half speed.

GPUs support bf16 too on newer chips (Ampere and later—think A100, H100). On older GPUs like the T4 we’ll use for comparison in Part 4, bf16 isn’t supported in hardware, so it gets emulated in software and is actually slower than fp32. The comparison gets interesting.

Practical implication: if you’re on a TPU, use bf16. It’s the point.

How you write code (this is the biggest one)

This one is subtle but it changes how you work day-to-day.

On a GPU with PyTorch, the default flow is “eager.” Each line of Python runs immediately. You write x = model(x), and the GPU starts working right away. You can drop a print(x) anywhere. You can step through code with a debugger. If your program crashes, the stack trace points at the Python line that broke.

On a TPU with JAX, the default flow is “compile then run.” You wrap your function in a @jit decorator. The first time you call it, JAX traces through the entire function, hands it to a compiler called XLA, and gets back one big optimized program. From then on, every call of that function runs that pre-compiled program.

Worth saying out loud: the line between the two worlds is fuzzier than it used to be. PyTorch now has torch.compile, CUDA graphs, and runs on XLA too, so GPUs can be driven in "compile first" style when you want. And JAX can be used eagerly when you're debugging. The contrast is about the default workflow, not a hard wall.

Pros of the TPU/JAX approach:

• The compiler sees the whole computation and can optimize aggressively, combine operations together, overlap memory moves with math, and pick smart memory layouts.
• Repeated calls are blazingly fast because there’s no Python overhead between operations.

Cons:

• The first call is slow — it’s compiling.
• Debugging inside the compiled function is harder. You can’t just drop a print there.
• The function has to be “pure” — it can’t have side effects or depend on anything that might change unpredictably.

Practical implication: the GPU/PyTorch world is great for rapid experimentation and debugging. The TPU/JAX world is great for running the same training step millions of times as fast as physically possible. Most real ML teams end up doing their prototyping in PyTorch and their big production runs in whichever stack is cheaper and faster.

So which one should I use?

Here’s a cheat sheet. Bookmark this.

A practical framing instead of a verdict: for very large, uniform workloads, LLM pretraining, big-batch inference, anything that’s a tall stack of predictable matmuls, TPUs are often more cost-efficient. For interactive research, odd model architectures, or anything that benefits from eager-mode debugging, GPUs are often more productive. Real cost and performance depend on model architecture, batch size, memory bandwidth, cluster topology, software stack, and pricing, so “it depends” is the honest answer, and most serious teams have access to both.

One clarification the analogy risks obscuring: TPUs aren’t only for transformers. Google uses them for CNNs, diffusion models, vision models, reinforcement learning, and recommendation systems. “Specialized” doesn’t mean “single-purpose” — it means “optimized for the math deep learning tends to do,” which covers a lot of ground.

Why JAX specifically?

You can technically run PyTorch on TPUs (there’s a project called PyTorch/XLA that makes it work), but the combination most people actually reach for is JAX on TPU. Three reasons:

1. JAX was built for the compiler. JAX’s design , @jit, grad, vmap, pmap, all these composable function transformations we'll meet in Part 3 , maps directly onto what XLA wants to see. The library and the compiler were designed to fit each other.
2. It’s widely used in Google Research and DeepMind. Google is actually multilingual internally; they also use TensorFlow, PyTorch, and various internal frameworks, but JAX has become the dominant language of the research groups producing a lot of the most interesting public AI work. If you want to read and reproduce that research, you’re going to see JAX.
3. It scales cleanly. Going from one chip to four chips to sixty-four chips, in JAX, is often one-line changes. In other frameworks, it’s a major project.

The learning curve is real, JAX asks you to think differently than PyTorch does, but the payoff is worth it, and building up to it is exactly what the next five parts do.

What’s coming

• Part 2: Getting a TPU and your first JAX program. Provision a TPU on Google Cloud (including how to deal with the infamous “no capacity” error), connect to it, run your first jnp operation.
• Part 3: JAX for PyTorch people. The seven core JAX concepts, each mapped back to its PyTorch equivalent so nothing feels foreign.
• Part 4: Inside the TPU. We’ll actually look at what @jit produces — the real compiled program. Then benchmark a TPU against a real GPU running the same code, with real numbers.
• Part 5: Your first real model. Train MNIST end-to-end on TPU. 97.6% accuracy in under 4 seconds.
• Part 6: Scaling out and shutting down. Training across multiple TPU chips, plus a cost audit so your cloud bill doesn’t surprise you.

What you’ll need to follow along

• A Google Cloud account with billing enabled. A fresh account gets $300 in free credits, which is way more than you need.
• A laptop with a terminal. Any OS.
• About 2–3 hours of reading. Add another hour or two if you run everything yourself.

Expected total cost if you run every script on Spot-priced TPUs (we’ll explain what those are in Part 2): $1.50 to $3.00. Less than a coffee, and you’ll walk away with a real understanding of what a TPU is and how to use one.

(Cost caveat: prices vary by region, availability, and TPU generation, and Spot pricing fluctuates. That estimate assumes the smallest TPU instances on Spot pricing in the region we’ll use. Your mileage will vary a little either way.)

Onward to Part 2, where we’ll get you a real TPU and run your first JAX program on it.