flax

a2af2ddc - Fix newlines

Fix newlines FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/google/flax/pull/1601 from google:marcvanzee-patch-10 3ea48ddf8e76ba67f79423c10c6b6d6291788cdb PiperOrigin-RevId: 400933467