flax

4b4b19f0 - Merge pull request #4322 from google:nnx-cleanup-nnx-basics