Merge pull request #464 from google:jax_nightly

PiperOrigin-RevId: 611623730
AI-Hypercomputer · Feb 29, 2024 · 10a7c47 · 10a7c47
2 parents 2e7e967 + a7f807d
commit 10a7c47
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/setup.sh b/setup.sh
@@ -148,8 +148,8 @@ elif [[ $MODE == "nightly" ]]; then
         exit 1
     fi
     echo "Installing jax-head, jaxlib-nightly"
-    # Install jax from GitHub head
-    pip3 install git+https://github.com/google/jax
+    # Install jax-nightly
+    pip3 install --pre -U jax -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
     # Install jaxlib-nightly
     pip3 install --pre -U jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html