flax

0ab03658 - Merge pull request #3791 from google:nnx-arrays-are-state