What does interpret=True actually do in a Pallas kernel?

It replaces the backend lowering with a plain JAX one. The docstring is literal about it: the call runs as a jax.jit of a scan over the grid whose body is the kernel lowered as a JAX function, it does not require a TPU or a GPU, and it is the only way to run Pallas kernels on CPU. The implementation matches. Your kernel jaxpr has its state effects discharged into a pure function, and a while loop runs that function once per grid point, slicing a block out of each operand with lax.dynamic_slice and writing the result back with lax.dynamic_update_slice. Two things follow. The arithmetic is genuinely real, because it is the same jaxpr on the same dtypes evaluated by the same JAX, so a rounding or reduction-order surprise there is a real one. And there is no memory system at all, because there are no buffers: an operand is one flat array in the loop carry and a ref is a slice of it. VMEM has no size, a DMA is a slice, and a semaphore is a value being carried. That is why an entire class of kernel bug is invisible on this path, and it is not a defect in the interpreter.

What is the difference between interpret=True and pltpu.InterpretParams?

They are two different interpreters selected by the same argument, which is typed Any rather than bool for exactly that reason. A truthy value gets the generic HLO interpreter described above. An instance of the TPU InterpretParams class gets a second interpreter whose own docstring describes it as a way to run Pallas TPU kernels on CPU while simulating a TPU shared memory (HBM, VMEM and so on), communication (remote and local DMAs), and synchronization operations (semaphores, barriers). So rung two adds a memory model: buffers have edges, DMAs are events that complete when something waits on them, semaphores count, and a dynamic happens-before race detector can watch several simulated cores at once. What it does not add is a compiler. The Mosaic toolchain never runs on either rung, so neither one can tell you whether your kernel will be accepted, what it does to your layouts, or how fast it is. There is a third interpreter, InterpretGPUParams, for the GPU side. All three are imported inside a try, so on a machine without those backends you fall back to the generic path silently.

Why does my Pallas kernel pass in interpret mode and fail on a TPU?

Usually one of three reasons, and they are worth separating because only one of them is a bug in JAX. First, the compiler rejected it. Rung one has no compiler in it, so a message about an unsupported dtype, shape, or operation is the first time anything has had an opinion, and interpret mode never promised to predict that. Second, the kernel depended on something the interpreter defines and the chip does not. The grid promises a count of invocations, not an order: the interpreter walks it sequentially in row-major order with the last axis fastest, while on TPUs programs execute in a combination of parallel and sequential, and the documentation states outright that when multiple invocations write to the same elements of the output the result is platform dependent. Similarly the interpreter fills fresh outputs and scratch with NaN, so an accumulator you forgot to zero looks broken on a laptop and fine on a chip. Third, the results simply differ. That case the project treats as a bug: the TPU details page says that if a kernel is accepted by the compiler it must return the expected results, and asks you to compare against interpret=True and file a report if the two diverge.

What happens if a Pallas BlockSpec index map goes out of bounds?

Nothing checks it, because there is nothing to check it against: the index map is traced into its own small jaxpr and its output is arithmetic rather than an assertion. On the interpret path a Blocked dimension of size b turns a block index i into a start index of b * i, and the block is fetched with an XLA DynamicSlice. That operation does not fail on an overrunning start, it moves the start back until the window fits, a behaviour JAX documents explicitly where Pallas defends against it for pl.load. So an index map one block past the end silently reads the last in-bounds window a second time and returns an output full of finite floats with nothing to look at. Under pltpu.InterpretParams the read is checked against a simulated buffer and the out_of_bounds_reads parameter, which defaults to raise, throws instead. On a chip it is undefined: the documentation says padding values are unspecified garbage and that at least one element of every block must be within bounds, and the Mosaic compiler parameters include a switch that turns bounds checks off.

How do I debug a Pallas kernel: debug=True, pl.debug_print, or something else?

Pick by the question. debug=True prints intermediate forms, and which forms depends on the rung: on the interpret path you get exactly one printout, the discharged kernel jaxpr, while on a TPU you get two, the kernel jaxpr and then the Mosaic module after a canonicalize pass has been run over it to make it readable. pl.debug_print prints a value from inside the kernel body, with format-string rules that differ per backend: Triton takes no placeholders at all, Mosaic GPU wants one per value with no format specs, and TPU splits its rules by whether the arguments are scalars or a single vector. pl.debug_check is a permanent assertion that compiles to nothing unless jax_pallas_enable_debug_checks is set. pl.enable_poison_buffers fills scratch with NaN at allocation on a real chip, which is the behaviour interpret mode already has unconditionally. detect_races on TPU interpret mode finds data races and sets a flag you can assert on, but it needs num_cores_or_threads above its default of one to have anything to race. And JAX_DUMP_IR_TO plus JAX_DUMP_IR_MODES write the surrounding modules to disk, both of which must be set before JAX is imported.

Writing Pallas Kernels in Colab: What Interpret Mode Proves, and What It Cannot

One keyword argument runs a TPU kernel on a laptop, and almost every confusing Pallas afternoon comes from not knowing what that run checked. This reads the machinery at a pinned commit of jax-ml/jax. There are three interpreters behind one argument, not one: a generic scan over the grid that models the arithmetic and removes memory entirely, a TPU mode that simulates HBM, VMEM, DMAs and semaphores and carries a happens-before race detector, and a real chip, the only rung with a compiler that can refuse you. What each rung catches and what it structurally cannot, drawn as a walk. Which import names which chip, with the memory-space tables side by side and the SMEM collision that means scalar memory on one and shared memory on the other. What interpret=True literally executes: a while loop over the grid, dynamic slices for blocks, outputs and scratch pre-filled with NaN. The padding the interpreter adds before the loop starts, computed live from the real function. Why the grid promises a count and not an order, and the one parameter that randomizes it. Then the case worth committing to: an index map one block past the end, which raises nothing, returns finite numbers, and quietly duplicates half the output. The debugging surface as a picker, one question in and one switch out. And the six-cell notebook shape, including the reset cell that exists because the interpreter keeps its simulated memory after an exception on purpose.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of Writing Pallas Kernels in Colab: What Interpret Mode Proves, and What It Cannot — built to be understood, not skimmed.

Questions

What does interpret=True actually do in a Pallas kernel?: It replaces the backend lowering with a plain JAX one. The docstring is literal about it: the call runs as a jax.jit of a scan over the grid whose body is the kernel lowered as a JAX function, it does not require a TPU or a GPU, and it is the only way to run Pallas kernels on CPU. The implementation matches. Your kernel jaxpr has its state effects discharged into a pure function, and a while loop runs that function once per grid point, slicing a block out of each operand with lax.dynamic_slice and writing the result back with lax.dynamic_update_slice. Two things follow. The arithmetic is genuinely real, because it is the same jaxpr on the same dtypes evaluated by the same JAX, so a rounding or reduction-order surprise there is a real one. And there is no memory system at all, because there are no buffers: an operand is one flat array in the loop carry and a ref is a slice of it. VMEM has no size, a DMA is a slice, and a semaphore is a value being carried. That is why an entire class of kernel bug is invisible on this path, and it is not a defect in the interpreter.
What is the difference between interpret=True and pltpu.InterpretParams?: They are two different interpreters selected by the same argument, which is typed Any rather than bool for exactly that reason. A truthy value gets the generic HLO interpreter described above. An instance of the TPU InterpretParams class gets a second interpreter whose own docstring describes it as a way to run Pallas TPU kernels on CPU while simulating a TPU shared memory (HBM, VMEM and so on), communication (remote and local DMAs), and synchronization operations (semaphores, barriers). So rung two adds a memory model: buffers have edges, DMAs are events that complete when something waits on them, semaphores count, and a dynamic happens-before race detector can watch several simulated cores at once. What it does not add is a compiler. The Mosaic toolchain never runs on either rung, so neither one can tell you whether your kernel will be accepted, what it does to your layouts, or how fast it is. There is a third interpreter, InterpretGPUParams, for the GPU side. All three are imported inside a try, so on a machine without those backends you fall back to the generic path silently.
Why does my Pallas kernel pass in interpret mode and fail on a TPU?: Usually one of three reasons, and they are worth separating because only one of them is a bug in JAX. First, the compiler rejected it. Rung one has no compiler in it, so a message about an unsupported dtype, shape, or operation is the first time anything has had an opinion, and interpret mode never promised to predict that. Second, the kernel depended on something the interpreter defines and the chip does not. The grid promises a count of invocations, not an order: the interpreter walks it sequentially in row-major order with the last axis fastest, while on TPUs programs execute in a combination of parallel and sequential, and the documentation states outright that when multiple invocations write to the same elements of the output the result is platform dependent. Similarly the interpreter fills fresh outputs and scratch with NaN, so an accumulator you forgot to zero looks broken on a laptop and fine on a chip. Third, the results simply differ. That case the project treats as a bug: the TPU details page says that if a kernel is accepted by the compiler it must return the expected results, and asks you to compare against interpret=True and file a report if the two diverge.
What happens if a Pallas BlockSpec index map goes out of bounds?: Nothing checks it, because there is nothing to check it against: the index map is traced into its own small jaxpr and its output is arithmetic rather than an assertion. On the interpret path a Blocked dimension of size b turns a block index i into a start index of b * i, and the block is fetched with an XLA DynamicSlice. That operation does not fail on an overrunning start, it moves the start back until the window fits, a behaviour JAX documents explicitly where Pallas defends against it for pl.load. So an index map one block past the end silently reads the last in-bounds window a second time and returns an output full of finite floats with nothing to look at. Under pltpu.InterpretParams the read is checked against a simulated buffer and the out_of_bounds_reads parameter, which defaults to raise, throws instead. On a chip it is undefined: the documentation says padding values are unspecified garbage and that at least one element of every block must be within bounds, and the Mosaic compiler parameters include a switch that turns bounds checks off.
How do I debug a Pallas kernel: debug=True, pl.debug_print, or something else?: Pick by the question. debug=True prints intermediate forms, and which forms depends on the rung: on the interpret path you get exactly one printout, the discharged kernel jaxpr, while on a TPU you get two, the kernel jaxpr and then the Mosaic module after a canonicalize pass has been run over it to make it readable. pl.debug_print prints a value from inside the kernel body, with format-string rules that differ per backend: Triton takes no placeholders at all, Mosaic GPU wants one per value with no format specs, and TPU splits its rules by whether the arguments are scalars or a single vector. pl.debug_check is a permanent assertion that compiles to nothing unless jax_pallas_enable_debug_checks is set. pl.enable_poison_buffers fills scratch with NaN at allocation on a real chip, which is the behaviour interpret mode already has unconditionally. detect_races on TPU interpret mode finds data races and sets a flag you can assert on, but it needs num_cores_or_threads above its default of one to have anything to race. And JAX_DUMP_IR_TO plus JAX_DUMP_IR_MODES write the surrounding modules to disk, both of which must be set before JAX is imported.

Writing Pallas Kernels in Colab: What Interpret Mode Proves, and What It Cannot

Questions

Related explainers