Loose up the JAX version to run JAX nightly TPU jobs

Author: judyzhaoxinwu

pyproject.toml

@@ -107,7 +107,7 @@ gcp = [
 # Note: Specify -f https://storage.googleapis.com/jax-releases/libtpu_releases.html during install.
 tpu = [
     "axlearn[gcp]",
-    "jax[tpu]>=0.9.1.dev20260130,<0.9.2",  # must be >=0.4.19 for compat with v5p.
+    "jax[tpu]>=0.9.0,<0.9.2",  # must be >=0.4.19 for compat with v5p.
     "pathwaysutils==0.1.2",  # For JAX+Pathways single-controller accelerator coordinator.
 ]
 # Vertex AI tensorboard. TODO(markblee): Merge with `gcp`.
@@ -133,7 +133,7 @@ dataflow = [
 # GPU custom kernel dependency.
 gpu = [
     "triton>=2.1.0",
-    "jax[cuda12]>=0.9.1.dev20260130,<0.9.2",
+    "jax[cuda12]>=0.9.0,<0.9.2",
     "nvidia-ml-py==12.560.30",
     # pin nccl version, otherwise jax[cuda12] will pull latest version
     "nvidia-nccl-cu12==2.27.5",