the museum

assigned reading in chapter 11 · every error below was reproduced for real; the text is what the interpreter said

Meet every failure in a toy first.

Each exhibit ran, failed, and was captured: the snippet, the verbatim error (paths shortened), the fix, and the rule that prevents it. Read them before your first kernel and again after; the error messages start making sense the second time.

exhibit
01/06

BlockSpec index_map returns the wrong rank

index_map runs once per grid step and it has to hand back one coordinate for every axis of the block shape you declared. A (2, 128) block is two-dimensional, so index_map needs to return a two-element tuple, but this one returns just `(i,)`. The error spells out the mismatch directly: it wants 2 values and gets 1. Whenever you see that pattern, count the axes in your block_shape and check the length of the tuple your index_map returns against it, they need to match even when one axis never moves.

the failing code

def broken_index_map_rank():
    def kernel(x_ref, o_ref):
        o_ref[...] = x_ref[...] * 2.0

    def index_map(i):
        return (i,)  # bug: a (2, 128) block over a 2D array needs two entries

    x = jnp.ones((8, 128), dtype=jnp.float32)
    return pl.pallas_call(
        kernel,
        grid=(4,),
        in_specs=[pl.BlockSpec((2, 128), index_map)],
        out_specs=pl.BlockSpec((2, 128), index_map),
        out_shape=jax.ShapeDtypeStruct((8, 128), jnp.float32),
        interpret=True,
    )(x)

Index map function index_map at scripts/gen_mistakes.py:28 for x_ref must return 2 values to match block_shape=(2, 128). Currently returning 1 values.

the fix

def fixed_index_map_rank():
    def kernel(x_ref, o_ref):
        o_ref[...] = x_ref[...] * 2.0

    def index_map(i):
        return (i, 0)  # one entry per block dimension

    x = jnp.ones((8, 128), dtype=jnp.float32)
    return pl.pallas_call(
        kernel,
        grid=(4,),
        in_specs=[pl.BlockSpec((2, 128), index_map)],
        out_specs=pl.BlockSpec((2, 128), index_map),
        out_shape=jax.ShapeDtypeStruct((8, 128), jnp.float32),
        interpret=True,
    )(x)

exhibit
02/06

out_shape disagrees with what the kernel writes

out_shape isn't a description of what the kernel reads, it's the caller's promise about the physical buffer the kernel writes into. Pallas allocates the output ref from that shape and then checks every write against it, so when the kernel concatenates two (8, 128) blocks into a (8, 256) value and tries to store it in a ref sized (8, 128), the write fails. The error names both shapes side by side, ref shape and value shape, which is the fastest way to spot this class of bug. The fix is to size out_shape from what the kernel actually produces, not from the shape of its inputs.

the failing code

def broken_out_shape():
    def kernel(x_ref, o_ref):
        # writes a (8, 256) value into a ref the caller sized as (8, 128)
        o_ref[...] = jnp.concatenate([x_ref[...], x_ref[...]], axis=1)

    x = jnp.ones((8, 128), dtype=jnp.float32)
    return pl.pallas_call(
        kernel,
        out_shape=jax.ShapeDtypeStruct((8, 128), jnp.float32),
        interpret=True,
    )(x)

Invalid shape for `swap`. Ref shape: (8, 128). Expected shape: (8, 128). Value shape: (8, 256). Transforms: (NDIndexer(indices=(Slice(start=0, size=8, stride=1), Slice(start=0, size=128, stride=1)), shape=(8, 128), int_indexer_shape=(), validate=False),).

the fix

def fixed_out_shape():
    def kernel(x_ref, o_ref):
        o_ref[...] = jnp.concatenate([x_ref[...], x_ref[...]], axis=1)

    x = jnp.ones((8, 128), dtype=jnp.float32)
    return pl.pallas_call(
        kernel,
        out_shape=jax.ShapeDtypeStruct((8, 256), jnp.float32),  # matches the write
        interpret=True,
    )(x)

exhibit
03/06

fori_loop carry structure mismatch

jax.lax.fori_loop traces the loop body once and reuses that trace for every iteration, so the carry it returns has to match the structure of the carry it received, same number of elements, same types, every time through. This body takes a two-element tuple in and returns a three-element tuple, adding an extra value to what it hands back. The error names the exact mismatch, input carry has length 2 and output carry has length 3, which is enough on its own to tell you where to look. Whatever a loop body accepts as carry, it must return in that same shape, nothing added and nothing dropped.

the failing code

def broken_carry_structure():
    def kernel(x_ref, o_ref):
        def body(i, carry):
            acc, count = carry
            return acc + x_ref[i, :], count + 1, acc  # bug: returns 3 values, carry is 2

        acc, count, _ = jax.lax.fori_loop(0, 8, body, (jnp.zeros((128,), jnp.float32), 0))
        o_ref[...] = acc

    x = jnp.ones((8, 128), dtype=jnp.float32)
    return pl.pallas_call(
        kernel,
        out_shape=jax.ShapeDtypeStruct((128,), jnp.float32),
        interpret=True,
    )(x)

scan body function carry input and carry output must have the same pytree structure, but they differ: The input carry component loop_carry[1] is a tuple of length 2 but the corresponding component of the carry output is a tuple of length 3, so the lengths do not match. Revise the function so that the carry output has the same pytree structure as the carry input.

the fix

def fixed_carry_structure():
    def kernel(x_ref, o_ref):
        def body(i, carry):
            acc, count = carry
            return acc + x_ref[i, :], count + 1  # same structure in and out

        acc, count = jax.lax.fori_loop(0, 8, body, (jnp.zeros((128,), jnp.float32), 0))
        o_ref[...] = acc

    x = jnp.ones((8, 128), dtype=jnp.float32)
    return pl.pallas_call(
        kernel,
        out_shape=jax.ShapeDtypeStruct((128,), jnp.float32),
        interpret=True,
    )(x)

exhibit
04/06

reading a ref with wrong-rank indexing

A ref only accepts as many indices as it has dimensions, and that dimension count comes from the block shape declared in the BlockSpec, not from the original array before blocking. Here x_ref is a 2D block, so indexing it with three entries, `x_ref[:, :, 0]`, asks for an axis that doesn't exist. The error reports this literally: indices must not be longer than shape, and then prints both lengths so you can compare them. When you write a ref index, count against the block's own rank, since that's the shape the kernel actually sees.

the failing code

def broken_ref_rank_indexing():
    def kernel(x_ref, o_ref):
        # x_ref is a 2D block; a three-index read assumes a rank it doesn't have
        o_ref[...] = x_ref[:, :, 0] * 2.0

    x = jnp.ones((8, 128), dtype=jnp.float32)
    return pl.pallas_call(
        kernel,
        out_shape=jax.ShapeDtypeStruct((8, 128), jnp.float32),
        interpret=True,
    )(x)

`indices` must not be longer than `shape`: indices=(slice(None, None, None), slice(None, None, None), 0), shape=(8, 128)

the fix

def fixed_ref_rank_indexing():
    def kernel(x_ref, o_ref):
        o_ref[...] = x_ref[:, :] * 2.0  # two indices for a 2D ref

    x = jnp.ones((8, 128), dtype=jnp.float32)
    return pl.pallas_call(
        kernel,
        out_shape=jax.ShapeDtypeStruct((8, 128), jnp.float32),
        interpret=True,
    )(x)

exhibit
05/06

index_map written for the wrong grid rank

Pallas calls index_map once per grid step and passes it one positional argument per grid dimension, so a 1D grid, declared as grid=(2,), only ever supplies one argument. This index_map is written for two grid axes, `index_map(i, j)`, but only `i` ever arrives. The error is plain Python surfacing through Pallas: missing 1 required positional argument, 'j', which is the same message you'd get calling any function with too few arguments. Match the parameter count of index_map to the number of entries in your grid tuple, not to the rank of the array you're blocking.

the failing code

def broken_grid_index_map_rank():
    def kernel(x_ref, o_ref):
        o_ref[...] = x_ref[...] * 2.0

    def index_map(i, j):
        return (i, j)

    x = jnp.arange(8 * 128, dtype=jnp.float32).reshape(8, 128)
    return pl.pallas_call(
        kernel,
        grid=(2,),  # declared a 1D grid...
        in_specs=[pl.BlockSpec((4, 128), index_map)],  # ...but index_map wants two grid axes
        out_specs=pl.BlockSpec((4, 128), lambda i: (i, 0)),
        out_shape=jax.ShapeDtypeStruct((8, 128), jnp.float32),
        interpret=True,
    )(x)

broken_grid_index_map_rank.<locals>.index_map() missing 1 required positional argument: 'j'

the fix

def fixed_grid_index_map_rank():
    def kernel(x_ref, o_ref):
        o_ref[...] = x_ref[...] * 2.0

    def index_map(i):
        return (i, 0)  # matches the declared 1D grid

    x = jnp.arange(8 * 128, dtype=jnp.float32).reshape(8, 128)
    return pl.pallas_call(
        kernel,
        grid=(2,),
        in_specs=[pl.BlockSpec((4, 128), index_map)],
        out_specs=pl.BlockSpec((4, 128), lambda i: (i, 0)),
        out_shape=jax.ShapeDtypeStruct((8, 128), jnp.float32),
        interpret=True,
    )(x)

exhibit
06/06

dtype mismatch on accumulation

Writing into a ref checks dtype as strictly as it checks shape, and ordinary arithmetic in JAX doesn't upcast a low-precision value just because you added a Python float to it. x_ref holds bfloat16, and `x_ref[...] + 1.0` stays bfloat16 under JAX's weak-type promotion rules, so storing that straight into an f32 output ref fails at the write. The error states both dtypes outright, ref dtype float32 against value dtype bfloat16, which tells you the fix without needing to read further into the trace. Before a store, cast explicitly to the ref's declared dtype with .astype rather than counting on the arithmetic to do it for you.

the failing code

def broken_dtype_accumulation():
    def kernel(x_ref, o_ref):
        # x_ref is bf16; o_ref is f32. Storing the bf16 sum straight into
        # the f32 ref, with no astype.
        o_ref[...] = x_ref[...] + 1.0

    x = jnp.ones((8, 128), dtype=jnp.bfloat16)
    return pl.pallas_call(
        kernel,
        out_shape=jax.ShapeDtypeStruct((8, 128), jnp.float32),
        interpret=True,
    )(x)

Invalid dtype for `swap`. Ref dtype: float32. Value dtype: bfloat16.

the fix

def fixed_dtype_accumulation():
    def kernel(x_ref, o_ref):
        o_ref[...] = (x_ref[...] + 1.0).astype(jnp.float32)

    x = jnp.ones((8, 128), dtype=jnp.bfloat16)
    return pl.pallas_call(
        kernel,
        out_shape=jax.ShapeDtypeStruct((8, 128), jnp.float32),
        interpret=True,
    )(x)