flax version>=0.11.0 requires python>=3.11

susanbao · susanbao · commit 092d7641269e · 2025-09-10T22:27:18.000Z
diff --git a/README.md b/README.md
@@ -71,7 +71,7 @@ MaxDiffusion supports
 
 We recommend starting with a single TPU host and then moving to multihost.
 
-Minimum requirements: Ubuntu Version 22.04, Python 3.10 and Tensorflow >= 2.12.0.
+Minimum requirements: Ubuntu Version 22.04, Python 3.11 and Tensorflow >= 2.12.0.
 
 ## Getting Started:
 
diff --git a/docs/getting_started/first_run.md b/docs/getting_started/first_run.md
@@ -10,7 +10,7 @@ multiple hosts.
 
 1. [Create and SSH to a single-host TPU (v6-8). ](https://cloud.google.com/tpu/docs/users-guide-tpu-vm#creating_a_cloud_tpu_vm_with_gcloud)
 * You can find here [here](https://cloud.google.com/tpu/docs/regions-zones) the list of zones that support the v6(Trillium) TPUs
-* We recommend using the base VM image "v2-alpha-tpuv6e", which meets the version requirements: Ubuntu Version 22.04, Python 3.10 and Tensorflow >= 2.12.0
+* We recommend using the base VM image "v2-alpha-tpuv6e", which meets the version requirements: Ubuntu Version 22.04, Python 3.11 and Tensorflow >= 2.12.0
    
 1. Clone MaxDiffusion in your TPU VM.
 ```bash
diff --git a/requirements.txt b/requirements.txt
@@ -5,7 +5,7 @@ grain
 google-cloud-storage>=2.17.0
 absl-py
 datasets
-flax>=0.10.2
+flax>=0.11.0
 optax>=0.2.3
 torch>=2.6.0
 torchvision>=0.20.1
