downgrade jax hpc cuda 11 test

syrkis · Jun 9, 2024 · a5db3c7 · a5db3c7
1 parent 24c3a12
commit a5db3c7
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 197 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -20,9 +20,8 @@ COPY requirements.txt .
 RUN python3.11 -m pip install -r requirements.txt
 
 RUN python3.11 -m pip install \
-    jaxlib==0.4.25+cuda11.cudnn86 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html \
-    jax==0.4.25 \
-    optax==0.2.2
+    pip install -U "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html \
+    optax==0.1.5
 
 ENV PYGLFW_PREVIEW=1
 

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -18,8 +18,6 @@ pyyaml = "^6.0.1"
 numpy-hilbert-curve = "^1.0.1"
 wandb = "^0.17.0"
 seaborn = "^0.13.2"
-jax = "^0.4.28"
-optax = "^0.2.2"
 scikit-learn = "^1.5.0"
 
 

diff --git a/requirements.txt b/requirements.txt
@@ -1,11 +1,9 @@
-absl-py==2.1.0
 altair==5.3.0
 attrs==23.2.0
 blinker==1.8.2
 cachetools==5.3.3
 certifi==2024.2.2
 charset-normalizer==3.3.2
-chex==0.1.86
 click==8.1.7
 contourpy==1.2.1
 cycler==0.12.1
@@ -16,8 +14,6 @@ fonttools==4.52.4
 gitdb==4.0.11
 GitPython==3.1.43
 idna==3.7
-jax==0.4.28
-jaxlib==0.4.28
 Jinja2==3.1.4
 joblib==1.4.2
 jsonschema==4.22.0
@@ -27,13 +23,10 @@ markdown-it-py==3.0.0
 MarkupSafe==2.1.5
 matplotlib==3.9.0
 mdurl==0.1.2
-ml-dtypes==0.4.0
 mpmath==1.3.0
 numpy==1.26.4
 numpy-hilbert-curve==1.0.1
 oeis==2023.3.10
-opt-einsum==3.3.0
-optax==0.2.2
 packaging==24.0
 pandas==2.2.2
 pillow==10.3.0