flax

a7157cfb - Merge pull request #4715 from google:nnx-mutable-array-p1