flax

692a62c5 - Merge pull request #648 from google:optimized_lstm