flax

7ce981bc - Update minimum python version info.

Update minimum python version info. Flax follows JAX (and numpy) minimum python version schedule. We weren't specifying the minimum version metadata in pyproject.toml Also includes some random autoformatting fixes.