flax

1367acb1 - Merge pull request #4049 from google:nnx-fix-variable-inplace-operators