flax

a3f5be21 - Merge pull request #4790 from google:enable-docs-nnx-remat