Loose up the JAX version to run JAX nightly TPU jobs

Author: judyzhaoxinwu

pyproject.toml

@@ -24,8 +24,8 @@ core = [
     "chex==0.1.91",
     "etils[epy]==1.13.0",  # pin to ensure importlib-resources is installed (needed by tensorflow-datasets)
     "importlab==0.8.1",  # breaks pytype on 0.8
-    "jax>=0.9.1.dev20260130,<0.9.2",
-    "jaxlib>=0.9.1.dev20260130,<0.9.2",
+    "jax>=0.9.0,<0.9.2",
+    "jaxlib>=0.9.0,<0.9.2",
     "ml-dtypes==0.5.3",
     "msgpack==1.1.0",  # for checkpointing.
     "nltk==3.7",  # for text preprocessing