Update Jax to release version

Author: andersensam

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "axlearn"
-version = "0.1.8"
+version = "0.1.9"
 description = "AXLearn"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -23,8 +23,8 @@ core = [
     "absl-py==2.3.1",
     "chex==0.1.91",
     "importlab==0.8.1",  # breaks pytype on 0.8
-    "jax==0.8.0.dev20251006",
-    "jaxlib==0.8.0.dev20251006",
+    "jax==0.8.0",
+    "jaxlib==0.8.0",
     "ml-dtypes==0.5.3",
     "msgpack==1.1.0",  # for checkpointing.
     "nltk==3.7",  # for text preprocessing
@@ -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.8.0.dev20251006",  # must be >=0.4.19 for compat with v5p.
+    "jax[tpu]==0.8.0",  # must be >=0.4.19 for compat with v5p.
     "pathwaysutils==0.1.1",  # 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.8.0.dev20251006",
+    "jax[cuda12]==0.8.0",
     "nvidia-ml-py==12.560.30",
     # pin nccl version, otherwise jax[cuda12] will pull latest version
     "nvidia-nccl-cu12==2.27.5",
@@ -210,14 +210,14 @@ junit_family="xunit2"
 line_length = 100
 profile = "black"
 
-[tool.pip.uv]
+[tool.uv.pip]
 find-links = [
     "https://storage.googleapis.com/axlearn-wheels/wheels.html",
     "https://storage.googleapis.com/jax-releases/libtpu_releases.html",
 ]
-extra-index-url = [
-    "https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/simple"
-]
+# extra-index-url = [
+#     "https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/simple"
+# ]
 no-build-isolation-package = [
     "apache-beam",
 ]