diff --git a/pyproject.toml b/pyproject.toml
index cd65095..bfaef5f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -26,7 +26,7 @@ dependencies = [
     "qnm==0.4.4",
     "seaborn~=0.13",
     "numpyro==0.19.0",
-    "jax>=0.4.25,<0.6.1",
+    "jax>=0.4.25,<0.8.3",
     "parse>=1.20",
     "jaxlib==0.4.38; sys_platform == 'darwin' and platform_machine == 'x86_64'",
     "tables~=3.10",
diff --git a/requirements.txt b/requirements.txt
index f4237b1..0463f3f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -9,7 +9,7 @@ jupyterlab
 ipywidgets
 ipykernel
 numpyro~=0.17
-jax>=0.4.25,<0.6.1
+jax>=0.4.25,<0.8.3
 parse>=1.20
 jaxlib==0.4.38; sys_platform == 'darwin' and platform_machine == 'x86_64'
 tables~=3.10