Bump Flax version for latest JAX compatibility

Author: bvandermoon

src/dependencies/requirements/generated_requirements/tpu-requirements.txt

@@ -51,7 +51,7 @@ execnet>=2.1.2
 fastapi>=0.122.0
 filelock>=3.20.0
 flatbuffers>=25.9.23
-flax>=0.12.1
+flax>=0.12.6
 fonttools>=4.60.1
 frozenlist>=1.8.0
 fsspec>=2025.10.0