flax

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