Aliasing, and the debugging toolkit

input_output_aliasesDonates an input buffer to an output so the kernel updates it in place instead of allocating a copy.taught in /l/pallas → tells the compiler that an output can reuse an input's buffer instead of allocating a new one. Pass the argument index pairs, and the kernel writes its result into memory the caller already owned: an in-place update rather than a fresh allocation. This is donation, the same idea jax.jit uses at a coarser grain, declared explicitly here at the kernel level, which matters whenever the array involved is large enough that a second allocation is the cost you actually want to avoid, not the compute itself.

def scale_rows_kernel(x_ref, s_ref, o_ref):
    o_ref[...] = x_ref[...] * s_ref[0]

def scale_rows_inplace(x, s, rows=64):
    n, d = x.shape
    return pl.pallas_call(
        scale_rows_kernel,
        grid=(n // rows,),
        in_specs=[pl.BlockSpec((rows, d), lambda i: (i, 0)),
                  pl.BlockSpec(memory_space=MS.SMEM)],
        out_specs=pl.BlockSpec((rows, d), lambda i: (i, 0)),
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        input_output_aliases={0: 0},
        interpret=True,
    )(x, s)

Aliasing changes what happens to a buffer, not what your kernel logic looks like, which is exactly why it deserves a debugging toolkit alongside it. A mistake in the aliasing declaration corrupts an input silently, since the caller's original array is now the same memory the kernel is writing into, and nothing about the kernel's own output looks wrong when that happens.

interpret=TrueRuns kernel logic on any machine for correctness work. It ignores memory spaces and never sees compile-time errors like lattice or VMEM violations.taught in /l/pallas → runs the kernel's logic as ordinary JAX, anywhere, without a TPU. That makes it the fast path for checking that your indexing and math are correct, since you get a normal Python traceback instead of a compiler error buried under a lowering pass. It also means it ignores memory spaces entirely: 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 →, SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu →, and ANYThe memory space that tells the compiler not to place a ref at all: a promise that the kernel will move the data itself with a manual DMA.taught in /l/pallas → all collapse into whatever plain JAX would do, so interpret mode cannot catch a lattice violation or a VMEM overflow. Both are properties of a memory layout the interpreter never builds in the first place, because there is no VMEM to overflow when the code underneath is just running as JAX.

debug=True catches what interpret mode cannot, because it runs the real lowering path instead of skipping it. It prints the kernel's jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → and the MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → module the compiler derives from it, the actual lowered representation, at the point where a lattice violation or a 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 would surface as a compile error rather than silently passing. Reading that printed module is slower going than reading a Python traceback, but it is the only place a memory-space failure actually shows up, since interpret mode never produces one to read.

# what Pallas built, printed at lowering time (no run needed)
jax.jit(f).lower(*args)   # on a TPU runtime, debug=True prints here

# on a machine with no TPU, cross-lower for the tpu platform; the
# Mosaic passes ship in jaxlib, so this works anywhere
import jax.export
jax.export.export(jax.jit(f), platforms=["tpu"])(*args)

Between the two flags, the workflow is straightforward: run interpret=TrueRuns kernel logic on any machine for correctness work. It ignores memory spaces and never sees compile-time errors like lattice or VMEM violations.taught in /l/pallas → first, while you are still checking that the kernel computes the right values at all, then drop it and turn on debug=True once you are chasing a compile-time failure that only shows up in the real memory layout. Neither flag replaces the other, because they are answering different questions: one asks whether the algebra is right, the other asks whether the schedule and the memory layout it produces are legal at all.

That split is worth keeping in mind whenever a kernel passes under interpret=TrueRuns kernel logic on any machine for correctness work. It ignores memory spaces and never sees compile-time errors like lattice or VMEM violations.taught in /l/pallas → and then fails to lower on real hardware. The instinct is to suspect the algebra again, since that is what you just finished checking, but a failure that only appears once memory spaces are real almost always belongs to the second question, not the first: a block that does not fit 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 →, or a shape that violates the tile rule from section one, neither of which interpret mode was ever in a position to see.