update ver requirement for jax

Author: ChiahsinChu

dmff/settings.py

@@ -1,4 +1,9 @@
-from jax.config import config
+try:
+    # jax < 0.4.20
+    from jax.config import config
+except:
+    # jax >= 0.4.20
+    from jax import config
 
 PRECISION = 'double'  # 'double'
 

requirements.txt

@@ -6,7 +6,7 @@ mkdocs-literate-nav>=0.4.1
 mkdocstrings>=0.19.0
 mkdocstrings-python>=0.7.0
 pygments>=2.12
-jax>=0.4.1, <=0.4.20
+jax>=0.4.1
 jaxlib>=0.4.1
 pymbar>=4.0.0
 tqdm