Support gradient accumulation using a forward fn decorator with lax.scan

[cemkoc]

This PR adds a new forward function decorator into the learner config which allows users to define their own decorators that could transform the forward function behaviour. One such example of this use case is gradient accumulation which is what this PR is about.

Specifically, this PR enables gradient accumulation using a a new forward function decorator called with_minibatch_steps which internally implements a lax.scan loop over minibatches sliced with dynamic input batch slicing. The new forward function decorator, called with_minibatch_steps, wraps the forward function and accumulates gradients using a jax.lax.scan for the number of gradient accumulation steps which the user specifies. To use this decorator the learner config is extended with an optional forward_fn_decorator parameter which will be specified by the user in the following way if they want to enable gradient accumulation:

Example:

gradient_accumulation_steps = 4

learner.forward_fn_decorator = config.config_for_function(with_minibatch_steps).set(
    steps=gradient_accumulation_steps,
    metric_accumulator=MetricAccumulator.default_config(),
)

Since we scan over the forward function (instead of the value_and_grad func) we need to have a way to compute the gradients in a minibatched manner therefore we wrap the forward function with a custom_vjp implementation to compute the gradients and accumulate during the forward phase. Since the gradients are accumulated and computed during the forward phase we simply pass them to the backward phase and use them as is. This allows us to compute and accumulate the gradients in a memory efficient (minibatched) manner.

A similar effort for implementing grad accumulation using a scan was mentioned by @apoorvtintin.

[markblee]

Sorry for missing this earlier, but we should document the other raise conditions too.

[cemkoc]

sounds good, let me update it now

[cemkoc]

done.