flax

b234192c - Merge pull request #3213 from google:fix-rnn-issue