flax

6382a2bf - Merge pull request #4377 from google:nnx-pytree-optimization