bump jax everywhere

stanford-crfm · Dec 24, 2024 · b7e6314 · b7e6314
1 parent e9f38cc
commit b7e6314
Show file tree

Hide file tree

Showing 7 changed files with 7 additions and 7 deletions.
diff --git a/.github/workflows/launch_small_fast.yaml b/.github/workflows/launch_small_fast.yaml
@@ -41,7 +41,7 @@ jobs:
       - name: Install locally
         run: |
           python -m pip install --upgrade pip
-          pip install -e .[test] "jax[cpu]==0.4.30"
+          pip install -e .[test] "jax[cpu]==0.4.38"
 
       - name: Launch Small Fast TPU Train LM job
         run: |

diff --git a/.github/workflows/run_entry_tests.yaml b/.github/workflows/run_entry_tests.yaml
@@ -9,7 +9,7 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.10"]
-        jax-version: ["0.4.26"]
+        jax-version: ["0.4.38"]
 
     steps:
       - uses: actions/checkout@v3

diff --git a/.github/workflows/run_pre_commit.yaml b/.github/workflows/run_pre_commit.yaml
@@ -10,7 +10,7 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.10"]
-        jax-version: ["0.4.14"]
+        jax-version: ["0.4.38"]
 
     steps:
       - uses: actions/checkout@v3

diff --git a/.github/workflows/run_ray_tests.yaml b/.github/workflows/run_ray_tests.yaml
@@ -9,7 +9,7 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.10"]
-        jax-version: ["0.4.26"]
+        jax-version: ["0.4.38"]
 
     steps:
       - uses: actions/checkout@v3

diff --git a/.github/workflows/run_tests.yaml b/.github/workflows/run_tests.yaml
@@ -10,7 +10,7 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.10"]
-        jax-version: ["0.4.26"]
+        jax-version: ["0.4.38"]
 
     steps:
       - uses: actions/checkout@v3

diff --git a/infra/helpers/setup-tpu-vm-tests.sh b/infra/helpers/setup-tpu-vm-tests.sh
@@ -105,7 +105,7 @@ pip install -U wheel
 
 # jax and jaxlib
 # libtpu sometimes has issues installing for clinical (probably firewall?)
-retry pip install -U "jax[tpu]@git+https://github.com/dlwh/jax@retry_refuse" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+retru pip install -U "jax[tpu]==0.4.38" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
 
 # clone levanter
 git clone $REPO levanter

diff --git a/infra/helpers/setup-tpu-vm.sh b/infra/helpers/setup-tpu-vm.sh
@@ -105,7 +105,7 @@ pip install -U wheel
 
 # jax and jaxlib
 # libtpu sometimes has issues installing for clinical (probably firewall?)
-retry pip install -U "jax[tpu]@git+https://github.com/dlwh/jax@retry_refuse" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+retru pip install -U "jax[tpu]==0.4.38" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
 
 # clone levanter
 git clone $REPO levanter