flax

f51707df - Merge pull request #609 from google:nlp_seq_linen