From 3841c87d61764d57c34387bc28c2c63b5da74697 Mon Sep 17 00:00:00 2001
From: Branden Vandermoon <103541588+bvandermoon@users.noreply.github.com>
Date: Fri, 20 Mar 2026 16:03:02 -0700
Subject: [PATCH] Bump Flax version for latest JAX compatibility

---
 .../requirements/generated_requirements/tpu-requirements.txt    | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/dependencies/requirements/generated_requirements/tpu-requirements.txt b/src/dependencies/requirements/generated_requirements/tpu-requirements.txt
index e1c715d3ba..08da4a3ab7 100644
--- a/src/dependencies/requirements/generated_requirements/tpu-requirements.txt
+++ b/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