What does pallas_call actually do when you call it?

Almost nothing at call time. pallas_call itself only sorts out its arguments: it flattens scratch_shapes, and then either builds a GridSpec out of grid, in_specs, out_specs and the flat scratch shapes, or, if you passed grid_spec yourself, checks that you left all four of those at their defaults and raises a specific ValueError naming whichever one you did not. Then it hands everything to a private function that returns a jitted wrapper, and returns that wrapper to you. Nothing is traced until you call the wrapper with real arguments. Inside the wrapper, the arguments are flattened with their key paths so error messages can name them, the out_shape tree becomes output avals, and get_grid_mapping turns the specs into a GridMapping plus the list of Ref avals the kernel will receive. Only then is your kernel traced to a jaxpr, and only then is the primitive bound. So a mistake in your BlockSpec surfaces on the first call, not on the pallas_call line.

Why does Pallas say my kernel function captures constants?

Because tracing the kernel produced a closed jaxpr with constants in it, and the tracer checks for that explicitly. The trace helper takes the constants out of the closed jaxpr, filters out any that are Refs, and if anything is left it pretty-prints those avals into a ValueError saying the kernel captures them and that you should pass them as inputs. The reason is not style. The backends consume the jaxpr, not your Python, and a captured array has no BlockSpec, so nothing tells the compiler how to slice it per grid step or which memory space it should live in. There is a second, separate check in the same function: if the kernel returns anything other than None, you get an error naming the PyTree it returned, because a Pallas kernel communicates only by writing into the output Refs it was given.

What happens if an input has no BlockSpec?

It is not an error, and the input is not skipped. A BlockSpec with no block_shape gets the whole array shape as its block, and a BlockSpec with no index_map gets a default one that returns a tuple of zeros as long as the array has dimensions. So the kernel sees the entire array in one block, at block index zero, on every grid step. The one thing that differs is the memory space, and it depends on how you omitted the spec. Leaving in_specs off entirely gives every input a BlockSpec whose memory space is DEFAULT. Passing a sequence of specs with pl.no_block_spec in one slot gives that one input a BlockSpec whose memory space is ANY instead, which means the block stays wherever the compiler put it and the kernel has to copy it in itself. That difference is the reason a kernel that works with no in_specs at all can fail when you spell out the other specs and leave one as no_block_spec.

How does vmap of a pallas_call work?

It grows the grid instead of running a loop. The batching rule prepends the batch size to the grid as a new, unnamed dimension, records it in vmapped_dims with every existing entry shifted up by one, and rewrites each BlockSpec index map so that it takes the new leading grid index and inserts it at the batched dimension of that operand. Block shapes gain a squeezed dimension at the same position, and array avals gain the real one, so the kernel body never changes. Cost estimates are multiplied by the axis size unless the axis size is a dynamic dimension. There are three ways out of that path. If nothing is batched, it binds the original primitive unchanged. If the axis size is one, it squeezes the batch dimension out, calls the kernel once, and expands the result back. And if the batching would have to cross a dynamic grid bound or a scalar prefetch operand of size greater than one, it falls back to a real fori_loop that dynamic-slices each input and dynamic-updates an output allocation, which is slower but correct.

Why is my Pallas GPU kernel compiled by Mosaic GPU instead of Triton?

Because that is the default at this commit and Triton is now opt-in. The shared GPU lowering picks Mosaic GPU when compiler_params is a Mosaic GPU CompilerParams, or when compiler_params is None and the jax_pallas_use_mosaic_gpu config flag is true, and that flag defaults to true. It picks Triton when compiler_params is a Triton CompilerParams, or when compiler_params is None and that same flag is false. Both branches now warn: using pl.pallas_call for a Mosaic GPU kernel is deprecated in favour of plgpu.kernel, and the Triton backend is deprecated altogether. Two more details decide real cases. Mosaic GPU raises on ROCm and tells you to pass Triton compiler params instead, and before either built-in backend is even imported the rule asks a registry whether some other backend has claimed this CompilerParams type for this platform, which is how an out-of-tree backend takes over.

pallas_call, Line by Line: One Call, Three Calling Conventions

Every Pallas kernel on every backend enters JAX through one 1381-line Python file. This walks it whole, at a pinned commit of jax-ml/jax. The file declares a single primitive, then spends the rest of itself teaching JAX what that primitive means under each transform: an abstract evaluation that decides the output types and refuses a non-manual mesh, a lowering rule that swaps its own body for a scan when you ask for interpret mode, a batching rule that answers vmap by growing the grid rather than looping, and a public function whose real work is turning a pile of keyword arguments into one GridMapping. The signature exhibit runs the real slot arithmetic out of that file: dial the grid, the scalar prefetch operands, the inputs, the outputs and the scratch, and watch the three orderings the primitive, the kernel and the index map each demand.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of pallas_call, Line by Line: One Call, Three Calling Conventions — built to be understood, not skimmed.

Questions

What does pallas_call actually do when you call it?: Almost nothing at call time. pallas_call itself only sorts out its arguments: it flattens scratch_shapes, and then either builds a GridSpec out of grid, in_specs, out_specs and the flat scratch shapes, or, if you passed grid_spec yourself, checks that you left all four of those at their defaults and raises a specific ValueError naming whichever one you did not. Then it hands everything to a private function that returns a jitted wrapper, and returns that wrapper to you. Nothing is traced until you call the wrapper with real arguments. Inside the wrapper, the arguments are flattened with their key paths so error messages can name them, the out_shape tree becomes output avals, and get_grid_mapping turns the specs into a GridMapping plus the list of Ref avals the kernel will receive. Only then is your kernel traced to a jaxpr, and only then is the primitive bound. So a mistake in your BlockSpec surfaces on the first call, not on the pallas_call line.
Why does Pallas say my kernel function captures constants?: Because tracing the kernel produced a closed jaxpr with constants in it, and the tracer checks for that explicitly. The trace helper takes the constants out of the closed jaxpr, filters out any that are Refs, and if anything is left it pretty-prints those avals into a ValueError saying the kernel captures them and that you should pass them as inputs. The reason is not style. The backends consume the jaxpr, not your Python, and a captured array has no BlockSpec, so nothing tells the compiler how to slice it per grid step or which memory space it should live in. There is a second, separate check in the same function: if the kernel returns anything other than None, you get an error naming the PyTree it returned, because a Pallas kernel communicates only by writing into the output Refs it was given.
What happens if an input has no BlockSpec?: It is not an error, and the input is not skipped. A BlockSpec with no block_shape gets the whole array shape as its block, and a BlockSpec with no index_map gets a default one that returns a tuple of zeros as long as the array has dimensions. So the kernel sees the entire array in one block, at block index zero, on every grid step. The one thing that differs is the memory space, and it depends on how you omitted the spec. Leaving in_specs off entirely gives every input a BlockSpec whose memory space is DEFAULT. Passing a sequence of specs with pl.no_block_spec in one slot gives that one input a BlockSpec whose memory space is ANY instead, which means the block stays wherever the compiler put it and the kernel has to copy it in itself. That difference is the reason a kernel that works with no in_specs at all can fail when you spell out the other specs and leave one as no_block_spec.
How does vmap of a pallas_call work?: It grows the grid instead of running a loop. The batching rule prepends the batch size to the grid as a new, unnamed dimension, records it in vmapped_dims with every existing entry shifted up by one, and rewrites each BlockSpec index map so that it takes the new leading grid index and inserts it at the batched dimension of that operand. Block shapes gain a squeezed dimension at the same position, and array avals gain the real one, so the kernel body never changes. Cost estimates are multiplied by the axis size unless the axis size is a dynamic dimension. There are three ways out of that path. If nothing is batched, it binds the original primitive unchanged. If the axis size is one, it squeezes the batch dimension out, calls the kernel once, and expands the result back. And if the batching would have to cross a dynamic grid bound or a scalar prefetch operand of size greater than one, it falls back to a real fori_loop that dynamic-slices each input and dynamic-updates an output allocation, which is slower but correct.
Why is my Pallas GPU kernel compiled by Mosaic GPU instead of Triton?: Because that is the default at this commit and Triton is now opt-in. The shared GPU lowering picks Mosaic GPU when compiler_params is a Mosaic GPU CompilerParams, or when compiler_params is None and the jax_pallas_use_mosaic_gpu config flag is true, and that flag defaults to true. It picks Triton when compiler_params is a Triton CompilerParams, or when compiler_params is None and that same flag is false. Both branches now warn: using pl.pallas_call for a Mosaic GPU kernel is deprecated in favour of plgpu.kernel, and the Triton backend is deprecated altogether. Two more details decide real cases. Mosaic GPU raises on ROCm and tells you to pass Triton compiler params instead, and before either built-in backend is even imported the rule asks a registry whether some other backend has claimed this CompilerParams type for this platform, which is how an out-of-tree backend takes over.

pallas_call, Line by Line: One Call, Three Calling Conventions

Questions

Related explainers