flax

d6fb0129 - Merge pull request #2540 from cgarciae:parametrize-dropout-rng-collection