Programming TPUs in JAX

Chapter 1: Auto Sharding

jax.jit plays two roles: it compiles Python to fast XLA code, and — when inputs are sharded — it distributes computation across devices. The XLA compiler's Shardy partitioner automatically decides how to split work and where to insert communication.

Here is a complete sharded matmul in auto mode:

import jax
import jax.numpy as jnp

# Create a 4x2 mesh (e.g., TPU v5e 4x2)
mesh = jax.make_mesh(axis_shapes=(4, 2), axis_names=('X', 'Y'))
jax.set_mesh(mesh)

# Create sharded arrays
In = jnp.zeros((8, 2048), dtype=jnp.bfloat16,
               device=jax.NamedSharding(mesh, jax.P('X', 'Y')))
W = jnp.zeros((2048, 8192), dtype=jnp.bfloat16,
              device=jax.NamedSharding(mesh, jax.P('Y', None)))

def matmul_square(In, W):
  return jnp.einsum('bd,df->bf', jnp.square(In), W)

# Compile with output sharding constraint
jit_matmul = jax.jit(matmul_square,
  out_shardings=jax.P('X', None)).lower(In, W).compile()

out = jit_matmul(In, W)

What is actually happening at the hardware level?

• W is sharded 2-way along the contracting dimension: W[DY, F]
• In is sharded 8-way: 4-way batch, 2-way contracting: In[BX, DY]
• The output is P('X', None) — sharded along batch, replicated across model
• This requires an AllReduce to sum partial products across the Y axis

Using our sharding notation, XLA will emit:

We can verify this by examining the HLO with jit_matmul.as_text():

# The matmul (local on each device)
%fusion = bf16[2,8192]{1,0} fusion(
  bf16[2,1024] %param,
  bf16[8192,1024] %copy-done)

# The AllReduce across devices
ROOT %AllReduce = bf16[2,8192]{1,0} AllReduce(
  bf16[2,8192] %fusion)

Note the per-device shapes: bf16[2, 1024] for activations (batch=8 split 4 ways, d_model=2048 split 2 ways).

Chapter 2: Explicit Sharding

Explicit sharding (or "sharding in types") looks a lot like auto sharding, but with one crucial difference: sharding propagation happens at the JAX level, not the XLA level. This means JAX can catch ambiguous decisions and ask you for clarification instead of silently guessing.

The key concept: each JAX operation has a sharding rule that takes input shardings and produces an output sharding. For most ops, there is only one reasonable choice. For some (like einsum), it is ambiguous.

import jax
import jax.sharding as shd

# Create mesh with Explicit axis types
mesh = jax.make_mesh(
  axis_shapes=(2, 2), axis_names=('X', 'Y'),
  axis_types=(shd.AxisType.Explicit, shd.AxisType.Explicit))
jax.set_mesh(mesh)

x = jax.device_put(np.arange(16).reshape(8, 2), jax.P('X', 'Y'))

@jax.jit
def f(x):
  print(jax.typeof(x))   # float32[8@X, 2@Y]
  out = x * 2
  print(jax.typeof(out))  # float32[8@X, 2@Y]
  return out

Notice the @X and @Y annotations in the type — the sharding is now part of the type system. For elementwise ops like x * 2, the output sharding is obvious: same as the input.

But for a matmul, things get ambiguous:

In = jnp.zeros((8, 2048), out_sharding=jax.P('X', 'Y'))
W = jnp.zeros((2048, 8192), out_sharding=jax.P('Y', None))

@jax.jit
def matmul_square(In, W):
  return jnp.einsum('bd,df->bf', jnp.square(In), W)

matmul_square(In, W)  # ERROR!

This produces a clear error:

Contracting dimensions are sharded and it is ambiguous
how the output should be sharded.
Please specify the output sharding via the `out_sharding` parameter.
Got lhs_contracting_spec=('Y',) and rhs_contracting_spec=('Y',)

The fix is to specify the output sharding:

@jax.jit
def matmul_square(In, W):
  return jnp.einsum('bd,df->bf', jnp.square(In), W,
                     out_sharding=jax.P('X', 'Y'))
# This induces a ReduceScatter (output sharded along both axes)

Auto and Explicit modes can be composed using jax.sharding.auto_axes and jax.sharding.explicit_axes APIs, so you can use explicit control on some mesh axes and let the compiler handle others.

Chapter 3: Manual Mode via shard_map

jax.shard_map is the "just let me write what I mean" mode. You get a per-device local view of the program and write all communication explicitly. No compiler surprises.

Here is a simple example — slicing and averaging across devices:

import jax
import jax.numpy as jnp
import jax.sharding as shd

mesh = jax.make_mesh((2, 4), ('x', 'y'),
  (shd.AxisType.Explicit, shd.AxisType.Explicit))
jax.set_mesh(mesh)

x = jnp.arange(0, 512, dtype=jnp.int32,
               out_sharding=jax.P(('x', 'y')))

# This function operates on 1/8th of the array
@jax.shard_map(in_specs=jax.P(('x', 'y')),
               out_specs=jax.P())
def slice_and_average(x):
  assert x.shape == (512 // 8,)  # Local view!
  return jax.lax.pmean(x[:4], axis_name=('x', 'y'))

out = slice_and_average(x)
assert out.shape == (4,)

What is happening? Each device sees only 64 elements (512 / 8). We slice the first 4 elements from each shard, then average them across all 8 devices using pmean (an AllReduce mean). The result is mean(x[:4], x[64:68], x[128:132], ...).

When to Use shard_map

Use shard_map when you need any of the following:

• Custom communication patterns: Circular shifts, butterfly patterns, or any communication that is not a standard AllGather/AllReduce
• Guaranteed no surprise communication: In auto/explicit mode, the compiler might insert unexpected collectives. With shard_map, what you write is what runs.
• Communication-computation overlap: The collective matmul pattern requires explicit control over when data is sent and when matmuls happen
• Debugging sharding issues: When auto mode produces a confusing profile, rewriting the critical section in shard_map gives you certainty

The key primitives available inside shard_map:

Chapter 4: The Collective Matmul

This is where shard_map really shines. Consider model parallelism where activations are sharded along the model dimension: A[BX, DY] * W[D, FY] → Out[BX, FY].

The naive approach first gathers all of A, then does a local matmul:

The problem: communication and computation happen sequentially. The TPU sits idle during the AllGather.

The collective matmul (Wang et al., 2023) overlaps them. The algorithm:

1. For each Y shard: perform a matmul of the local chunk of A with the corresponding chunk of W, producing a partial result
2. Simultaneously, use ppermute to circularly shift A to get the next chunk
3. Accumulate partial results

def collective_matmul_allgather_lhs(lhs, rhs):
  axis_size = jax.lax.axis_size('Y')
  idx = jax.lax.axis_index('Y')
  chunk_size = lhs.shape[1]

  def f(i, carrys):
    accum, lhs = carrys
    # Matmul with the right chunk of rhs
    rhs_chunk = jax.lax.dynamic_slice_in_dim(
      rhs, (idx + i) % axis_size * chunk_size, chunk_size)
    update = lhs @ rhs_chunk
    # Circular shift lhs to get next chunk
    lhs = jax.lax.ppermute(lhs, axis_name='Y',
      perm=[(j, (j - 1) % axis_size)
            for j in range(axis_size)])
    return accum + update, lhs

  accum = jnp.zeros((lhs.shape[0], rhs.shape[1]))
  accum = jax.lax.pvary(accum, ('X', 'Y'))
  accum, lhs = jax.lax.fori_loop(
    0, axis_size - 1, f, (accum, lhs), unroll=True)

  # Last chunk
  i = axis_size - 1
  rhs_chunk = jax.lax.dynamic_slice_in_dim(
    rhs, (idx + i) % axis_size * chunk_size, chunk_size)
  return accum + lhs @ rhs_chunk

The profile for the collective matmul shows no AllGather op — it is all useful compute with communication overlapped underneath. FLOPs utilization is much higher.

Chapter 5: Meshes and PartitionSpecs

Before we go further, let us solidify the two core abstractions that all three modes share: meshes and PartitionSpecs.

Meshes

A mesh is a named, shaped arrangement of your TPU devices. It assigns human-readable names to the physical axes of your hardware:

# A 4x2 mesh: 4 devices along 'X', 2 along 'Y'
mesh = jax.make_mesh(
  axis_shapes=(4, 2),
  axis_names=('X', 'Y')
)

Common mesh configurations and their use:

PartitionSpecs

A PartitionSpec (jax.P) maps array dimensions to mesh axes. Each position corresponds to an array dimension:

# For a 2D array [batch, features]:
jax.P('X', 'Y')     # Shard dim 0 along X, dim 1 along Y
jax.P('X', None)   # Shard dim 0 along X, replicate dim 1
jax.P(None, 'Y')   # Replicate dim 0, shard dim 1 along Y
jax.P()            # Fully replicated

Ways to Create Sharded Arrays

1. At creation time: jnp.zeros(..., device=jax.NamedSharding(mesh, P(...)))
2. Via device_put: jax.device_put(array, P(...))
3. Via out_sharding: jnp.zeros(..., out_sharding=P(...))
4. Via jit output: jax.jit(f, out_shardings=P(...))

Chapter 6: MoE with shard_map

Mixture of Experts (MoE) models are one of the most compelling use cases for shard_map. The routing is inherently irregular — each token goes to a different expert — which makes it hard for automatic partitioners to do the right thing.

The Problem Setup

We have:

• W: float32[EX, D, F] — E expert weight matrices, sharded across X
• A: float32[SX, D] — activations, sharded across X
• B: int32[SX] — routing assignments (which expert for each token)

Goal: compute Out[i] = A[i] @ W[B[i]] — each token is processed by its assigned expert.

Local Implementation (No Sharding)

def moe_local(W, A, B):
  S, _ = A.shape
  E, _, F = W.shape

  def expert_forward(carry, e):
    output = carry
    mask = (B == e)[:, None]     # [S, 1]
    expert_result = A @ W[e]    # [S, F]
    output = output + expert_result * mask
    return output, None

  output = jnp.zeros((S, F))
  output, _ = jax.lax.scan(expert_forward, output, jnp.arange(E))
  return output

This iterates over experts with masking. Each expert computes its transform of all tokens, then masks out the irrelevant ones. Simple but wasteful — we compute E times more FLOPs than needed.

The Sharded Challenge

If you just jax.jit the local implementation with sharded inputs, XLA will likely AllGather the full activations A locally. This is expensive in both communication and memory.

The better approach uses shard_map with an AllToAll to route tokens to the correct devices:

You can also use jax.lax.ragged_dot for a more efficient implementation that avoids the masking overhead entirely.

Chapter 7: End-to-End Transformer

Let us put everything together and build a complete Transformer MLP block with overlapped communication — the kind of code that runs in production systems.

The Two Matmuls

A Transformer FFN with Megatron-style model parallelism has two matmuls with different communication patterns:

AllGather Collective Matmul (Up-projection)

This is what we implemented in Chapter 4: circular shift + chunk matmul to overlap the AllGather with computation.

ReduceScatter Collective Matmul (Down-projection)

The complement: instead of gathering inputs, we scatter results. The idea:

1. Tile over the input (F) dimension
2. Compute partial products for each tile
3. Permute the accumulator as you go, so each device accumulates only its output shard

Combined Pipeline

Putting both together for a full FFN block:

# Full overlapped FFN: In[BX, DY] -> Out[BX, DY]
@jax.shard_map(
  in_specs=(jax.P('X', 'Y'), jax.P(None, 'Y'), jax.P('Y', None)),
  out_specs=jax.P('X', 'Y'))
def ffn_collective(x, w_up, w_down):
  # Up-projection: AG-matmul overlaps AllGather with compute
  h = collective_matmul_allgather(x, w_up)
  h = jax.nn.gelu(h)
  # Down-projection: RS-matmul overlaps compute with ReduceScatter
  out = collective_matmul_reducescatter(h, w_down)
  return out

Bidirectional Communication

The collective matmuls above use unidirectional circular shifts. An extension uses bidirectional communication — sending data both left and right simultaneously. This can further reduce latency by overlapping both directions of the ring, though the implementation is more complex.

Choosing the Right Approach

Chapter 8: Worked Problems

Here are several problems to test your understanding. These are meant to be implemented on a real TPU (free TPU v2-8 available on Google Colab).

Problem 1: Shard-local Operations

Let A: float32[SX, DY] with X * Y = N.

Part a: Write a function that computes the average within each (X, Y) shard, returning an array of shape [X, Y]. Implement in both jax.jit and shard_map. Profile each. Was communication added?

Part b: Write a function that returns roll(x, shift, axis=0) - x for some shift, applied within each shard along X. Use shard_map.

Problem 2: Mixture of Experts

Build a complete MoE forward pass:

1. Start with a local (single-device) implementation using scan + masking
2. Profile the jax.jit version with sharded inputs — what communication does it add?
3. Implement with shard_map and explicit AllToAll to avoid materializing a full [E, S, D] buffer
4. Extend to top-k routing where each token goes to k experts

Problem 3: Collective Matmul Variations

1. AllReduce collective matmul: Implement A[BX, DY] * W[DY, F] → Out[BX, F] with overlapped communication. Tile over the output dimension and use jax.lax.psum.
2. ReduceScatter collective matmul: Implement Tmp[BX, FY] * W2[FY, D] → Out[BX, DY]. Be careful about passing only the minimal data needed. Hint: permute the result as you accumulate.
3. End-to-end: Combine both into a full Transformer block: In[BX, DY] * Win * Wout → Out[BX, DY]

Problem 4: Bidirectional Collective Matmul

Rewrite the collective AllReduce and ReduceScatter matmuls to use bidirectional communication (sending data both left and right simultaneously). How much faster are these?