flax

a307d473 - Flax: avoid key reuse in tests

Flax: avoid key reuse in tests Detected via running tests with configuration `jax_enable_key_reuse_checks=true`, available since https://github.com/google/jax/pull/19795 PiperOrigin-RevId: 612604299