Merge pull request #22448 from hawkinsp:tpuci

PiperOrigin-RevId: 652520107
jax-ml · Jul 15, 2024 · cab1f85 · cab1f85
2 parents 7255ab1 + 3019467
commit cab1f85
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 7 deletions.
diff --git a/.github/workflows/cloud-tpu-ci-nightly.yml b/.github/workflows/cloud-tpu-ci-nightly.yml
@@ -56,17 +56,13 @@ jobs:
               -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
 
           elif [ "${{ matrix.jaxlib-version }}" == "nightly" ]; then
-            pip install .
-            pip install --pre jaxlib \
-              -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
+            pip install --pre . -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
             pip install --pre libtpu-nightly \
               -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
             pip install requests
 
           elif [ "${{ matrix.jaxlib-version }}" == "nightly+oldest_supported_libtpu" ]; then
-            pip install .
-            pip install --pre jaxlib \
-              -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
+            pip install --pre . -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
             pip install --pre libtpu-nightly==0.1.dev${{ env.LIBTPU_OLDEST_VERSION_DATE }} \
               -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
             pip install requests

diff --git a/.github/workflows/metal_plugin_ci.yml b/.github/workflows/metal_plugin_ci.yml
@@ -35,7 +35,7 @@ jobs:
           pip install absl-py pytest
           if [[ "${{ matrix.jaxlib-version }}" == "nightly" ]]; then
             pip install --pre jaxlib \
-              -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
+              -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
           fi;
           cd jax
           pip install .