Params, and the drills that fix them

Params are where a primitive keeps whatever isn't itself a traced value, the configuration that has to be fixed before the primitive can run rather than computed while it runs. dot_general's params hold dimension_numbers, the exact statement of which axes contract against which and which axes broadcast, the information that turns two arrays into a matmul, a batched matmul, or a plain contraction, depending on what's in there. A reduction like sum or max holds axes in its params, the list of dimensions being collapsed. A type conversion holds new_dtype, the target type values get cast into. None of these are variables flowing through the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →; they are settings the equation carries alongside it.

Reading params fluently is less about memorizing a table and more about repetition against real jaxprsThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →, enough of them that dimension_numbers or axes stops requiring a lookup. That's what the drills are for: decoder exercises where you predict what an equation does from its params alone, oracle checks that confirm or correct the guess, and spot-the-decision exercises that isolate exactly which param changed the behavior. None of it is complicated once you've done it a dozen times. The difficulty is entirely in not having done it yet.

The same primitives keep showing up once you move into kernel material. A Pallas kernel body still lowers through dot_general, still carries dimension_numbers, still reduces along explicit axes. Nothing about the grammar changes there, only the shapes and the register tiling underneath it. That's the payoff of getting fluent with params here: by the time you're staring at a kernel jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → with a grid and a block spec wrapped around it, the params inside are already familiar, and what's new is only the pipelining around them.