flax

Copies `dynamic_scale.py` from `optim/` to `training/`.

#2375

Merged