flax

2bdcee76 - Merge pull request #3921 from google:nnx-compact