flax

e4418e27 - Merge pull request #4486 from google:fix-to-nnx-linen-attributes