From ac2f79780eb23abf71aec568ea52eecd81e383f3 Mon Sep 17 00:00:00 2001
From: Archis Joglekar <joglekara@users.noreply.github.com>
Date: Sun, 21 Jul 2024 17:52:49 -0700
Subject: [PATCH] bump JAX version (#52)

---
 env.yaml                                | 2 +-
 env_gpu.yaml                            | 6 +++---
 tests/test_lpse2d/test_tpd_threshold.py | 5 +++--
 3 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/env.yaml b/env.yaml
index 7ec51bb..7fcf986 100644
--- 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
index c86d445..fb843dd 100644
--- 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
index 59958cd..7b60c60 100644
--- 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