The carve, stated exactly
The layer guide at /l/pallas already makes the central correction: an index mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → returns block coordinates, not element offsets, and Pallas multiplies by the block shape to get the address. The reference states the full rule as executable code, and two details in it are easy to miss.
First, the multiplication is unconditional and per axis, so every axis of the returned tuple is scaled by that axis's block size. Second, there is an assertion: at least one element of the block must be within bounds. A start index past the end of the array is an error, but a block that starts inside and runs off the end is legal and common.
block_indices = x_spec.index_map(*invocation_indices)
elem_indices = []
for x_size, block_size, block_idx in zip(x_shape, x_spec.block_shape, block_indices):
start_idx = block_idx * block_size
# At least one element of the block must be within bounds
assert start_idx < x_size
elem_indices.append(slice(start_idx, start_idx + block_size))The last block is a promise about garbage
When the block shape does not divide the array evenly, the final iteration on that axis still receives a full-shape block. Out-of-bounds elements are padded on input and discarded on output, and the reference is explicit that the padding values are unspecified and you should assume they are garbage.
Interpret mode pads with NaN for floating-point values so you have a chance to spot the access, and the same paragraph says this behavior should not be depended upon. Both halves matter. The NaN is a debugging affordance, so treat a NaN in interpret mode as a real finding; and never write a kernel whose correctness needs the padding to be any particular value.
That is why a mask built from an iota against the array's true bound is not defensive style. It is the only thing standing between a reduction and unspecified bytes. The museum keeps this family of rank mistakes, an index mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → handing back one coordinate for a two-axis block among them, at /mistakes/kernels.
None axes, whole-array specs, and the default map
A None entry in block_shape behaves as the value 1, except that the axis is squeezed out of the ref the kernel body sees. You can write pl.Squeezed() for the same thing. A three-axis array carved with (None, 2) on its last two axes hands the body a rank-1 ref, which is usually what you wanted when you wrote a leading batch axis into the grid.
Two defaults save typing and are worth knowing by name. block_shape=None means the whole array shape. index_map=None means the map that returns zeros for every axis regardless of grid position, which is the "stage it once, everybody reads the same copy" spec written the short way.
There is also a second indexing mode most kernels never touch. With pl.Element(block_size), the values your index mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → returns are used directly as array indices with no scaling, and you may declare virtual low-high padding per dimension as though the array were padded on input. Element mode is TPU only. Reach for it when the natural expression of your access is an element offset and the block arithmetic is fighting you.
def kernel(o_ref):
assert o_ref.shape == (2,) # (None, 2) arrived as rank 1
o_ref[...] = jnp.full((2,), 10 * pl.program_id(1) + pl.program_id(0))
pl.pallas_call(kernel,
jax.ShapeDtypeStruct((3, 4), dtype=np.int32),
out_specs=pl.BlockSpec((None, 2), lambda i, j: (i, j)),
grid=(3, 2), interpret=True)()What makes a carve good
Three constraints decide a block shape, and they pull in different directions. The lattice is the hard one: the last two dimensions must be divisible by 8 and 128, or equal the array's dimensions exactly. Rank-1 blocks have their own rule that the path never states, namely that the block dimension must equal the array dimension, or be a multiple of 1024, or be a power of two and at least 128 * (32 / bitwidth(dtype)).
The budget is the soft one. Larger windows generally give better hardware utilization, so the pull is upward, and the ceiling is that a window plus the space for spilled vector registers can exceed VMEM. What you get then is a low-level compiler error about memory, which is the museum's VMEMThe TPU’s software-managed vector scratchpad, about 128 MiB. Blocks must be staged here before compute touches them; what is resident is what your schedule staged.taught in /l/tpu → overflow exhibit and the reason a budget line belongs in your notes before you compile.
Reuse is the one people forget, because it is a property of the grid order rather than the block shape. When two lexicographically consecutive grid indices map to the same slice of an input, the HBMThe chip’s main memory: large, far, and the resource memory-bound ops spend. 8.2e11 bytes per second on v5e, 1.6e12 on v6e.taught in /l/tpu → transfer for the second one is skipped, since the data is already there. Your grid order therefore sets how many transfers your carve costs, and reordering axes can change traffic without changing a single block dimension.
| constraint | rule | how it fails |
|---|---|---|
| lattice | last two dims divisible by 8 and 128, or equal the array dims | compile error naming both numbers |
| rank-1 blocks | equal the array dim, or a multiple of 1024, or a power of two at least `128 * (32 / bitwidth)` | compile error at lowering |
| VMEM budget | window plus spilled registers must fit | low-level out-of-memory from the backend |
| reuse | consecutive steps on the same slice skip the transfer | no error at all, just traffic you paid for |
Writes have an ordering rule that reads do not
Because the TPU grid runs sequentially, several invocations may write the same slice of the output with no risk of a race. The reference attaches one condition to that permission: all invocations that write a particular slice must be consecutive. Break the run and you are outside what the backend guarantees.
This is where the matmul convention comes from, and it is a correctness argument rather than a performance one. Some prefix of the grid axes varies the output slice; the remaining suffix leaves the output window fixed. A reduction axis leaves the output fixed, so it has to be last, and the output ref then works as an accumulator across it. Put K first and the writes to a given output block are no longer consecutive.
The general statement is stricter still: when multiple invocations write to the same elements of the output, the result is platform dependent. Nothing about your shapes or dtypes is invalid, so nothing warns you.
BlockSpecs, precisely
You already know that a BlockSpecHow one array is carved for the grid: a block shape plus an index map saying which block each grid step sees.taught in /l/pallas → pairs a block_shape with an index_mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → that changes as the grid advances. What the mental model skips is exactly how those two combine into an address. The index map does not return an offset in elements. It returns a block coordinate: which numbered block along that axis, not which row or column. Pallas then multiplies each returned coordinate by the matching entry in block_shape to get the element offset into the ref. Grid step (2, 0) paired with block shape (128, 512) reads starting at element (256, 0), because 2 times 128 is 256.
This is the detail that trips people up the first time they write an index mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → by hand. It feels natural to write the arithmetic yourself, to return (256, 0) directly, since that is the offset you actually want. Do that and every block after the first one lands in the wrong place, because Pallas multiplies again on top of an offset you already computed. The index map's job is to count blocks, not elements. Let the framework do the multiplication it is already going to do.
Sometimes you want an entire array staged in VMEMThe TPU’s software-managed vector scratchpad, about 128 MiB. Blocks must be staged here before compute touches them; what is resident is what your schedule staged.taught in /l/tpu → rather than a sliced piece across the grid, a weight matrix every step needs in full, for instance. Write an index mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → that returns zeros regardless of the grid coordinate, a constant (0, 0) on every step. Set the block shape to match the full array shape on that axis. The compiler then stops walking that axis in blocks: it stages the array once and every grid step reads the same VMEM copy. It is the same offset mechanism as a blocked spec, just with a constant index map instead of one that tracks the grid coordinate.
Nothing stops you from mixing the two within a single BlockSpec. A weight matrix might be blocked along its output-feature axis, where each grid step really does need a different slice, while its input-feature axis stays whole, because every step needs the full contraction dimension to do a matmul at all. The index mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → returns a real coordinate on one axis and a constant zero on the other, in the same tuple.
The block shape itself is not free-form. TPU vector registers store data in a tile fixed by dtype: (8, 128) for f32, (16, 128) for bf16, (32, 128) for int8. The last two dimensions of any block shape must be a whole multiple of that tile, or must equal the full array dimension on that axis exactly when you are not blocking it further. There is no third option. A block that sits between those two cases, smaller than the full axis but not a multiple of the tile, does not compile.
The rule exists because the compiler lays vector registers out in that exact tile shape underneath every block you ask for. A block that does not respect the tile has no consistent way to map onto real registers, and there is no fallback layout for the compiler to fall back to instead: it rejects the shape outright, at compile time, rather than guess at a layout that might be wrong.
Three failures in this area show up often enough that the museum tracks each by name: a mismatch between the rank the index mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → returns and the rank of the block shape, a mismatch between the number of arguments the index map accepts and the number of grid axes, and a block shape that breaks the divisibility rule outright. All three are caught at compile time, not at run time, which is the one piece of good news: none of them corrupts data silently, they just stop the kernel from lowering at all.
Check yourself
01 What does an index map return, and who does the multiplication?
Block coordinates per axis; Pallas multiplies by the block shape unconditionally to get element offsets.
02 What may the padding past the array's edge contain?
Garbage, by contract: unspecified on input, discarded on output. Interpret mode's NaN is a debugging aid you may not rely on.
Readings
- Grids and BlockSpecs ↗ the exact slice semantics, None axes, and the Element indexing mode
- TPU kernels: noteworthy properties and restrictions ↗ the consecutive-write rule and the revisit skip, in the backend's own words