flax

abc11552 - Merge pull request #4384 from google:nnx-performance-guide