Merge pull request #464 from google:jax_nightly

PiperOrigin-RevId: 611623730
google · 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