Chapter 1: Advanced Transformations and Control Flow

Course: "Advanced JAX: Performance, Optimization, and Scale", Chapter 1 (course home) Environment: JAX / jaxlib 0.12.0.1, Python 3.14, CPU only. Every conclusion in this tutorial was run and verified on this version.

How to read this: This tutorial follows the 8 pages of Chapter 1 of the course. The explanations and code were rewritten from scratch, with extra material added. Outdated or error-prone patterns go in # ❌ cells, followed right away by the # ✅ recommended version.

Learning goals¶

After this chapter you should be able to:

1. Explain the tracing model behind jit / grad / vmap and what they mean when composed;
2. Use lax.scan fluently (including reverse, unroll, length, xs=None) and read the error you get when carry types don't match;
3. Pick between lax.cond / lax.switch / lax.select_n / jnp.where for a given situation, and know that cond turns into select under vmap;
4. Write loops with a dynamic stopping condition using lax.while_loop, know that it does not support reverse-mode differentiation, and know what to use instead;
5. Understand how control flow behaves, and what it costs, when combined with jit / grad / vmap / jax.checkpoint;
6. Handle variable-length data with padding + mask and jax.ops.segment_sum;
7. Understand closures and staging: how constants get embedded in the program, what constvars in make_jaxpr are, and why you shouldn't capture large arrays in closures;
8. Implement a GRU, an LSTM and an ODE integrator with lax.scan, and cross-check them value by value against Python-loop versions.

Sections¶

How this connects to the Basics part¶

• Tracing, recompilation, static arguments: Basics 02 JIT compilation. Section 2.4 there already shows minimal examples of fori_loop, lax.switch and while_loop; this chapter goes much deeper.
• Reverse mode, jvp / vjp, NaN gradients from jnp.where and the "double where" trick: Basics 03 Automatic differentiation with grad
• in_axes / out_axes and batching semantics: Basics 04 Automatic vectorization with vmap
• Managing random keys, purely functional state: Basics 01 Introduction to JAX, Basics 06 State management

⚠️ A heads-up about this environment: the CPU backend of jaxlib 0.12.0.1 has a bug that reproduces reliably. If a jitted function captures a constant matrix through a closure, uses it as the left operand of a matmul (W @ x) and then reduces, the result silently comes out wrong. The trigger threshold depends on the shape; 128×128 already triggers it. So in this tutorial weights and data are always passed in as function arguments, and key numbers are cross-checked against NumPy float64. For the symptoms, the trigger boundary and workarounds, see "Known issues in this environment" in Basics 00 Overview; for the root-cause analysis, see Advanced Chapter 2 §2.3.

1.1 Core transformations review: jit, grad, vmap¶

jax.jit, jax.grad and jax.vmap are JAX's three most basic function transformations: they take a Python function and return a new function. The lax.scan, lax.cond and lax.while_loop covered in this chapter are all built on top of their tracing machinery, so let's review that first.

jax.jit: tracing + XLA compilation. The first time a function is called with a given input shape and dtype, JAX runs the Python code once with abstract tracers in place of real arrays, records the primitive operations it hits as an intermediate representation called a jaxpr, and hands that to XLA to compile into machine code for CPU, GPU or TPU. After that, as long as the shapes and dtypes match, the cached compiled program runs directly and the Python function body does not run again. That means:

• The function must be a pure function. Side effects such as print or modifying globals happen only once, during tracing;
• A change in input shape or dtype, or in the value of a static argument (static_argnums), triggers retracing and recompilation.

⚠️ Note: It's easy to assume that changing a Python constant captured in a closure triggers recompilation. In fact, values captured by a closure or read as globals are not part of jit's cache key. When they change, jit does not recompile; it keeps using the old value from tracing time. This is the stale-value problem. Only creating a new function object (for example a new closure) causes retracing. Section 1.7 proves this with asserts.

jax.grad: reverse-mode automatic differentiation. The function being differentiated must return a scalar. grad returns a new function that computes the gradient, by default with respect to the first argument; use argnums to pick other arguments. Internally it is built on VJPs (vector-Jacobian products), and nesting it gives higher-order derivatives. If the output isn't a scalar, reduce it to a scalar first, or use jax.jacrev / jax.jacfwd; you can also combine with vmap to compute Jacobians or Hessians.

jax.vmap: automatic vectorization. It turns a function that handles one example into a function that handles a batch of examples, without writing a Python loop. The signature is jax.vmap(fun, in_axes=0, out_axes=0). An integer in in_axes means "map along this axis"; None means this argument isn't mapped and the whole batch shares one copy (like broadcasting). out_axes says where the mapped axis goes in the output.

⚠️ Note: The "vectorization" in vmap is vectorization. Some materials mistranslate it as quantization, which is a completely different thing.

Composition. The three transformations nest freely, e.g. jit(grad(f)), vmap(jit(f)), jit(vmap(grad(f))). The nesting order determines the meaning. For example, vmap(grad(f)) gives per-example gradients, while grad(lambda p: mean(vmap(f)(p))) gives the gradient of the mean loss.

Course code 1: speeding up with jit¶

Random keys: old style vs recommended¶

🔍 Extra: choosing data for the jit example. ⚪ Optional: read this when you write your own eager-vs-jit numerical cross-checks.

If you generate the data with the normal distribution random.normal, about 16% of the elements are below −1, and log(x+1) gives NaN there. jnp.allclose defaults to equal_nan=False, so it returns False as soon as either array contains a NaN. Even if you switch to equal_nan=True to set the NaNs aside, there's another problem: at some positions the two terms cancel and the function value is close to 0 (this data has one such spot, x≈3.58). There rtol does almost nothing, and the normal rounding difference between eager and jit can exceed the default atol=1e-8. So when cross-checking, keep the inputs inside the function's domain. That's why the code above uses uniform.

Course code 2: gradient of a linear model with grad¶

The model and loss are

$$\hat y = \mathbf w^\top \mathbf x + b,\qquad L(\mathbf w,b)=(\hat y-y)^2,$$

and the analytic gradients are $\partial L/\partial \mathbf w = 2(\hat y-y)\,\mathbf x$, $\partial L/\partial b = 2(\hat y-y)$.

⚠️ Note: With this data, $\hat y = 1.5\times0.2-0.5\times0.8+0.3 = 0.2$ and the residual is $-2.3$, so $\partial L/\partial\mathbf w=[-0.92,\,-3.68]$ and $\partial L/\partial b=-4.6$. If you see elsewhere a w-gradient of [-0.392 -1.568] and a b-gradient of -1.96, that's a miscalculation. The cell above already verifies the result against the analytic formula with assert.

Course code 3: batched affine transform with vmap¶

Course code 4: per-example gradients with jit(vmap(grad))¶

🔍 Extra: composition order = mathematical meaning. ⚪ Optional: skip this if you've done the "per-example gradients" topic in Basics Chapter 4.

vmap(grad(loss)) gives the gradient for each example (shape (B, ...)), while grad(lambda p: mean(vmap(loss)(p))) gives the gradient of the mean loss. Differentiation is linear, so the latter equals the former averaged over the batch dimension; the cell below verifies this. For the principle and uses (differential privacy, per-example clipping, etc.), see the per-example gradients topic in §4.6 of Basics 04 Automatic vectorization with vmap.

1.2 Mastering lax.scan for sequential operations¶

jit, grad and vmap are great at processing data "as a whole batch". But many algorithms have dependencies between time steps, such as RNNs, running sums, ODE integration and discounted returns in reinforcement learning: step $t$ needs the result of step $t-1$. jax.lax.scan is the primitive designed for exactly this kind of computation.

Why not just write a Python for loop inside jit? During tracing, a Python loop runs as-is. If the trip count is a Python integer (a static value), every iteration's operations get unrolled into the jaxpr: loop 1000 times and the graph contains 1000 copies of the loop body, which means slow compilation and high memory use. If the trip count is a traced value, Python's range() can't run at all and you get an error. lax.scan instead traces the loop body only once and compiles it into a single XLA loop primitive.

⚠️ Note: People often say "if you use a Python for loop inside jit and its length depends on a runtime value, JAX unrolls the loop". In fact it only unrolls when the length is a static Python integer. When the length depends on a runtime (traced) value, range(n) immediately raises TracerIntegerConversionError, so it never gets as far as unrolling. The cell below verifies both cases.

Core idea and signature¶

final_carry, ys = jax.lax.scan(f, init, xs=None, length=None, reverse=False, unroll=1)

• f(carry, x) -> (new_carry, y): the loop body. carry is the state passed from step to step, x is the current slice of xs along axis 0, and y is this step's output;
• init: the initial carry. Its pytree structure, and the shape and dtype of every leaf, must exactly match the new_carry returned by f;
• xs: optional, can be any pytree; all leaves are iterated in lockstep along axis 0. When it's None, the loop body receives x as None;
• length: number of iterations. Usually inferred from xs; required when xs=None. If given, it must equal the size of xs along axis 0;
• reverse: when True, scans from the last element backward; the outputs ys are still stored at their original index positions;
• unroll: a performance tuning knob, the number of copies of the loop body placed in one XLA loop iteration. It can be a positive integer or True (fully unrolled);
• Return values: final_carry and ys. ys stacks each step's y along a new axis 0 and can be a pytree; if y is None, ys is None too.

In plain Python, the semantics are:

def scan(f, init, xs, length=None, reverse=False):
    if xs is None:
        xs = [None] * length
    carry, ys = init, []
    for x in (reversed(xs) if reverse else xs):
        carry, y = f(carry, x)
        ys.append(y)
    if reverse:
        ys = ys[::-1]
    return carry, np.stack(ys)

As a recurrence: $c_t = f_c(c_{t-1}, x_t)$, $y_t = f_y(c_{t-1}, x_t)$, $t=1,\dots,T$.

Course code 1: running sum¶

Course code 2: a simple RNN with scan¶

$$h_t = \tanh\left(W_{hh}\,h_{t-1} + W_{xh}\,x_t + b_h\right)$$

The carry is the hidden state $h$ and xs is the input sequence. Each step returns the new hidden state as both the carry and the output. There are two easy-to-hit pitfalls with keys here, so let's compare first:

⚠️ Note: If you use the same w_key to generate both W_hh and W_xh, W_xh comes out exactly equal to the first 15 elements of flattened W_hh, so the two matrices are fully correlated. Also, x_key = random.split(key) gives a key array, and passing it to random.normal raises ValueError: normal accepts a single key.... The right way is key, x_key = split(key), or one split(key, n) followed by unpacking.

Another common pattern: define the parameters outside rnn_step and let JAX pick them up through the closure. That runs, but the weights get embedded as constants in the compiled program (see Section 1.7 for why). On this machine's jaxlib 0.12.0.1, capturing a large matrix in a closure can also silently give wrong results. So the recommended version passes the parameters in as arguments of the jitted function.

Benefits of scan over a Python loop (course key points):

1. No unrolling: the loop body is traced once and compiled into an XLA While loop, so compile time and program size are basically independent of the number of steps;
2. Memory efficiency: XLA can reuse buffers across iterations;
3. Loop-body optimization: XLA can apply optimizations such as operator fusion inside the loop body.

Relationship with other transformations: it works best inside jit. grad can differentiate with respect to init, xs and parameters captured by the loop body, which amounts to BPTT (backpropagation through time). vmap can process a batch of sequences in parallel, e.g. jax.vmap(lambda seq: lax.scan(step, h0, seq))(batch). If each sequence also has its own initial state, add h0 to in_axes too.

Course code 3: xs=None and length¶

⚠️ Note: Even when xs is None, the loop body can't be written as just f(carry). scan always calls the loop body with two arguments, f(carry, x); with xs=None, x is simply None. Writing only one parameter raises TypeError: <lambda>() takes 1 positional argument but 2 were given.

🔍 Extra: details of reverse, unroll, length and pytrees¶

⚪ Optional: read this when you use reverse, unroll, or pytree-shaped xs / ys.

reverse=True scans from back to front, but ys[t] still corresponds to xs[t]; the output is not flipped. It's equivalent to "flip xs, scan forward, then flip ys back". Typical uses are discounted returns in RL, $G_t = r_t + \gamma G_{t+1}$, the backward branch of a bidirectional RNN, and suffix sums.

unroll doesn't change the result, only the structure of the compiled program. unroll=k means each XLA loop iteration runs k time steps; True means fully unrolled. When the loop body is tiny (scalar ops or small matrix ops), the fixed per-iteration loop overhead is a large fraction, and a moderate unroll (2, 4, 8) may be faster. When the loop body is large, unrolling only bloats compile time and code size. Always measure it yourself. Below is one set of numbers from this machine's CPU (for comparison only, no assertions). If you run it a few times, you'll see that in this example the different unroll settings take about the same time, basically within timing noise, and full unrolling is sometimes even slower. That tells you the bottleneck here isn't loop overhead:

Compile time of scan vs a Python loop (same 500 steps, printed only, no assertions):

A length that disagrees with the length of xs raises an error at trace time. Pytree-shaped xs and ys are very common: you can pack several sequences into a dict or tuple and iterate over them in lockstep.

🔍 Extra: reading carry-type mismatch errors¶

🟢 Must-read

scan requires init and the new_carry returned by the loop body to have exactly the same type: pytree structure, shape and dtype. This is the most common scan error, and JAX 0.12 tells you precisely which part doesn't match:

One easy trap: init=0 (a Python int, weakly typed) with floating-point xs does not raise; JAX promotes the weakly typed carry to float32. But if you write jnp.int32(0) (strongly typed), you get a dtype mismatch. So the safest habit is to spell out the dtype of init, e.g. jnp.zeros((), xs.dtype).

🔍 Extra: lax.associative_scan. ⚪ Optional: read this when you work on long associative recurrences such as linear RNNs or SSMs.

If the recurrence is associative, e.g. sums, products, or the linear recurrence $h_t = a_t h_{t-1} + b_t$, you can use lax.associative_scan for a parallel prefix scan. Its depth is $O(\log T)$, and on GPU/TPU it's usually much faster than a sequential scan for long sequences (this machine is CPU-only, so the speedup wasn't measured). Linear RNNs and SSMs (such as the parallel mode of S5 and Mamba) rely on exactly this. Ordinary nonlinear RNNs aren't associative and generally need scan.

1.3 Conditional execution with lax.cond¶

During jit tracing, a Python if can only choose a branch based on concrete Python-level values. If the condition depends on the runtime value of a JAX array (a traced value), if has to convert the tracer to bool, and tracing fails. lax.cond instead puts the conditional logic into the computation graph:

jax.lax.cond(pred, true_fun, false_fun, *operands)

• pred: a scalar boolean with shape (). It can be a traced value or a Python bool;
• true_fun / false_fun: the two branch functions, both taking *operands;
• *operands: inputs passed to the branches; can be several arrays or any pytree.

Key constraint: the outputs of the two branches must have exactly the same pytree structure, shape and dtype, because the output type of the compiled program has to be fixed at compile time.

Under the hood: during tracing both branches are traced, to check and fix the output type. The compiled program contains a real runtime branch, and when pred has no vmap batch dimension only the selected branch runs (for the batched case, see the Extra in Section 1.3).

Course code: Python if vs lax.cond¶

⚠️ Note: Many materials say you get a ConcretizationTypeError here. What JAX 0.12 actually raises is the more specific TracerBoolConversionError. It's a subclass of ConcretizationTypeError, so except ConcretizationTypeError still catches it.

Multiple operands are simply passed in one after another. Branch outputs with different types and a non-scalar pred both raise errors at trace time:

Relationship with other transformations:

• vmap: judging by the result, each element of the batch "picks its own branch" based on its own condition. But the implementation isn't what you'd intuitively expect; see the Extra below;
• grad: gradients flow only through the branch that actually ran.

⚠️ Note: It's easy to think that when pred depends on a differentiable variable, gradients flow through it too. In fact pred comes from a comparison (>, >=, etc.), and comparisons output booleans, which are not differentiable, so gradients do not flow through pred. The derivative with respect to a threshold parameter that decides the branch is always 0. The real issue is that the function may be discontinuous at the switch point, and the gradient can't see the change caused by "switching branches". If you need a derivative with respect to the threshold, use a smooth relaxation such as sigmoid. Section 1.5 proves with assert that the threshold gradient is 0.

lax.cond vs jnp.where:

🔍 Extra: lax.switch and lax.select_n¶

🟢 Must-read

• lax.switch(index, branches, *operands): the multi-branch version of cond; index is an integer scalar. An out-of-range index is clamped to [0, len(branches)-1] instead of raising. This behavior can easily hide bugs, so add your own check when needed;
• lax.select_n(which, *cases): elementwise multi-way selection, the primitive underneath jnp.where. It doesn't broadcast; all cases must have the same shape. When which is bool, False picks case 0 and True picks case 1, which is the opposite order of jnp.where(c, x, y). When which is an integer, it picks case number which.

🔍 Extra: cond turns into select under vmap¶

🟢 Must-read

vmap treats the whole batch as one array and processes it with a single program, so there's no way for different elements of the batch to take different branches. This is a vmap transformation rule and has nothing to do with the backend; the jaxpr we see below on CPU shows exactly this. So if the vmap batch dimension reaches pred, JAX rewrites cond as: run both branches on the whole batch, then pick elementwise with select_n. At that point cond's advantage of "computing only one branch" is gone, and the cost is the sum of both branches. Be especially careful when the branches have side effects (such as jax.debug.print) or are expensive.

If pred is not batched, for example when only the operand has a batch dimension, cond is kept and still runs only one branch.

🔍 Extra: jnp.where and NaN in gradients. ⚪ Optional: skip this if you already know the "double where" trick from Basics Chapter 3.

jnp.where computes both sides, so a NaN/Inf on the side that isn't selected still pollutes the gradient; for example, the gradient of where(x>0, sqrt(x), 0) at x=0 is NaN. The fix is the "double where": first replace invalid inputs with safe values, then compute. For the principle and a demo, see §3.8 of Basics 03 Automatic differentiation with grad; for a more advanced discussion, see topic 2 in Basics 07 Practical topics.

1.4 Looping with lax.while_loop¶

The number of steps of lax.scan is fixed at compile time. Some iterations need to decide when to stop based on runtime intermediate results, such as Newton's method converging, binary search, or the Collatz sequence. That's what lax.while_loop is for:

final_val = jax.lax.while_loop(cond_fun, body_fun, init_val)

• cond_fun(carry) -> scalar bool: keep looping while it's True;
• body_fun(carry) -> new_carry: the returned state must have exactly the same structure, shape and dtype as the input;
• init_val: the initial state; can be a scalar, an array or a pytree.

The execution order is cond(init) → body(init) → cond(next) → … until cond returns False, and then it returns the state that made cond False. It's equivalent to this Python code:

val = init_val
while cond_fun(val):
    val = body_fun(val)
return val

The state is functional: body doesn't modify its input; it returns a new state. cond_fun and body_fun are each traced only once. The carry's shape and dtype must be static, but its values can be dynamic, and the whole loop compiles into a single XLA While op.

Compared with scan: scan has a fixed number of steps and can collect each step's output ys; while_loop has a variable number of steps and returns only the final state.

Caveats: cond_fun must eventually return False, or the compiled program loops forever and the process hangs. And Ctrl-C usually can't interrupt it either, because control is inside the compiled XLA loop and never returns to Python (we tried sending SIGINT to the process, and it was still running 10 seconds later). The only way out is to restart the kernel or kill the process. In real code, add an iteration counter to the carry and set a maximum number of steps as a safety net.

Course code: the smallest power of 2 that is at least n¶

while_loop can't do reverse-mode differentiation¶

⚠️ Note: You'll often see claims like "while_loop supports jax.grad; backprop unrolls the iterations that ran", but that's not true. Reverse mode (grad / vjp) on lax.while_loop raises an error right away: ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. The reason is that backprop needs to save each step's intermediate values (residuals), and since the number of while_loop steps isn't known at compile time, XLA can't preallocate buffers. Forward mode (jvp / jacfwd) doesn't need to keep history and works fine.

What to use when the number of steps is fixed?

You can see fori_loop's "dual identity" directly in the jaxpr:

🔍 Extra: one detail. ⚪ Optional: useful when you write library code and need fori_loop to be reverse-mode differentiable.

Calling jax.grad(fl_dynamic)(2.0, 5) outside jit does succeed, because there n=5 is a concrete Python integer and fori_loop treats it as static. The same code fails as soon as you add jit. So in library code, don't rely on behavior that depends on a value "happening to be concrete": either declare the number of steps in static_argnums, or just use scan.

🔍 Extra: under vmap, while_loop runs until "the slowest element" finishes¶

⚪ Optional: read this when you run a while_loop over a batch and the elements need very different numbers of iterations.

If the vmap batch dimension reaches cond_fun, different elements may need different numbers of iterations. JAX's approach: as long as any element still needs to continue, the whole batch keeps looping. In the lowered HLO, this shows up as a reduce or over the predicate. Finished elements have their state frozen with select; they still take part in the computation, but the results are thrown away. So when iteration counts vary a lot within a batch, a lot of compute is wasted, and you may want to bucket elements by their expected number of steps.

1.5 Combining control flow with transformations¶

What makes JAX powerful is that the lax control-flow primitives compose freely with jit / grad / vmap. This section goes through the behavior and cost of each combination.

With jit¶

jit traces the scan / while loop body, and each branch of a cond, into jaxprs, which XLA then compiles into HLO. The loop body is compiled only once; both branches of a cond are compiled, and at runtime one of them runs based on the predicate.

With grad¶

scan: backprop applies the chain rule in reverse time order, i.e. BPTT. The intermediate values of every forward step must be saved for the backward pass, so memory grows linearly with sequence length. For long sequences you can ease this with gradient checkpointing (jax.checkpoint); see the Extra later in this section.

Course example: a scalar RNN, $h_t = \tanh(W x_t + U h_{t-1} + b)$ with $W=0.5,\ U=0.8,\ b=0.1$, loss $L=\sum_t h_t$, differentiated with respect to the initial hidden state $h_0$. Deriving by hand gives:

$$\frac{\partial L}{\partial h_0}=\sum_{t=1}^{T}\prod_{k=1}^{t} U\,(1-h_k^2).$$

⚠️ Note: Here dL/dh0 ≈ 1.0749918, in full agreement with the hand-derived formula and the float64 central difference. Some materials give about 0.313..., which is wrong.

cond: gradients pass only through the branch that actually ran. In the course example, when sum(x)>0 the loss is $\sum p_i x_i$ and the gradient is $x$; otherwise the loss is $\sum p_i/(x_i+10^{-5})$ and the gradient is $1/(x+10^{-5})$.

⚠️ Note: The negative-case gradient is $1/(-1+10^{-5}) = -1/0.99999 \approx$ -1.00001, whose absolute value is greater than 1. Don't write it as [-0.99999 -0.99999] by mistake.

while_loop: no reverse mode; see Section 1.4.

With vmap¶

scan: each sequence in the batch is scanned independently. vmap adds a batch dimension to both carry and xs; there's still just one loop, and each step processes the whole batch.

cond: when the predicate depends on the mapped axis, JAX computes both branches over the whole batch and then picks results by each element's predicate (that's the select_n from the Extra in Section 1.3). The cost is the sum of both branches.

while_loop: when different elements need different numbers of iterations, the whole batch keeps running until the largest iteration count in the batch, with finished elements frozen by a mask; see the Collatz example in Section 1.4.

Combining several transformations¶

You can write multi-level combinations such as jax.grad(jax.vmap(...)) or jax.jit(jax.grad(conditional_loss)). The example below has two things to watch out for, though:

⚠️ Note: Writing jax.grad(jax.vmap(threshold_from_params, ...)) directly doesn't run; it raises TypeError: Gradient only defined for scalar-output functions. Output had shape: (2, 3). If you sum first and then take grad, you get a result, but the gradient with respect to params is always 0. That's because the threshold affects branch selection only through the >= comparison, and comparisons aren't differentiable. This confirms the conclusion from Section 1.3.

🔍 Extra: jax.checkpoint with scan: trading compute for memory¶

⚪ Optional: read this when backprop over long sequences runs out of memory; the full treatment is in Advanced Chapter 6 §6.6.

When computing reverse-mode gradients through scan, JAX by default saves all intermediate values (residuals) of every step's loop body. If you write lax.scan(jax.checkpoint(step), init, xs), the forward pass saves only each step's carry, and the backward pass recomputes that step's intermediates. The gradients are the same; the price is roughly one extra forward pass (the LSTM in hands-on A of Section 1.8 verifies with assert that gradients match with and without checkpointing). For inspecting residuals with print_saved_residuals, policy options, measured memory comparisons, and the standard "scan over stacked layers + checkpointing" structure, see Advanced Chapter 6 §6.6.

1.6 Advanced masking techniques¶

XLA requires static shapes, but real data often comes in different lengths: sentences of different lengths, episodes that end early, graphs of different sizes. The standard approach is to pad to a common length and pair it with a boolean mask that marks which positions are valid. Mask operations are elementwise with no branches, so they compile efficiently under jit, and vmap and grad handle them directly.

Why do we need masks?

• When computing means or attention weights, padding positions must be excluded, or the result gets "diluted" by the padding values;
• In scan, once a shorter sequence has ended, its hidden state and outputs shouldn't be updated anymore.

The core API is jnp.where(condition, x, y): take x where condition is True, otherwise y; all three can broadcast.

Course code: padding and sequence masks¶

Pad [1,2,3], [4,5], [6,7,8,9] with 0 to length 4:

The masked mean is $\bar x_i = \dfrac{\sum_j m_{ij}x_{ij}}{\max(1,\sum_j m_{ij})}$. The $\max(1,\cdot)$ in the denominator prevents division by zero for completely empty rows.

🔍 Extra: don't build masks from "value != 0". 🟢 Must-read

Building the mask with padded != 0 above is just for the demo. In real data, 0 can be a legitimate value, such as token id 0 or a feature that happens to be 0. In real projects, build the mask from a lengths array: mask = jnp.arange(T)[None, :] < lengths[:, None].

Course code: masking attention scores¶

Set the scores at padding positions to a very large negative number, and after softmax their weights will be close to 0.

🔍 Extra: when "the whole row is masked", -1e9 and -inf behave differently. 🟢 Must-read

With a causal mask on top of padding, or an empty sequence in the batch, an entire row can be False:

• Fill with -1e9: softmax gives a uniform distribution [0.25, 0.25, ...]. No error, but these weights are meaningless, and you usually need to multiply by the mask again afterward;
• Fill with -inf: mathematically this is 0/0. jax.nn.softmax first subtracts the row maximum, giving $-\infty-(-\infty)$, so the result is also NaN, and the NaN keeps propagating into the gradients;
• Use jax.nn.softmax(..., where=mask): fully masked rows output all zeros, with no NaN (verified in the cell below).

So in practice people use a finite large negative number, such as jnp.finfo(dtype).min or -1e9. At low precision, check that the value itself doesn't overflow: for example, -1e9 cast to float16 becomes -inf (float16 tops out around 65504), while bfloat16 has the same exponent range as float32 and doesn't overflow.

Course code: arithmetic masks¶

Convert the mask to floats and multiply it with the data, and you get masked sums and means as well. The course specifically warns that multiplying by 0 isn't always safe. If a masked position holds inf or NaN, $0\times\infty=\text{NaN}$ and the result is polluted anyway. In log space, log(0) causes trouble too. So prefer jnp.where.

Masks with control flow and transformations¶

• jit: mask operations are ordinary elementwise ops, so compilation has no trouble at all;
• vmap: mask operations are vectorized automatically;
• grad: gradients flow only through the side that jnp.where selects; the unselected side gets gradient 0. But watch out for NaN; see the "double where" in §3.8 of Basics 03 Automatic differentiation with grad;
• scan: at each time step, the mask decides whether to "update the state" or "keep the old state".

There's one easy mistake when using masks in scan: the mask changes with the time step, so it should go into xs together with the inputs and be read step by step, not into the carry, because the carry is just passed along unchanged from step to step. Below is the full version, which also verifies that "pad, then scan with a mask" gives exactly the same final state as "scan only over the real length".

⚠️ Note: A common broken version puts current_mask in the carry, says xs=(inputs, masks) in a comment, and uses x as a single input in the loop body. Those three things contradict each other, and the code doesn't run. The right approach is to read the mask step by step as part of xs, i.e. xs=(inputs, masks), and unpack it in the loop body as (input_element, current_mask). Written that way, the result has been cross-checked against a per-sequence Python loop.

Performance trade-offs (course key points): masking computes every position and then picks results. On GPU/TPU this is usually faster than conditional branching because it doesn't cause thread divergence. The cost is that the excluded elements are computed for nothing. If one branch is expensive and often skipped, and the predicate is a scalar (not on a vmap batch dimension), lax.cond may be faster. Let profiling decide.

🔍 Extra: segment operations with jax.ops.segment_sum¶

⚪ Optional: read this for graph neural networks, variable-length data with huge length differences, or before Exercise 3.

The compute wasted by padding + mask is proportional to "max length − mean length". If lengths vary hugely, such as node degrees in a graph neural network or the number of user actions in a recommender system, you can take a different approach and not pad at all: concatenate all elements into one long array, tag each element with the "segment id" it belongs to, and reduce per segment:

$$\text{out}[k] = \sum_{i:\ \text{seg}[i]=k} \text{data}[i].$$

jax.ops provides four functions, segment_sum / segment_max / segment_min / segment_prod. Keep these points in mind when using them:

1. num_segments must be static under jit. If you don't pass it, the default is max(segment_ids)+1, which needs a concrete value, so under jit it raises ConcretizationTypeError;
2. ids in segment_ids that are ≥ num_segments are silently dropped, and so are negative ids;
3. An empty segment gives the identity of the reduction: 0 for sum, 1 for prod, -inf for max, +inf for min (for floats; with integer dtypes, the empty-segment result of max / min is that type's minimum / maximum value);
4. If you know the ids are sorted, pass indices_are_sorted=True; some backends are faster that way.