flax

85eb8c05 - Merge pull request #3591 from shoyer:patch-1