diff --git a/setup.py b/setup.py
index 8da22ef88..1e19603c4 100644
--- a/setup.py
+++ b/setup.py
@@ -9,7 +9,7 @@
 # `pyproject.toml`, once optional dependencies defined there can reference one another
 extras_require = dict()
 extras_require["jax"] = [
-    "jax[cpu]<0.4.3; platform_system!='Windows'",
+    "jax[cpu]<0.4.4; platform_system!='Windows'",
 ]
 extras_require["zoo"] = [
     "tqdm>=4.0",