flax

89c08028 - Merge pull request #3964 from google:nnx-generalize-state