flax

9b5cd533 - per-example based update rule for stochastic mirror descent.

per-example based update rule for stochastic mirror descent. Also added a geometric lr schedule, and added num_cycles to the cosine_lr_schedule function (just for completeness). PiperOrigin-RevId: 381300586