flax

Explicitly mark the lazy-rng with legacy_config, and all new configs should use flax.config.

#2895

Merged