flax

5408c40c - Merge pull request #4407 from google:rnn-broadcast-rngs