4-Step Training Loop with Scan + Loss Curve

Train one scalar weight on four batches and record the loss before each update. The model predicts w * x, with no bias term. Each batch uses the sum of squared errors:

loss = sum((w * x - y)**2)

Use a sum, not a mean. Changing the reduction changes both the reported loss and the gradient used by SGD.

One batch, one update

At each of the four steps:

1. Take the next row of x_batches and y_batches.
2. Calculate and record its loss using the current weight.
3. Differentiate that loss with respect to the weight.
4. Apply one SGD update and carry the new weight to the next step.

For w0 = 0, x = [1, 1], y = [1, 1] and lr = 0.1, the first loss is 2. The gradient is -4, so the weight becomes 0.4. If the second batch is identical, its pre-update loss is 2 * (0.4 - 1)**2 = 0.72. With that same batch four times, return approximately [2.0, 0.72, 0.2592, 0.093312].

Different rows may contain different examples, so the recorded losses need not decrease. You are returning four batch losses, not evaluating a fixed validation set after every update.

Express the loop with scan

Use jax.lax.scan with (w, opt_state) as the carry and (x_batches, y_batches) as the inputs. Scan slices both input arrays along axis zero, passing one pair of rows to the body at a time.

The body returns ((new_w, new_opt_state), loss_before_update). Scan passes the new carry onward and stacks the scalar losses into one array. Carry structure, shapes and dtypes must remain consistent between iterations. The values are passed explicitly rather than mutated in place.

Initialize optax.sgd(lr) once before the scan. Plain SGD with a fixed rate has no gradient history, but carrying its returned state follows the same pattern needed for stateful optimizers. Scan represents the loop as a JAX operation instead of unrolling a Python loop under JIT; it does not guarantee that all operations fuse or that the loop runs faster.