flax

6a75a1d3 - Merge pull request #4830 from google:fix-example-gemma2