TPU 101 — Part 3: JAX for PyTorch People

A beginner-friendly series on Google TPUs and JAX. This part: seven core JAX concepts, each mapped back to its PyTorch equivalent, so the transition feels like a nudge rather than a leap.

In Part 2 you got a TPU running and printed your first TpuDevice(...). Good. Now we learn JAX properly before we try to do anything real with it.

Here’s the thing about JAX: it looks a lot like NumPy, which tempts PyTorch users to treat it as a drop-in replacement. It isn’t. JAX is a functional library built around seven core ideas, and once you internalize those ideas the rest of the library falls into place. If you don’t internalize them, you will spend a lot of afternoons confused about error messages.

Where to run this part: Part 3 is conceptual — none of the snippets need TPU hardware to teach what they’re teaching. If you deleted your TPU at the end of Part 2 (good, you saved money), you can read this part as-is, or run the snippets on Colab’s free CPU runtime, or any local Python install with pip install jax. The benchmark numbers I report were measured on a v6e-1, so your timings on CPU will look different — but the lessons (immutability, the jit speedup pattern, key-splitting, etc.) land the same. We'll need TPU silicon again in Part 4, where chip-specific behavior actually matters.

The seven ideas:

1. Arrays are immutable (no in-place ops).
2. jit — trace once, compile once, reuse while shapes stay the same.
3. grad — autodiff on a pure function.
4. vmap — auto-batching for free.
5. PRNG keys — randomness is explicit and stateless.
6. Pytrees — parameters are nested dicts, not objects.
7. Functional training step — no nn.Module, no optimizer.step().

Let’s walk through each with runnable code.

1. Immutability

PyTorch happily lets you do this:

x[0] = 99.0   # fine, modifies x in place

JAX doesn’t:

import jax.numpy as jnp
x = jnp.array([1.0, 2.0, 3.0])
x[0] = 99.0                    # TypeError!
x2 = x.at[0].set(99.0)         # returns a NEW array; x is unchanged

That x.at[0].set(99.0) idiom is how you express "I want a modified copy." There's also .at[idx].add(...), .at[idx].mul(...), and so on.

Why the restriction? Because JAX’s power features — jit, grad, vmap — require functions to be pure (same inputs always produce same outputs, no side effects). An in-place mutation is a side effect on the array that was passed in. That would break everything.

It feels awkward at first and then you stop noticing it.

2. jit — the headline speedup

This is the feature that makes JAX on TPU actually fast. Here’s the setup:

import time
import jax
from jax import jit
import jax.numpy as jnp
def slow_fn(x):
    # A silly sequence of ops, good enough to see the effect.
    for _ in range(100):
        x = jnp.sin(x) + jnp.cos(x)
    return x
x = jnp.ones((1000, 1000))
# Eager (un-jit'd) version.
t0 = time.time()
slow_fn(x).block_until_ready()
print(f"eager  : {(time.time() - t0) * 1000:7.2f} ms")
# Jit'd version - first call includes compile time.
fast_fn = jit(slow_fn)
t0 = time.time()
fast_fn(x).block_until_ready()
print(f"jit #1 : {(time.time() - t0) * 1000:7.2f} ms")
# Second call: no compile, just run.
t0 = time.time()
fast_fn(x).block_until_ready()
print(f"jit #2 : {(time.time() - t0) * 1000:7.2f} ms")

A couple of lines deserve explanation:

• fast_fn = jit(slow_fn) — this doesn't compile anything yet. It wraps the function. Actual tracing and compilation happens on the first call.
• .block_until_ready() — same as Part 2. JAX dispatch is async, so without this our timings would be meaningless (we'd be timing the dispatch, not the computation).

Running this on a v6e-1, I got:

eager   :  518.52 ms
jit #1  :  153.36 ms   (includes compile)
jit #2  :    0.21 ms   (pure run)

Read those numbers carefully. The eager version takes 518 ms. The first jit call takes 153 ms — that already includes XLA compiling the whole sequence, and it’s still faster than eager because XLA is aggressively fusing the ops. The second jit call takes 0.21 ms. That’s a ~2500× speedup over eager after the one-time compile cost.

This is why, on TPU, basically everything interesting lives inside a @jit.

Gotcha to remember: JIT caches compiled programs based on the function, input shapes, dtypes, and any arguments marked static. If you call the same jit’d function with a different array shape or dtype, it recompiles. Ordinary array values usually do not cause recompilation by themselves. Python scalars are fine when they behave like normal numeric inputs, but they become compile-time/static values if you use them in Python control flow, array shape construction, or mark them with static_argnames. If a value truly has to control a Python if, range, or array shape, mark it static — but remember that each different static value gets its own compilation.

3. grad — autodiff on a function

PyTorch autodiff is centered on tensors: you set requires_grad=True, do a forward pass, call .backward(), and read .grad off the tensor.

JAX autodiff is centered on functions. You take a function, and you get a new function that computes its gradient.

import jax
import jax.numpy as jnp
from jax import grad
def f(x):
    return jnp.sin(x) ** 2
df = grad(f)        # df is a new function: the derivative of f
print(f(1.0))       # f(1.0)  = sin(1.0)^2 ≈ 0.708
print(df(1.0))      # f'(1.0) = 2*sin(1.0)*cos(1.0) ≈ 0.909

grad(f) returns a new function that computes df/dx. That's the whole API. A few things worth knowing:

• grad(grad(f)) — gives you the second derivative, for free.
• grad(f, argnums=1) — differentiate with respect to the second argument.
• jax.value_and_grad(f) — returns a function that gives you both f(x) and grad(f)(x) in a single forward/backward pass. Use this in training loops — it saves a forward pass.

There’s no requires_grad, no .backward(), no .grad attribute. Just functions mapping to functions.

4. vmap — auto-batching for free

Write a function for one example, get a batched version for free:

import jax.numpy as jnp
from jax import vmap
def euclid(a, b):
    # Distance between two points.
    return jnp.sqrt(jnp.sum((a - b) ** 2))
# Single-pair example.
print(euclid(jnp.array([0.0, 0.0]), jnp.array([3.0, 4.0])))  # 5.0
# Now batch it: pa and pb are (N, 2) - N points each.
pa = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]])
pb = jnp.array([[3.0, 4.0], [1.0, 2.0], [2.0, 3.0]])
distances = vmap(euclid)(pa, pb)
print(distances)   # [5. 1. 1.]

vmap(euclid) returns a new function that, when given arrays with a leading batch dimension, runs euclid on each slice along that dimension and stacks the results. You didn't write any batching loops.

Why does this matter?

• In PyTorch, you often write the batched version by hand: add a leading dim, broadcast carefully, and hope the shapes line up. That’s error-prone.
• With vmap, you write the per-example logic, which is usually simpler, and get batching for free.

And vmap composes with jit and grad:

# Per-example gradients: params are shared, x and y are batched.
per_example_grads = jit(vmap(grad(loss_one), in_axes=(None, 0, 0)))

The in_axes=(None, 0, 0) means: don't batch params, but do batch the inputs x and labels y.

Get Roya’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

PyTorch has torch.func.vmap these days, and it works well. In JAX, vmap has been there since day one, and it composes freely with every other JAX transformation (jit, grad, pmap) — so it feels more central to the library than it does in PyTorch.

5. PRNG keys — explicit randomness

Here’s a subtle but important one.

PyTorch keeps RNG state in a global variable. When you call torch.randn(3), you're implicitly consuming and advancing that state. That's convenient for imperative code, but it breaks under jit: a jit'd function is supposed to be pure (same inputs → same outputs), and a function that invisibly reads and writes a global RNG isn't pure.

JAX handles this by making RNG state an explicit value you pass around:

from jax import random
key = random.key(42)          # modern JAX key API
# random.PRNGKey(42) also works in older code
k1, k2 = random.split(key)    # derive two new keys from it
r1 = random.normal(k1, (3,))
r2 = random.normal(k2, (3,))
print(r1)   # [ 0.076 -0.486  1.290]
print(r2)   # [ 0.606  0.799 -0.909]

Notes on these lines:

• random.key(42) — makes a "key" from a seed. Think of it as a handle to a specific point in the RNG sequence.
• random.split(key) — the critical one. You can reuse a key syntactically — JAX won't stop you — but you should treat it as a bug: same key + same random function = same random values. split derives new independent keys from an existing one.
• random.normal(k1, (3,)) — draw samples. The key is the first argument, which feels weird at first.

The mental rule: every time you need randomness, split first, then use. If you forget to split and reuse key in two places, you'll get identical "random" numbers in both places and wonder why your training is broken.

6. Pytrees — parameters are nested containers

PyTorch organizes model parameters in nn.Module objects. You write a class, define self.weight = nn.Parameter(...), and call self.parameters() to iterate over them.

JAX has none of that. A “model” is just a nested dict (or list, or dataclass) of arrays. Collectively, any nested container of arrays is called a pytree.

import jax.numpy as jnp
params = {
    "w1": jnp.ones((4, 3)),
    "b1": jnp.zeros(3),
    "w2": jnp.ones((3, 1)),
    "b2": jnp.zeros(1),
}

That’s it. That’s a model’s parameters. To operate on every array in a pytree at once:

from jax import tree_util
# Map a function over every leaf array in the pytree.
shapes = tree_util.tree_map(lambda a: a.shape, params)
print(shapes)
# {'b1': (3,), 'b2': (1,), 'w1': (4, 3), 'w2': (3, 1)}
total = sum(a.size for a in tree_util.tree_leaves(params))
print("total params:", total)   # 19

• tree_util.tree_map(f, tree) — applies f to every leaf array in the pytree and returns a new pytree with the same structure.
• tree_util.tree_leaves(tree) — flattens the pytree into a list of its leaf arrays.

This looks minimal and you might think “surely real models need more.” Real projects often layer libraries like Flax, Equinox, or Haiku on top — but underneath, the parameters still behave like pytrees. The next section shows a full training step using nothing but a pytree of params.

7. A functional training step

All six ideas above compose into the canonical JAX training loop. This assumes we’re using the params pytree from the previous section.

import jax.numpy as jnp
from jax import jit, grad, tree_util
def mlp(params, x):
    # Forward pass: a 2-layer MLP with tanh activation.
    h = jnp.tanh(x @ params["w1"] + params["b1"])
    return h @ params["w2"] + params["b2"]
def loss(params, x, y):
    return jnp.mean((mlp(params, x) - y) ** 2)
@jit
def step(params, x, y, lr=0.1):
    g = grad(loss)(params, x, y)
    # Return NEW params - no in-place update.
    return tree_util.tree_map(lambda p, dp: p - lr * dp, params, g)

Look at step carefully:

• It’s a pure function: takes params in, returns new params out.
• No self. No class. No nn.Module.
• No optimizer.zero_grad() (there's no state to zero).
• No optimizer.step() (we're literally just computing p - lr * dp for every leaf).
• The entire thing is @jit'd: XLA sees the whole training step and compiles it once.

A training loop using it:

# Fake data.
x_data = jnp.ones((16, 4))
y_data = jnp.ones((16, 1))
for i in range(5):
    params = step(params, x_data, y_data)
    if i % 4 == 0:
        print(f"step {i}: loss = {loss(params, x_data, y_data):.4f}")

The key mental flip from PyTorch: state flows through your functions rather than living on objects. Params go in, new params come out, you reassign.

For real training you’d replace the hand-rolled SGD with Optax (JAX’s equivalent of torch.optim, with Adam, LR schedules, grad clipping, etc.). We'll use Optax in Part 5.

The complete PyTorch ↔ JAX cheat sheet

Things that go wrong

ConcretizationTypeError: Abstract tracer value... — you used a traced array in a Python if, while, or len(). Inside a jit'd function, shapes must be known statically; tensor values aren't available. Use jnp.where for conditionals, jax.lax.cond for branching, or mark the argument as static with static_argnames.

“My model trains but produces the same output every time.” — you reused a PRNG key instead of splitting it. Search your code for every call to a random.* function and make sure each one gets a fresh split.

Everything is suspiciously slow. — you forgot @jit. JAX eager mode is useful for debugging, but many small eager operations pay Python dispatch overhead. For training, put the whole step under @jit.

“My function recompiles on every call.” — most likely your input shapes or dtypes are changing between calls, or you’re passing static arguments whose values keep changing. Check whether a Python value is being used in shape construction, Python control flow, or static_argnames.

What you actually need to remember

If you remember only four things from this part:

1. JAX arrays don’t mutate. Use .at[...].set(...) for updates.
2. @jit the training step. Eat the compile cost once; reap the speedup forever.
3. Randomness takes a key. Always split before use.
4. Your “model” is a pytree. State flows through functions, not through objects.

Everything else is variation on those four.

In Part 4, we’ll stop talking about JAX abstractly and crack open what @jit actually produces: the compiled program that XLA ships to the TPU. We'll also measure real bf16 vs fp32 throughput on a v6e and compare against a Colab T4 GPU — so you see the performance story concretely.