Learning rate scheduling

[kazewong]

@marylou-gabrie mentioned it is unlikely a uniform learning rate is going to work out of the box, so we should look into learning rate scheduling at some point

https://flax.readthedocs.io/en/latest/guides/lr_schedule.html