the museum
xla wing

serves the xla path · three reproduced on 8 host-platform devices (jax 0.4.38, CPU), one captured on a Colab TPU v6e and labelled where it sits · fixes run and proven · all wings

The partitioner says no, and it means something.

Three exhibits come from the SPMD seam, where your sharding meets the compiler's rules: a spec the shape cannot satisfy, an axis the mesh does not have, and a collective called with no mesh in sight. All three reproduce on a laptop with eight fake host devices. The fourth could only come from real hardware, and it did: LAB·X2 asks a Colab TPU for more HBM than it has, and the refusal arrives with both numbers attached. More of those land here as the lab gets run.

exhibit
01/04

A sharding the array's shape cannot satisfy

A PartitionSpec is a promise that each named mesh axis divides its array dimension. Ten rows over four data shards leaves a remainder, and the runtime refuses rather than silently pad. Pad or reshape to a multiple, or shard a different dimension.

the failing code

import jax
import jax.numpy as jnp
import numpy as np
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

mesh = Mesh(np.array(jax.devices()[:8]).reshape(4, 2), ("data", "model"))

x = jax.device_put(jnp.ones((10, 16)), NamedSharding(mesh, P("data", "model")))

the error, verbatimValueError: One of device_put args was given the sharding of NamedSharding(mesh=Mesh('data': 4, 'model': 2), spec=PartitionSpec('data', 'model'), memory_kind=unpinned_host), which implies that the global size of its dimension 0 should be divisible by 4, but it is equal to 10 (full shape: (10, 16))

the fix

import jax
import jax.numpy as jnp
import numpy as np
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

mesh = Mesh(np.array(jax.devices()[:8]).reshape(4, 2), ("data", "model"))

x = jax.device_put(jnp.ones((12, 16)), NamedSharding(mesh, P("data", "model")))

exhibit
02/04

A constraint naming an axis the mesh does not have

Sharding constraints are resolved against the mesh's axis names, and this mesh knows data and model, not batch. The partitioner cannot invent an axis; the fix is spelling the mesh's own vocabulary.

the failing code

import jax
import jax.numpy as jnp
import numpy as np
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

mesh = Mesh(np.array(jax.devices()[:8]).reshape(4, 2), ("data", "model"))

@jax.jit
def f(x):
    return jax.lax.with_sharding_constraint(x, NamedSharding(mesh, P("batch")))

f(jnp.ones((8, 16)))

the error, verbatimValueError: Resource axis: batch of PartitionSpec('batch',) is not found in mesh: ('data', 'model').

the fix

import jax
import jax.numpy as jnp
import numpy as np
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

mesh = Mesh(np.array(jax.devices()[:8]).reshape(4, 2), ("data", "model"))

@jax.jit
def f(x):
    return jax.lax.with_sharding_constraint(x, NamedSharding(mesh, P("data")))

f(jnp.ones((8, 16)))

exhibit
03/04

A collective called outside any mesh context

psum sums across a named mesh axis, and axis names only exist inside a mapped context: shard_map, pmap, or a mesh-aware jit. On its own the name is unbound. Collectives are kernels over the mesh, not array ops; give them the mesh.

the failing code

import jax
import jax.numpy as jnp

jax.lax.psum(jnp.ones(4), axis_name="data")

the error, verbatimNameError: unbound axis name: data

the fix

import jax
import jax.numpy as jnp
import numpy as np
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

mesh = Mesh(np.array(jax.devices()[:8]).reshape(4, 2), ("data", "model"))

from functools import partial
from jax.experimental.shard_map import shard_map

@partial(shard_map, mesh=mesh, in_specs=P("data", None), out_specs=P())
def total(block):
    return jax.lax.psum(block.sum(), axis_name="data")

out = total(jnp.ones((8, 16)))

exhibit
04/04

Asking for more HBM than the chip has

captured on a Colab TPU v6e (reported as TPU v6 lite) during LAB·X2, 2026-07-27; the fix is the same program sized against the budget the error itself reports

Buffer assignment places every temporary before a single instruction runs, so a program that cannot fit is refused at compile time rather than discovered halfway through. Read the refusal as a measurement: it names what the program asked for and what the chip has, and the ratio between those two numbers is your first estimate of how much smaller the shapes need to be.

the failing code

BIG = 300000
huge = jnp.ones((BIG, BIG), dtype=jnp.float32)
out = jax.jit(lambda x: x + 1.0)(huge)

the error, verbatimRESOURCE_EXHAUSTED: Ran out of memory on HBM, the total memory required for HLO temporaries (335.31G) exceeds available HBM (31.24G).

the fix

# size the program to the budget the chip just told you about:
# 31.24G of HBM, float32 at 4 bytes, and room for more than one temporary
BIG = 40_000                       # 40000 * 40000 * 4 bytes = 6.4 GB
huge = jnp.ones((BIG, BIG), dtype=jnp.float32)
out = jax.jit(lambda x: x + 1.0)(huge)