flax

8447c032 - Merge pull request #4820 from google:nnx-fix-rng-dtype-check