flax

a47c375d - Merge pull request #626 from google:nlp_seq_linen