Bump jax 38 (#846)

stanford-crfm · Dec 24, 2024 · 736d9ea · 736d9ea
1 parent d9678f5
commit 736d9ea
Show file tree

Hide file tree

Showing 8 changed files with 8 additions and 9 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/docker/tpu/Dockerfile.base b/docker/tpu/Dockerfile.base
@@ -5,8 +5,7 @@ RUN pip install virtualenv
 # venv binaries encode their directory, so we need to setup the venv in the final location
 RUN virtualenv -p python3.10 /opt/levanter/.venv
 ENV PATH /opt/levanter/.venv/bin:$PATH
-#RUN /opt/levanter/.venv/bin/pip install -U "jax[tpu]==0.4.34" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-RUN /opt/levanter/.venv/bin/pip install -U "jax[tpu]@git+https://github.com/dlwh/jax@retry_refuse" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+RUN /opt/levanter/.venv/bin/pip install -U "jax[tpu]==0.4.38" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
 
 # Install package dependencies to make incremental builds faster.
 WORKDIR /tmp/

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