diff --git a/pyproject.toml b/pyproject.toml
index dbfc3278..c2f60f89 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,7 +37,7 @@ dev = [
     "beartype>=0.19,<0.23",
     "ipykernel>=6,<7",
     "ipython>=8,<9",
-    "jax[cpu]>=0.4,<0.6",
+    "jax[cpu]>=0.4,<0.9",
     "jupyter>=1,<2",
     "mypy>=1,<2",
     "nbval>=0.11,<0.12",