At https://github.com/google/flax/blob/main/flax/nnx/training/optimizer.py#L221 it would be nice if nnx.Optimizer.update could return the underlying updates PyTree from optax. All existing callers are free to ignore it, but this would allow logging and diagnostics on the weight update in a straightforward way without the caller needing to drop down to pure optax.