Support Optax extra arguments in StatefulTrainer

[sylvesterkaczmarek]

Summary

Fixes #126.

StatefulTrainer currently calls optimizer_def.update(grads, state, params) without any way to pass keyword-only extra arguments, which prevents using Optax GradientTransformationExtraArgs transformations that require additional signals.

This change:

• wraps standard Optax transformations with optax.with_extra_args_support() so existing optimizers keep their current behavior;
• adds a keyword-only optimizer_extra_args mapping to StatefulTrainer.step();
• forwards that mapping only to optimizer_def.update(), while the existing **kwargs continue to go only to the loss function.

Testing

Added a regression test using a GradientTransformationExtraArgs update that requires a scale keyword argument and verifies a training step succeeds when the extra argument is supplied.

The existing deterministic MLP training test continues to exercise a normal optax.adam optimizer through the default path.