flax

2e779f57 - Merge pull request #4800 from google:nnx-remove-variable-state