docker gpu attempt

Dockerfile

@@ -20,9 +20,9 @@ COPY requirements.txt .
 RUN python3.11 -m pip install -r requirements.txt
 
 RUN python3.11 -m pip install \
-    jaxlib==0.3.25+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html \
-    jax==0.3.25 \
-    optax==0.1.5
+    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
 
 ENV PYGLFW_PREVIEW=1
 

pyproject.toml

@@ -19,8 +19,8 @@ numpy-hilbert-curve = "^1.0.1"
 wandb = "^0.17.0"
 seaborn = "^0.13.2"
 jax = "^0.4.28"
-jaxlib = "^0.4.28"
 optax = "^0.2.2"
+scikit-learn = "^1.5.0"
 
 
 [build-system]

requirements.txt

@@ -19,6 +19,7 @@ idna==3.7
 jax==0.4.28
 jaxlib==0.4.28
 Jinja2==3.1.4
+joblib==1.4.2
 jsonschema==4.22.0
 jsonschema-specifications==2023.12.1
 kiwisolver==1.4.5
@@ -50,6 +51,7 @@ referencing==0.35.1
 requests==2.32.2
 rich==13.7.1
 rpds-py==0.18.1
+scikit-learn==1.5.0
 scipy==1.13.1
 seaborn==0.13.2
 sentry-sdk==2.4.0
@@ -59,6 +61,7 @@ smmap==5.0.1
 streamlit==1.35.0
 sympy==1.12
 tenacity==8.3.0
+threadpoolctl==3.5.0
 toml==0.10.2
 toolz==0.12.1
 tornado==6.4

src/plots.py

@@ -3,7 +3,6 @@
 # by: Noah Syrkis
 
 # imports
-import plotly as py
 import pandas as pd
 import jax.numpy as jnp
 import numpy as np