Scan, while loops and structured state

A running total needs yesterday’s state

The increments are [1, 2, 3, 4]. Their running totals are [1, 3, 6, 10]. Each output depends on the previous carry, so independently mapping a function across the increments would miss the dependency. A scan expresses the sequence of state updates and collects an output from every step.

A while loop answers another iteration question: double a value until it reaches or exceeds 10. Starting at 2 gives 2 → 4 → 8 → 16. Its number of steps depends on runtime state. Both APIs keep the carry's shape fixed; the choice is about how iteration is determined, not permission to grow a tensor.

Choose Run all, then change the sequence or initial state. This notebook uses jax4j 0.2.1 on the emulated CPU. Structured regions are interpreted by that backend; these are shape and derivative examples rather than device performance measurements. Reloading discards declarations, edits and Java state.

1. Describe one scan step and one while loop

A scan step returns both its next carry and its per-step output. Those can be different tensors. Here they are the same expression: the updated total. The scan API returns separate final-carry and stacked-output tensors afterward.

The while predicate is positive when its scalar state is below 10. Taking sign().maximum(0) turns that condition into the zero/nonzero convention from the conditional lesson. The body doubles the state without changing its shape.

These typed declarations persist for the following cells. The record will organize a second scan's state. A Java int counter in that record would not be a runtime tensor leaf; we use a one-element count tensor so the update belongs to the structured calculation.

The callbacks build supported tensor operations. Keep host reads and mutable Java effects outside them. Their internal temporaries are managed through the region capture or enclosing transformation context; ordinary inputs and returned outputs still need explicit resource owners.

2. Scan over a leading sequence axis

The input has shape [4, 1]: four steps, each receiving a one-element sample. The carry has shape [1] throughout. The per-step output also has shape [1], so collecting all four outputs produces shape [4, 1].

Expect final carry 10 and outputs [1, 3, 6, 10] with shape [4, 1]. Change the initial carry to 10. The outputs become [11, 13, 16, 20] and the final carry becomes 20. The sequence is unchanged; every step inherits the different starting state.

Try returning next.mul(next) as the step output while keeping next as the carry. With the original zero start, the collected outputs become [1, 9, 36, 100], while the final carry remains 10. This separates state needed by the next step from observations you want to retain.

Scan captures its body once and represents the fixed sequence length as a region loop. You do not write a Java loop that expands one copy of the body per step. This release can expand static scans for higher-order derivative construction, so that representation statement should not be generalized to every nested derivative at every order.

3. Carry a named total and count

A running mean needs two pieces of state: total and count. scanTree lets us keep those in a record while emitting the mean after each increment. The carry's tree structure and leaf shapes remain fixed across steps.

Expect total 10, count 4 and means [1, 1.5, 2, 2.5]. Change the second increment from 2 to 6: the new means are [1, 3.5, 10 / 3, 3.5]. Predict the total and count before running it. The count advances through tensor arithmetic, not a Java side effect.

The result record contains a new carry tree plus collected outputs. It does not close its tensor leaves automatically. Whole-tree cleanup is appropriate here because these returned leaves are newly computed; the input tensors remain owned by the outer resource list. The same ownership distinction from Chapter 3 still applies when a different function returns borrowed inputs.

4. Push a tangent through a dynamic while loop

A while loop repeatedly evaluates its predicate on the current carry. Starting at 2, doubling takes three steps and finishes at 16. Forward differentiation can propagate a tangent through those executed updates: a tangent of 1 becomes 8 after three doublings.

Expect final value 16 and tangent 8. Try starting at 3: the loop runs twice, finishes at 12 and has tangent 4. These derivatives describe the path where the number of iterations is locally stable. At inputs where that count changes, the complete function can jump; a path derivative is not a derivative of a discrete iteration-count decision.

The carry must keep its shape. Use structured state when you need both a value and an explicit iteration counter, and design a terminating predicate. The public whileLoop call has no separate maximum-iteration argument. For a known number of sequence steps, scan makes that bound explicit in its input.

5. Reverse mode for dynamic while is rejected

jax4j 0.2.1 supports forward-mode differentiation of these dynamic while regions, but rejects reverse mode. The next cell intentionally asks grad for that unsupported transformation and catches the exception. A finite answer to the forward example does not imply that reverse mode is available.

Do not replace a genuinely dynamic loop with an arbitrary fixed iteration count just to silence this error. If a fixed horizon is part of the model, express it as a scan and differentiate that computation. If you only need a few input directions through a dynamic process, a JVP may answer your actual question directly.

6. Reverse-differentiate a fixed scan

For the additive scan, the final total is the initial value plus every input increment. Its gradient with respect to the four increments is therefore [1, 1, 1, 1]. This fixed sequence supports the reverse-mode question that was unavailable for the dynamic while loop.

Expect shape [4, 1] and four ones. Try differentiating the sum of all collected running totals instead, using .outputs().sum() in the lambda. The answer becomes [4, 3, 2, 1]: the first increment influences four reported totals, while the last influences only one. The choice of output objective determines which paths contribute to the gradient.

Continue with custom JVPs, VJPs and checkpoints when an operation needs an explicit derivative contract or recomputation policy. Revisit structured parameters for the distinction between container structure and runtime tensor leaves.