Gradient Accumulation via MultiSteps

Collect two supplied gradient arrays, average them, and apply one SGD update. Gradient accumulation can let a training loop process smaller batches before updating parameters. This task starts after those gradients have been computed.

Two calls, one parameter change

Wrap optax.sgd(lr) with optax.MultiSteps(..., every_k_schedule=2). Keep its default use_grad_mean=True.

1. Initialize the wrapper's state from params.
2. Call update with grads_a. Save the returned state. The returned updates are zero, so applying them leaves the parameters unchanged.
3. Call update with grads_b and that saved state. This time the wrapper sends the mean gradient to SGD and returns the parameter updates.
4. Apply those updates and return the resulting parameters.

The resulting rule is:

mean_grad = (grads_a + grads_b) / 2
new_params = params - lr * mean_grad

For params = [1, 2], grads_a = [1, -1], grads_b = [3, 1] and lr = 0.1, the mean is [2, 0], giving [0.8, 2.0]. Opposing gradients can cancel. Using their sum instead would double the update.

Keep the accumulation state

Do not initialize again between the two calls. The returned state contains the first gradient contribution and the count needed to trigger the second call's update. After a completed pair, MultiSteps clears the accumulation for the next pair.

This task gives both gradient arrays equal weight. If they represent mean loss gradients from differently sized batches, their simple average is not generally the mean gradient over all examples. No batch sizes or reweighting are part of this exercise.