downgrade jax hpc cuda 11 test

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

Hide file tree

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

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -19,6 +19,9 @@ numpy-hilbert-curve = "^1.0.1"
 wandb = "^0.17.0"
 seaborn = "^0.13.2"
 scikit-learn = "^1.5.0"
+jax = "0.4.10"
+jaxlib = "0.4.10"
+optax = "0.1.5"
 
 
 [build-system]

diff --git a/requirements.txt b/requirements.txt
@@ -1,9 +1,11 @@
+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.82
 click==8.1.7
 contourpy==1.2.1
 cycler==0.12.1
@@ -14,6 +16,8 @@ fonttools==4.52.4
 gitdb==4.0.11
 GitPython==3.1.43
 idna==3.7
+jax==0.4.10
+jaxlib==0.4.10
 Jinja2==3.1.4
 joblib==1.4.2
 jsonschema==4.22.0
@@ -23,10 +27,13 @@ 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.1.5
 packaging==24.0
 pandas==2.2.2
 pillow==10.3.0