flax

eb1c5ee2 - Merge pull request #3186 from google:fix-rnn-docs