Fix learning rate schedule uses unscaled learning rate instead of batch size adjusted one.
Fix learning rate schedule uses unscaled learning rate instead of batch size adjusted one.

PiperOrigin-RevId: 356209733

google/flax,google/flax