flax

1b72435a - Merge pull request #4185 from google:nnx-transform-metadata-issue