flax

a9ef59b2 - Merge pull request #4641 from google:nnx-improve-variable