flax

b87f5b32 - Merge pull request #3114 from google:explicit-dropout-rng