flax

2d15e85f - Merge pull request #3578 from google:nnx-fix-variable-unflatten