flax

868d3595 - Copies `dynamic_scale.py` from `optim/` to `training/`.

Copies `dynamic_scale.py` from `optim/` to `training/`. This is a preparation for removing `dynamic_scale.py` from `optim/` before removing the entire `optim/` package. PiperOrigin-RevId: 465563581