Bump Flax version for latest JAX compatibility#3476

copybara-service[bot] merged 1 commit intomainAI-Hypercomputer/maxtext:mainfrom

bvandermoon-bump-flaxAI-Hypercomputer/maxtext:bvandermoon-bump-flaxCopy head branch name to clipboard