bump JAX version (#52)

ergodicio · Jul 22, 2024 · ac2f797 · ac2f797
1 parent 2f6e93d
commit ac2f797
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 6 deletions.
diff --git a/env.yaml b/env.yaml
@@ -2,7 +2,7 @@ name: adept
 channels:
   - conda-forge
 dependencies:
-  - python=3.10
+  - python=3.12
   - ipython
   - ipykernel
   - pip

diff --git a/env_gpu.yaml b/env_gpu.yaml
@@ -2,15 +2,15 @@ name: adept-gpu
 channels:
   - conda-forge
 dependencies:
-  - python=3.10
+  - python=3.12
   - ipython
   - ipykernel
   - pip
   - pip:
       # works for regular pip packages
       - --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-      - jaxlib==0.4.20+cuda12.cudnn89
-      - jax==0.4.20
+      - jaxlib==0.4.28+cuda12.cudnn89
+      - jax==0.4.28
       - jaxopt
       - numpy
       - scipy

diff --git a/tests/test_lpse2d/test_tpd_threshold.py b/tests/test_lpse2d/test_tpd_threshold.py
@@ -3,10 +3,11 @@
 
 import numpy as np
 from utils.misc import setup_parsl
-from parsl.app.app import python_app
 
+# from parsl.app.app import python_app
 
-@python_app
+
+# @python_app
 def run_once(L, Te, I0):
     import yaml
     from adept import ergoExo