flax

c0a2657a - Merge pull request #3518 from google:nnx-fix-graph-flatten