Use of Optax "Reduce on plateau" on Flax NNX

[JavierJJR]

Hello all,

I am quite new to Optax and Flax. To learn how to use Optax, I have begun implementing models that I had previously built in PyTorch into Flax NNX. However, I'm encountering significant difficulties when implementing the Reduce-On-Plateau Learning Rate Scheduler. At https://optax.readthedocs.io/en/latest/_collections/examples/contrib/reduce_on_plateau.html, an example is provided for using it with Flax Linen, but I haven't found any examples for Flax NNX. For this reason, I was wondering if anyone has a simple example of how to use this scheduler in Flax NNX.

Thanks a lot in advance! 🙏