flax

f7c0669d - Merge pull request #3859 from google:nnx-improve-nnx-basics