TPU 101 — Part 2: Getting a TPU and Running Your First JAX Program

A beginner-friendly series on Google TPUs and JAX. This part: provisioning a TPU on GCP, installing the tools (with every real error you’ll hit), and running your first import jax on actual TPU hardware.

Part 1 laid out the why. Now the how. By the end of this part you’ll have:

1. Your Google Cloud project set up correctly (billing, the right API, the right permissions).
2. A real TPU VM running in that project.
3. gcloud installed on your machine (or you'll be using Cloud Shell instead, both work).
4. A working SSH connection into the TPU.
5. Your first JAX program confirming the accelerator is alive.

I’ll warn you upfront: provisioning cloud infrastructure is where tutorials usually fall apart. The commands look simple, but cloud capacity, region availability, and install quirks can eat an afternoon. I’ll flag every error I actually hit while writing this, so you don’t have to discover them yourself.

Before you begin

Before any of the commands below will work, you need:

• A Google Cloud project with billing enabled. If you’ve never used GCP, sign up at console.cloud.google.com, fresh accounts get $300 in free credits.
• The Cloud TPU API and the Compute Engine API enabled on your project. The TPU API is the obvious one; the Compute Engine API matters too because TPU VM commands are managed through the broader gcloud compute surface. You can enable both in the console, go to console.cloud.google.com/apis/library, search for each one, and click Enable, or, once you've set PROJECT_ID (in Step 3), enable both in one CLI command:
• gcloud services enable tpu.googleapis.com compute.googleapis.com \ --project="$PROJECT_ID"
• If you skip these, the error you eventually get may not clearly say “enable the Cloud TPU API” , it might be a generic permission-denied message , so it’s worth doing upfront.
• Permission to create TPU VMs and SSH into them. If this is your own personal project, the Owner role you got automatically when you created the project covers everything. If this is a company or university project, ask whoever administers it to grant you these three IAM roles on the project:
• TPU Admin (roles/tpu.admin) , to create, delete, and describe TPUs.
• Service Account User (roles/iam.serviceAccountUser) ,TPU VMs run under a service account; you need permission to "act as" it.
• Compute Viewer (roles/compute.viewer), read access to compute resources, used during SSH.
• These are the three roles the official Google Cloud TPU + JAX setup guide lists, so they’re the canonical answer if your admin asks “which exactly?”
• Either Cloud Shell access (in your browser, no install required, Step 2 explains how to use it) or a local gcloud installation (Step 2 walks through this too).

If you skip any of the above, the very first command in Step 3 will fail in a way that doesn’t make it obvious which prerequisite is missing. Worth ten minutes upfront to save an hour of debugging.

Step 1 — Understand the “No capacity” problem before you hit it

The first command we’ll run tries to create a TPU. The first error you might see looks like this:

ERROR: ... there is no capacity ...

If you come from AWS, your instinct is “it’s a quota issue, let me request more.” If the message explicitly says there is no capacity, that’s a different problem , it usually means Google’s available pool for that TPU type in that zone is temporarily full. That’s distinct from quota or IAM errors, which usually mention “quota,” “permission,” “roles,” or “access” by name. The fix for “no capacity” is to wait, try a different zone, or use Spot , none of which would help with quota or IAM.

Three ways to work around a real “no capacity” error:

1. Wait and try again, capacity comes and goes hour to hour.
2. Try a different zone , us-central1-b full? Try us-east5-a.
3. Use Spot pricing, Spot is dramatically cheaper than on-demand (commonly 50–80% off, depending on TPU generation and region) and sometimes has different availability characteristics, but it’s not a magic bypass for capacity issues: Spot and on-demand often draw from related capacity pools, so a “no capacity” error on on-demand frequently means Spot is tight too. Spot also doesn’t bypass quota or IAM problems , those are separate failure modes you’d need to fix separately. The catch with Spot: Google can preempt (shut down) your VM at any time. For a learning session, that’s fine; just don’t rely on Spot for anything that needs to stay up for days without checkpointing. And don’t treat preemption as a cost safety net ,Spot VMs can also happily run for a long time without being preempted, so you still need to delete yours when you’re done.

We’ll use Spot.

Step 2 — Install gcloud on your Mac (or skip this entirely)

You have two options. Pick one and use it for the rest of this tutorial , every later step assumes you have a working terminal with gcloud available, and either option gives you that. They're functionally equivalent; the only difference is whether the terminal lives in your browser (Option A) or on your laptop (Option B).

If you’re not sure, pick Option A. Cloud Shell is zero-install, pre-authenticated, and you can always set up the local install later.

A quick note if you’re on Linux or Windows: the local-install walkthrough below is macOS-specific. The same family of official installer works on Linux , see the official Google Cloud CLI install for Linux (one-liner curl + a shell script, takes 2 minutes). If you're on Windows, use WSL2 and follow the Linux instructions, or use the Windows installer. In all cases, skip to Step 3 once gcloud --version works. Or , easiest of all , use Option A below and don't install anything.

Option A: Skip the install, use Cloud Shell

If you just want to follow along, you don’t need gcloud on your Mac at all. Google Cloud Shell is a Linux terminal in your browser with gcloud already installed and authenticated.

1. Go to console.cloud.google.com and log in.
2. Make sure the project selector at the top shows your project.
3. In the top-right, click the >_ icon ("Activate Cloud Shell"). A terminal opens at the bottom of the browser.

Done. Skip to Step 3.

Option B: Install gcloud locally

Nicer long-term, gcloud lives on your laptop and works without a browser tab. The cleanest way to install it on macOS is Google's own installer script, which can install or use a supported Python for the SDK and doesn't depend on Homebrew or whatever Python setup you already have:

curl https://sdk.cloud.google.com | bash

(A note on curl ... | bash: piping a downloaded script straight into bash is a pattern security-conscious folks reasonably raise eyebrows at, because you're trusting the script wasn't tampered with in transit. The mitigation is that sdk.cloud.google.com is served over HTTPS from Google's own domain. For a personal learning machine, this is fine; if you'd rather audit first, run curl -o install.sh https://sdk.cloud.google.com, read the script, then bash install.sh.)

The installer will ask a few questions:

• “Do you want to help improve the Google Cloud CLI (y/N)?” — answer however you want; it’s just usage telemetry.
• “Installation directory” — accept the default ($HOME), which puts the SDK at ~/google-cloud-sdk.
• It may prompt for your Mac password if it needs to install or configure dependencies (such as a supported Python interpreter for the SDK). Heads up: you won’t see characters as you type; that’s normal.
• “Modify profile to update your $PATH and enable shell command completion?” — answer Y. This is what makes gcloud available in new terminal windows.
• “Enter a path to an rc file to update” — accept the default (usually ~/.zshrc on modern macOS).

The install takes a few minutes, it’s downloading and setting up the SDK and its dependencies in ~/.config/gcloud/virtenv.

When it finishes, open a new terminal window (the simplest way to pick up the PATH update). Or if you want to stay in the current terminal:

exec -l $SHELL

That re-execs your shell, so it re-reads .zshrc and sees the updated PATH. Either way, verify:

which gcloud
gcloud --version

You should see /Users/yourname/google-cloud-sdk/bin/gcloud and a version banner like Google Cloud SDK 5xx.0.0.

Log in and pick your project:

gcloud init

• Choose “Log in with a new account”, a browser opens.
• Log in with your Google account.
• When asked which project to use, pick yours.
• Skip the default region prompt (just press Enter).

You only do this once per machine.

Forward reference: gcloud init is enough auth for everything in this part, and for Parts 3–6 of the series. The moment you want to read datasets or write checkpoints to a GCS bucket from inside Python, you'll also need gcloud auth application-default login, that's a separate auth flow that gives client libraries (not just the gcloud CLI) credentials. We don't need it yet, but if you see DefaultCredentialsError later, that's the fix.

Things that can go wrong

A few things to watch for, drawn from real install attempts:

**which gcloud returns "not found" right after the install finished.** Your current shell still has the old PATH. The installer added the new entries to .zshrc but only newly-opened shells pick that up. Fix: open a new terminal window, or run exec -l $SHELL in the current one.

The installer prompts for your password. That’s the installer running a macOS pkg installer for a dependency it needs (typically a supported Python interpreter). If you’re on a managed company laptop where you don’t have admin rights, you have two options: ask IT to grant local-admin temporarily, or use Cloud Shell (Option A) instead, it’s functionally equivalent for this whole series.

You’d rather use Homebrew because the rest of your tooling is brewed. There’s a google-cloud-sdk cask (brew install --cask google-cloud-sdk), but it's been intermittently broken on macOS, at the time of writing, on a clean install, it can fail with virtualenv: command not found or with pyexpat/libexpat symbol-mismatch errors stemming from how the cask depends on Homebrew Python. The official installer above sidesteps all of that by bringing its own Python. If you have a working brew cask install already, great, leave it alone, but for a fresh setup, the curl ... | bash path is currently the more reliable one.

Step 3 — Create the TPU

Where to run this: in whichever shell you set up in Step 2, either Cloud Shell (in your browser) or your local terminal with gcloud installed. Both work identically for creating the TPU; pick whichever one you already have open. Cloud Shell is the safer default if you didn't install gcloud locally , it's already authenticated, and there's nothing to configure.

A small note about Cloud Shell specifically: sessions can disconnect when idle, and they’re capped at 12 hours total. If your browser tab sits idle while you’re reading, the session may disconnect, and your exported shell variables ($ZONE, $PROJECT_ID, etc.) are lost. The TPU itself keeps running, it's not affected, and your gcloud config set project setting also persists. Only the shell variables disappear. If that happens, just paste the export block again before running any later commands.

Before the create command, do two things: one persistent, one for the current session.

One-time, persistent: tell gcloud which project you're working in. This is saved to your gcloud config and survives across reboots, shell restarts, and Cloud Shell timeouts. You only ever do this once per machine per project:

gcloud config set project paste-your-real-project-id-here

Replace paste-your-real-project-id-here with your actual GCP project ID. You can find it in the top bar of the Cloud Console , it usually looks like my-project-name-123456.

Per-session: export shell variables we’ll reuse for every later command. export makes them visible to subprocesses (including scripts you run later), and they live as long as the current terminal session does:

export PROJECT_ID=paste-your-real-project-id-here    # ← actually replace this
export TPU_NAME=jax-tpu
export ZONE=us-central1-b
export ACCELERATOR_TYPE=v6e-1
export RUNTIME_VERSION=v2-alpha-tpuv6e

(If you close the terminal or your Cloud Shell session times out, you’ll need to re-run this export block. The gcloud config set project from above doesn't go away , only the shell variables do.)

Critical: the PROJECT_ID line is a placeholder. You have to replace paste-your-real-project-id-here with your actual GCP project ID , the same one you used in gcloud config set project above. If you leave the placeholder text in, the next command will fail with The project property must be set to a valid project ID, not the project name [paste-your-real-project-id-here] , which is gcloud's way of telling you that the string isn't a real project.

Sanity check before continuing , paste this and confirm it prints your actual project ID, not the placeholder:

echo "PROJECT_ID=$PROJECT_ID, ZONE=$ZONE"

You should see something like PROJECT_ID=my-project-name-123456, ZONE=us-central1-b. If PROJECT_ID is empty or still says paste-your-real-project-id-here, the variables didn't get set , re-run the export block above. (Common mistake: writing echo PROJECT_ID=foo only prints the text; it doesn't set anything. Use export PROJECT_ID=foo. Technically, PROJECT_ID=foo on its own is enough for commands typed in the same shell , shell expansion happens before gcloud runs , but export is the safer habit because scripts and subprocesses can see it too.)

Now create the TPU:

gcloud compute tpus tpu-vm create "$TPU_NAME" \
  --zone="$ZONE" \
  --accelerator-type="$ACCELERATOR_TYPE" \
  --version="$RUNTIME_VERSION" \
  --project="$PROJECT_ID" \
  --spot

Let me break down each flag , these aren’t obvious:

• **$TPU_NAME** (jax-tpu) , just a name for the VM. Call it whatever you want.
• **$ZONE** (us-central1-b), the GCP zone. TPU availability varies by zone us-central1-b is usually a safe bet for v6e.
• **$ACCELERATOR_TYPE** (v6e-1) , requests the smallest v6e slice: one v6e chip attached to one TPU VM. In v6e (Trillium), each chip contains one TensorCore with multiple MXUs plus vector and scalar units , which ties back to the architecture sketch in Part 1. v6e-1 is a good fit for testing and learning because it's the smallest v6e slice. You can ask for -4, -8, etc., for bigger slices, but one chip is plenty for now and much cheaper. (v6e is the current widely-available generation as of writing; newer generations are launching with more emphasis on GKE and queued-resource workflows, but the TPU VM workflow used here is still the right mental model for this beginner tutorial.)
• **$RUNTIME_VERSION* (v2-alpha-tpuv6e) — this is the runtime image* that ships on the VM. It has to match the TPU generation, or things will silently not work. Different TPU generations want different runtime strings, and Google occasionally moves images out of alpha or renames them. If the value below fails with an error about an unknown runtime version, check the current Cloud TPU runtime versions docs for the latest string for v6e + JAX. (As of writing: v2-alpha-tpuv6e for v6e, v2-alpha-tpuv5-lite for v5e, v2-alpha-tpuv5 for v5p, tpu-ubuntu2204-base for v4 , these can shift.)
• **--spot** — use Spot pricing. It's much cheaper than on-demand, but availability is still dynamic, and quota/IAM rules still apply.

Note on older tutorials: the command above uses gcloud compute tpus tpu-vm, which is the modern TPU VM architecture. A lot of older blog posts and Stack Overflow answers use plain gcloud compute tpus create (no tpu-vm), that's the legacy "TPU node" architecture, where the TPU is a separate device you talk to over the network from a different VM. Don't mix the two; if you copy a command from an older post and it doesn't have tpu-vm in it, ignore that post.

If the first zone is full

If the first zone returns a “no capacity” error, here’s a loop that tries other zones in order and reassigns ZONE to whichever one actually worked, so every later command in this tutorial keeps using the right zone with no extra effort from you:

# Stick to v6e zones — mixing in other generations would break later steps
# (device output, performance numbers, and the runtime image string all
# assume v6e).
ZONES_TO_TRY=(us-central1-b us-east5-a)
# TPU zone availability changes over time. If both zones below fail,
# check Google's current TPU regions/zones page
# (https://cloud.google.com/tpu/docs/regions-zones) and add another
# v6e-supported zone to the list.
for Z in "${ZONES_TO_TRY[@]}"; do
  echo "Trying $Z..."
  if gcloud compute tpus tpu-vm create "$TPU_NAME" \
    --zone="$Z" \
    --accelerator-type="$ACCELERATOR_TYPE" \
    --version="$RUNTIME_VERSION" \
    --project="$PROJECT_ID" \
    --spot; then
      echo "Created TPU in zone: $Z"
      export ZONE="$Z"      # update the exported ZONE for later commands
      break
  fi
done

The if … then break means: if the create command succeeds, update ZONE and stop looping. Because every later command in this tutorial uses "$ZONE", they'll automatically point at the right place, no need for you to remember which zone you ended up in.

If every zone in your list fails on the same day, that’s most often a transient capacity issue. Wait an hour and try again, or expand the list with another v6e-supporting zone from Google’s TPU regions/zones page. And remember: Spot is much cheaper than on-demand but doesn’t bypass quota or IAM problems, so if the failures look like permission or quota errors rather than capacity ones, fix those first (see Step 1).

Confirm the TPU is ready

gcloud compute tpus tpu-vm describe "$TPU_NAME" \
  --zone="$ZONE" \
  --project="$PROJECT_ID" \
  --format="value(state)"

You want it to print READY. If it says CREATING, wait 30 seconds and try again.

ERROR: TPU node jax-tpu already exists

If you re-run the create command after a partial success (or after closing your terminal mid-create), you may see:

ERROR: TPU node jax-tpu already exists

This means a TPU with that name is still around in the project. Don’t try to create another one , describe it first to see what state it’s in:

gcloud compute tpus tpu-vm describe "$TPU_NAME" \
  --zone="$ZONE" \
  --project="$PROJECT_ID"

If it’s READY, great, you can skip the create step entirely and SSH straight in (Step 4). If it's stuck or you don't recognize it, delete it cleanly and recreate:

gcloud compute tpus tpu-vm delete "$TPU_NAME" \
  --zone="$ZONE" \
  --project="$PROJECT_ID"

Then re-run the create command from above.

Step 4 — SSH into the TPU

Here’s a thing most people don’t realize: a TPU VM has the TPU chips attached directly to it. You don’t SSH into a host and then connect to a separate TPU over the network — the VM is the TPU’s host.

This was different a few years ago, and a lot of stale tutorials still describe the old architecture. Ignore them.

SSH in:

gcloud compute tpus tpu-vm ssh "$TPU_NAME" \
  --zone="$ZONE" \
  --project="$PROJECT_ID"

What you’ll see the very first time

If this is the first time you’ve ever SSH’d into any Google Cloud VM from this machine, gcloud will generate an SSH key pair for you. It saves the keys to ~/.ssh/google_compute_engine (private) and ~/.ssh/google_compute_engine.pub (public), then uploads the public key to the TPU. You'll see a couple of prompts:

Generating public/private rsa key pair.
Enter passphrase (empty for no passphrase):

Just press Enter — don’t type a passphrase. Then it asks again to confirm:

Enter same passphrase again:

Press Enter again. This creates an SSH key with no passphrase, which is convenient for a short-lived learning TPU. (A passphrase is a password that locks the key file itself; if you set one, you’ll typically be prompted for it on every connection unless you set up an ssh-agent to cache it — which is more setup than this tutorial covers.) For long-lived or production environments, follow your organization's SSH key policy — they'll usually require a passphrase, and many shops also use a separate identity-management system (like OS Login or IAP-tunneled SSH) instead of long-lived per-user keys.

After that, gcloud uploads the public key to the TPU. The first time can take 30–60 seconds. Subsequent SSH connections are instant.

What you’ll see on subsequent connections

After the first time, gcloud reuses the key it generated. You’ll just see a brief “Propagating SSH public key…” message and then drop straight into the TPU’s shell. Your terminal prompt will change to something like roya@t1v-n-abcd1234-w-0:~$ — that remote-looking hostname means you're on the TPU, not your laptop.

“Enter passphrase for key” — what to do

If your SSH attempt asks for a passphrase like this:

Enter passphrase for key '/Users/your-name/.ssh/google_compute_engine':

…it means you (or an old gcloud invocation, possibly years ago) set a passphrase on this key in the past. If you remember it, type it and press Enter. If you don't — easiest fix is to delete the keys and let gcloud generate fresh ones with no passphrase:

1. Press Ctrl+C to cancel the passphrase prompt.
2. Delete the existing keys:

• rm ~/.ssh/google_compute_engine ~/.ssh/google_compute_engine.pub

1. Re-run the SSH command. When prompted “Enter passphrase (empty for no passphrase),” press Enter twice.

For most personal learning setups, this is safe, these keys are used by gcloud for connecting to GCP VMs, and deleting them just means gcloud regenerates them on next use. If you use the same key for other workflows or are on a managed company machine, check before deleting it.

Other things that can go wrong

**Permission denied (publickey) on your very first try.** The TPU finished CREATING but the key upload to it is still propagating. Wait 30 seconds and retry, usually works on the second attempt.

**Connection refused.** The TPU exists, but its SSH service isn't up yet. Wait another 30 seconds.

Step 5 — Install JAX and run your first program

You’re now on the TPU. Install JAX with TPU support:

python3 -m pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

This is the install command from Google’s official JAX-on-Cloud-TPU guide. A few things worth understanding here:

• **python3 -m pip* instead of just pip , the TPU VM image ships with both python3.10 and python3.11, and a bare pip might install into a different* Python than the one python3 runs. Using python3 -m pip is "use the pip that belongs to the same Python you'll use to run scripts," which guarantees they match. This avoids the classic gotcha where pip says it installed JAX successfully, but then python3 hello_jax.py says ModuleNotFoundError: No module named 'jax'.
• The [tpu] extra installs the TPU runtime dependencies.
• The -f flag points pip at the JAX TPU/libtpu release index , pip needs that to find the right matched versions of jax, jaxlib, and libtpu for the TPU runtime.

If you skip the -f flag and just install plain jax, JAX may only see the CPU or fail to initialize the TPU backend.

Create the test script on the TPU

You’re now on a remote Linux shell, not in your local IDE. So we need to create the Python file on the TPU itself. There are two easy ways. Pick whichever you’re more comfortable with.

Option 1: nano (most beginner-friendly)

nano is a built-in terminal text editor . It shows the keyboard shortcuts at the bottom of the screen, so you don't need to memorize anything. Open a new file:

nano hello_jax.py

A blank editor will open. Paste in the Python code below (Cmd+V on Mac terminal, or right-click → Paste in most terminals). Then save and exit:

• Save: Press Ctrl+O, then Enter to confirm the filename.
• Exit: Press Ctrl+X.

You’ll be back at the shell prompt.

Option 2: heredoc (one-shot, no editor needed)

If you’d rather skip the editor entirely, you can write the whole file in one paste. Copy this entire block , opening cat <<'EOF' line through the closing EOF line , and paste into the TPU shell:

cat > hello_jax.py <<'EOF'
import jax
import jax.numpy as jnp
print("jax.__version__ :", jax.__version__)
print("jax.devices()   :", jax.devices())
print("device count    :", jax.device_count())
print("default backend :", jax.default_backend())
x = jnp.arange(8)
y = x * 2
print("x               :", x, "dtype=", x.dtype)
print("y = x * 2       :", y)
a = jnp.ones((1024, 1024))
b = jnp.ones((1024, 1024))
c = (a @ b).block_until_ready()
print("1024x1024 matmul OK, c[0,0] =", float(c[0, 0]))
EOF

The <<'EOF' ... EOF syntax means "everything between these markers becomes the content of hello_jax.py." The single quotes around 'EOF' matter, they tell the shell not to try to interpret things like $x inside the script as shell variables.

Either way, confirm the file exists:

ls -l hello_jax.py
cat hello_jax.py

You should see the file listed and the Python code printed back.

What’s actually in the file

Here’s that script, in one block, so you can read it without the heredoc wrapping:

import jax
import jax.numpy as jnp
print("jax.__version__ :", jax.__version__)
print("jax.devices()   :", jax.devices())
print("device count    :", jax.device_count())
print("default backend :", jax.default_backend())
x = jnp.arange(8)
y = x * 2
print("x               :", x, "dtype=", x.dtype)
print("y = x * 2       :", y)
a = jnp.ones((1024, 1024))
b = jnp.ones((1024, 1024))
c = (a @ b).block_until_ready()
print("1024x1024 matmul OK, c[0,0] =", float(c[0, 0]))

Before we run it, let me point out the non-obvious lines:

• **import jax.numpy as jnp** — JAX provides a NumPy-like API. Arrays are placed on JAX devices, and most normal jnp operations get dispatched to the active backend (TPU here) and can be traced by @jit/grad/etc. A handful of operations may fall back to CPU or trigger transfers in edge cases, but for everything in this tutorial, "your jnp calls run on the TPU" is a fine mental model.
• **jax.devices()** — lists every accelerator JAX can see. On a TPU VM you want this to print something containing TpuDevice(...).
• **jax.default_backend()** — which backend JAX will target by default. Should print tpu.
• **jnp.arange(8)* — notice it creates an int32 array. JAX arrays default to 32-bit precision (int32 for ints, fp32 for floats). On TPU, though, matrix multiplication itself uses bf16 inputs with fp32 accumulation by default , the dtype of the array and the dtype the MXU uses internally* aren't the same thing. In Part 4 we'll make dtype and precision explicit so you can see the performance and accuracy tradeoffs directly.
• **(a @ b).block_until_ready()* — this is the one line that most surprises PyTorch users. JAX dispatches operations asynchronously*: the Python line returns before the accelerator has actually finished the work. block_until_ready() forces Python to wait. You need this whenever you're timing things or printing a result you want to be sure is computed.

Run it:

python3 hello_jax.py

You should see output similar to:

jax.__version__ : 0.x.x
jax.devices()   : [TpuDevice(...)]
device count    : 1
default backend : tpu
x               : [0 1 2 3 4 5 6 7] dtype= int32
y = x * 2       : [ 0  2  4  6  8 10 12 14]
1024x1024 matmul OK, c[0,0] = 1024.0

The exact formatting varies by JAX version and TPU generation , newer JAX releases print device info differently than older ones. The important checks are:

• jax.devices() includes TpuDevice (not CpuDevice, not GpuDevice)
• jax.default_backend() prints tpu
• the matrix multiplication completes and c[0, 0] is 1024.0

If you see those three things, you’re successfully running on TPU silicon.

If your output for jax.devices() includes fields like coords=..., those describe where the device sits in the TPU topology. For a one-chip tutorial, you can ignore those coordinates — they only become useful when we scale across multiple chips in Part 6.

PyTorch → JAX cheat sheet (for this part)

Things that go wrong

A short catalog of errors you might hit after everything above, with fixes:

**No GPU/TPU found, falling back to CPU** — your runtime image doesn't match the TPU generation. Confirm --version=v2-alpha-tpuv6e for v6e. You'll need to recreate the VM with the correct version.

**ModuleNotFoundError: No module named 'libtpu'** — you installed jax without the [tpu] extra, or without the -f libtpu releases URL. Rerun the full install command from Step 5.

SSH-related errors — see Step 4 for Permission denied (publickey), Connection refused, and the "Enter passphrase" prompt.

Multiple Python versions are confusing me — use python3 explicitly. The default TPU VM image ships both python3.10 and python3.11.

Before you walk away from your laptop

This is a big deal. A running TPU costs money. Cloud TPU is billed per chip-hour, and Spot prices change over time, so always check the Cloud TPU pricing page for current rates. At the time of writing, a small v6e-1 on Spot costs well under a dollar per hour in the regions we use, but that adds up fast if you leave it running, even at a hypothetical $0.40/hour, a forgotten week (168 hours) is $67. Your real rate may be higher; check the pricing page.

Step A: Exit the TPU first

If you’ve been working on the TPU shell, your prompt will look something like roya@t1v-n-abcd1234-w-0:~$ , that remote-looking hostname means you're still on the TPU. The delete command needs to run from your local shell (the one with the $PROJECT_ID, $TPU_NAME, $ZONE variables you exported earlier), not from inside the TPU.

Exit back to your laptop:

exit

Your prompt should change back to whatever your laptop or Cloud Shell prompt looks like (e.g., (base) roya@Royas-Laptop ~ %). If it still looks like a remote machine prompt, run exit again until it's gone.

Step B: Verify your variables are still set

Quick sanity check before deleting , make sure the variables from Step 3 are still defined in your current shell:

echo "PROJECT_ID=$PROJECT_ID, TPU_NAME=$TPU_NAME, ZONE=$ZONE"

If anything is empty (Cloud Shell may have timed out, or you opened a new terminal), re-export them:

export PROJECT_ID=your-real-project-id
export TPU_NAME=jax-tpu
export ZONE=us-central1-b   # or whichever zone the TPU actually lives in

Why this matters: if you accidentally run the delete command on the TPU shell or in a fresh terminal where the variables aren’t set, gcloud will give you The project property is set to the empty string, which is invalid. That's not a real failure , your TPU is still alive and still costing money , it just means the variables didn't expand to anything. Verify, then delete.

Step C: Delete the TPU

gcloud compute tpus tpu-vm delete "$TPU_NAME" \
  --zone="$ZONE" \
  --project="$PROJECT_ID"

It’ll ask “Do you want to continue (Y/n)?” , type Y and press Enter. Takes about 30 seconds.

Step D: Confirm it’s gone

gcloud compute tpus tpu-vm list --zone="$ZONE" --project="$PROJECT_ID"

If this prints Listed 0 items. (or no rows other than the header), you're clean , no more billing for this TPU.

Don’t worry , you can create another one in 30 seconds whenever you want to keep going (just re-run the create command from Step 3). In Part 6 I’ll show you how to set up a billing budget alert so this is guaranteed never to bite you.

You now have working TPU infrastructure and a successful jnp call. In Part 3, we'll dive into JAX itself — how it differs from PyTorch, the seven core concepts, and how to write code that makes XLA happy.