Fluency is drills, not prose.

The op reference

Back on paper to read: every op the corpus uses, one card each. Filter by name or family; each card ends with the misread to avoid.

transpose shape

Permutes an array's axes according to permutation, a tuple stating where each output axis comes from in the input. It looks free because no arithmetic happens, but on TPU it is not: a transpose that swaps the last two lane dimensions forces the compiler to physically move data across the (8, 128) tiling lattice, and that shuffle costs real HBM or VMEM bandwidth. The misread is treating every transpose as metadata-only; watch for one sitting right before a dot_general in an HLO dump, since XLA sometimes fuses it into the matmul and sometimes cannot.

dot_general compute

The one primitive every matmul, batched matmul, and attention score lowers to. Its dimension_numbers argument is a pair of pairs: which axes contract, summed away like the shared K dimension in a matmul, and which axes batch, kept independent like the head dimension in multi-head attention, each given as (lhs_axes, rhs_axes). The misread is skimming past dimension_numbers=(([1], [0]), ([], [])) as decoration when it is the entire semantics of the op; decode contracting axes first, batch axes second, and the output shape falls out on its own.

reduce_max reduction

Collapses one or more axes of an array down to their maximum, with axes naming which ones disappear. In softmax it produces the per-row max used for numerical stability before the exponential runs. The misread is assuming the reduced axis stays as size 1 for broadcasting; reduce_max drops the axis entirely, and the broadcast_in_dim you see right after it in a jaxpr is what puts a size-1 dimension back.

reduce_sum reduction

Sums an array over the axes in axes, dropping them from the output shape exactly like reduce_max does. It shows up wherever a normalization needs a denominator, softmax's row sum, LayerNorm's mean and variance, and wherever a loss reduces a batch to a scalar. The one thing that trips people up: reduce_sum on a low-precision input often runs after a convert_element_type up to f32, because summing many bf16 values in bf16 accumulates rounding error fast.

broadcast_in_dim shape

Takes an array and places it into a larger shape, with broadcast_dimensions saying which output axes the input's existing axes map to; any output axis left out gets filled by repeating the value along it. It is the primitive behind ordinary NumPy broadcasting once JAX has made every implicit dimension explicit. The misread is assuming it always adds axes at the end; a bias vector broadcasting over rows and a scalar broadcasting into every element use the same primitive with very different broadcast_dimensions, and reading that argument is the only way to tell which is which.

convert_element_type shape

Casts an array from one dtype to another, with new_dtype naming the target. It appears constantly and mostly invisibly: accumulating a reduction in f32 before casting back to bf16, staging an int8 weight up to bf16 before a matmul, or JAX inserting a cast you did not write because two operands disagreed in precision. The misread is skimming past it as boilerplate; a convert_element_type sitting between a reduction and the next op is usually the whole reason a kernel is numerically stable, and one sitting inside a hot loop is usually a cast worth hoisting out.

exp compute

The elementwise exponential, applied independently to every element. It is one of the more expensive elementwise ops on the VPU because transcendentals lack the one-cycle hardware path that add or multiply get, so a kernel with a lot of exp, softmax, GELU's tanh approximation, any sigmoid, spends real VPU cycles even though nothing there touches the MXU. The misread is treating all elementwise ops as equally cheap; exp and tanh are the ones worth counting separately when estimating a kernel's compute time.

sub compute

Elementwise subtraction of two arrays, or an array and a scalar, of matching or broadcastable shape. By itself it is trivial, but in a jaxpr a sub right after a reduce_max and right before an exp is the signature of numerically stable softmax: subtracting the row max before exponentiating so the largest exponent computed is exp(0) instead of something that overflows. Reading sub in isolation misses that; read it next to its neighbors.

div compute

Elementwise division. It ends softmax, dividing by the row sum, ends normalization layers, dividing by the standard deviation, and shows up anywhere a ratio is computed directly instead of via a reciprocal. The misread is assuming div and multiplying by a reciprocal cost the same; division is more expensive than a multiply on the VPU, which is exactly why performance-sensitive kernels compute rsqrt once and multiply by it repeatedly instead of dividing repeatedly.

add compute

Elementwise addition, the most common op in any jaxpr because it carries bias terms, residual connections, and accumulator updates. Nothing about it is subtle in isolation, but a chain of adds inside a remat block or a grid-carried accumulator is usually the actual computation a kernel exists to do; the surrounding ops are staging. The misread is skimming past every add as filler when a few of them are the point.

mul compute

Elementwise multiplication, functionally symmetric with add but with one sharp edge: a mul by a broadcasted scalar constant, a GELU coefficient, a scaling factor, an inverse temperature, is cheap on the VPU, while a mul between two full-size arrays of the same shape is not free just because it is 'only' elementwise. It still touches every element, so it scales with array size the same way a reduction does; only matmul-shaped contractions get the MXU's throughput advantage.

max compute

Elementwise maximum between two arrays or an array and a scalar. It is easy to confuse with reduce_max, a different primitive: max compares two same-shaped things element by element and keeps the larger, reduce_max collapses one array along an axis. Seeing max -inf b in a jaxpr is not a reduction; it is guarding against a row of all-masked scores producing a -inf max that would turn the softmax into 0/0.

select_n control

A generalized where: given an integer index array and N candidate arrays of the same shape, it picks element-wise from candidate index at each position. Two candidates, a boolean index, is the common case and is exactly how masking works: a causal mask compiles down to a select_n choosing between the real score and -inf at each position, driven by a comparison built from iota. The misread is looking for a where primitive in the jaxpr and not finding one; both jnp.where and lax.select lower to select_n.

iota shape

Generates an array of sequential integers along one axis, with dimension naming which axis counts up and shape naming the full output shape. On its own it looks like nothing, but two iotas compared against each other, one counting rows, one counting columns, is how a causal mask, a one-hot index, or a scatter target gets built without any Python control flow. The misread is expecting a mask to arrive as a boolean parameter; more often it is constructed on the fly from a pair of iotas and a comparison.

reshape shape

Reinterprets an array's shape without changing its data, as long as the total element count matches. On TPU, whether a reshape is actually free depends on the tiling lattice: reshapes that keep the last two dimensions' packing intact are metadata-only, but a reshape that merges or splits the lane dimension can force a physical relayout, the same cost class as a transpose. The misread is assuming reshape is always free because that is true off-TPU and false the moment tiling gets involved.

concatenate shape

Joins arrays along an existing axis named by dimension, producing one array whose size along that axis is the sum of the inputs'. It is the primitive behind splitting rotary embeddings into two halves, rotating them, and stitching them back together, and behind any kernel that builds an output from pieces computed separately. The misread is assuming concatenation of small pieces is cheap; each piece still needs its own memory movement, so several small concatenates can add up to more HBM traffic than one contiguous write.

slice shape

Extracts a fixed, statically known rectangular region from an array using start_indices, limit_indices, and strides. Because the indices are compile-time constants, slice is what an associative-scan unrolling produces: a log-depth tree of strided slices and adds, with no scan primitive anywhere in the jaxpr. The misread is expecting lax.associative_scan to leave a scan primitive behind; XLA unrolls it into exactly this shape once the sequence length is known.

dynamic_slice shape

Like slice, but the start offsets are runtime values instead of compile-time constants, which is what makes it the primitive behind gather-like indexing: reading one page from a paged-attention KV cache, or reading a runtime-chosen block inside a Pallas kernel. The misread is treating it as interchangeable with slice; a dynamic_slice cannot be constant-folded, and its bounds must be clamped to stay in range, which is why clamping arithmetic often feeds its start-index arguments.

argmax reduction

Returns the index of the maximum value along axes, not the value itself. It is the routing primitive behind mixture-of-experts dispatch, choosing which expert each token goes to, and behind any greedy decode step. The misread is forgetting that argmax breaks ties by taking the first maximal index, which matters when a routing kernel is differential-tested against a reference and two experts score exactly equal.

reduce reduction

The general-purpose reduction primitive: it takes an arbitrary binary combiner function and an initial value, then folds an array down along the named axes using that combiner. reduce_max and reduce_sum are both just reduce with a fixed combiner baked in and given their own primitive name for the compiler's benefit. The misread is expecting to see a bare reduce in most jaxprs; JAX prefers the specialized forms, so a bare reduce usually signals a custom combiner, like the online-softmax combine, with no specialized primitive of its own.

custom_jvp_call control

Wraps a function with a hand-supplied forward-mode derivative rule instead of letting JAX differentiate the traced primitives underneath it automatically. It shows up around numerically delicate functions, erf, some GELU formulations, where the naive derivative through the primitive chain is less accurate or less stable than a rule someone derived by hand. The misread is assuming the wrapped function's body tells you its gradient; with custom_jvp_call, the body computes the forward value and a separate rule, invisible in this jaxpr, computes the gradient.

remat control

Marks a block of computation to be recomputed during the backward pass instead of having its intermediates saved from the forward pass. jax.checkpoint is what produces it, and the trade is explicit: less memory held between forward and backward, more FLOPs spent recomputing. The misread is assuming a remat block changes what gets computed; it changes only when the forward's intermediates are computed, once eagerly and once again during backprop, not the values themselves.

scan control

Runs a function repeatedly over the leading axis of an array, carrying an accumulator forward from one iteration to the next, JAX's primitive for a loop with a known, static trip count. It is the shape a scanned transformer layer stack takes, looping over layers without Python unrolling them one by one. The misread: lax.associative_scan, despite the name, is a different code path and leaves no scan primitive in the jaxpr at all, it unrolls into slices and adds instead; a literal scan primitive means lax.scan, not any function with 'scan' in its name.

while control

Runs a loop whose trip count is not known until runtime, driven by a condition function that reads the loop state and returns a boolean. It is what a kernel needs when the number of iterations depends on data, a page table's length, a convergence check, rather than on a shape known at trace time. The misread is expecting while to be common in kernel jaxprs; because trip count usually is known from shapes, most kernels use scan or a Pallas grid instead, and a while you do find is worth reading closely for what makes it truly data-dependent.

cond control

Runtime branching between two traced functions, chosen by a predicate, with both branches traced ahead of time even though only one executes. That last part is the sharp edge: both branches must produce the same output shape and dtype, and both get compiled, so cond costs code size and any FLOPs a branch needs to allocate even when it never runs. The misread is treating cond like Python's if; pl.when inside a Pallas kernel is closer to that intuition, since it can skip work rather than merely skip using the result.

psum distributed

An all-reduce sum across a named mesh axis, appearing in a jaxpr whenever a dot_general's contracting dimension is sharded and each device only holds a partial sum that needs combining. Seeing psum2, the two-argument, axis-index-groups form, inside a shard_map block is what a sharded matmul looks like once GSPMD's implicit collectives get made explicit. The misread is thinking of psum as a special communication step bolted onto the matmul; it is the correctness-restoring last piece of the same computation, not an optional extra.

all_gather distributed

Collects shards of an array spread across devices along a mesh axis into one full copy on every device. Physically, on a TPU ring, it runs as N-1 hops: each step, every device forwards the shard it just received to its neighbor while also computing on what it already has, which is why a correct all-gather kernel is really a ppermute loop with a grid-carried accumulator, the same double-buffering pattern as any Pallas pipeline. The misread is treating all_gather as an atomic primitive with no internal structure; underneath, it is a schedule you can build yourself from remote DMAs.

ppermute distributed

Sends each device's data to a fixed permutation of neighbors along a mesh axis in one step, the primitive a ring topology's single hop lowers to. Chained across N-1 steps with an accumulator carried between them, it is exactly how a ring all-gather gets built from scratch. The misread is expecting ppermute to name which physical link carries the data; it names only the logical source-to-destination mapping, and the mesh topology decides whether that mapping happens to land on a physical neighbor or a longer hop.

one_hot shape

Not a single XLA primitive: jax.nn.one_hot lowers to a broadcast_in_dim of the index array, an iota counting the target classes, an eq comparison between the two, and a convert_element_type to the output dtype. What matters here is what happens next: multiplying a one-hot matrix into a dot_general is how MoE dispatch and gather-by-index both get implemented as dense matmuls instead of a scatter, trading wasted FLOPs on the zero entries for a shape the MXU can execute. The misread is expecting to find a gather primitive in a dispatch kernel's jaxpr; production MoE code more often routes through this one-hot-times-dot pattern precisely because it is MXU-friendly.

rsqrt compute

Computes 1 / sqrt(x) elementwise, and exists as its own primitive rather than a composition of sqrt and div because the combination is one hardware-supported operation cheaper than a real division. It is the last step of every normalization layer, applied to the variance before the elementwise scale. The misread is writing 1 / jnp.sqrt(x) and assuming JAX always recognizes and fuses it back into rsqrt; it usually does, but checking the actual jaxpr is the only way to be sure for a given JAX version.

log compute

The elementwise natural logarithm, another transcendental with real VPU cost. Its main appearance in this track is inside logsumexp, computed as the row max plus the log of the sum of rescaled exponentials, the numerically stable way to compute a log-sum-exp without ever materializing the unstable direct form. The misread is assuming log and exp always appear together in a 1:1 ratio; a stable softmax needs exp but no log at all, while log-sum-exp needs exactly one log at the end regardless of how many terms were summed.

erf compute

The Gauss error function, computed elementwise, and the exact ingredient in GELU's precise definition: x * 0.5 * (1 + erf(x / sqrt(2))). Because erf is expensive and has no simple closed form, most production GELU implementations swap it for the tanh approximation instead, which is why the MLP jaxpr you actually trace in this track shows tanh, not erf, even though both compute 'GELU'. The misread is expecting the two formulations to be numerically identical; they agree closely but not exactly, which matters at the tolerances a differential test checks.

tanh compute

The elementwise hyperbolic tangent, another VPU transcendental. It appears directly in the tanh-approximate form of GELU, 0.5 * x * (1 + tanh(...)), the version XLA actually lowers to when you write jax.nn.gelu, and less obviously inside normalization math on some backends. The misread is assuming tanh's presence always means an activation function; in a jaxpr it can also show up as part of a fused epilogue the compiler chose, not necessarily code you wrote yourself.

no ops match