diff --git a/setup.cfg b/setup.cfg
index 5d630d7..93efbad 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -66,7 +66,7 @@ GBDT =
     lightgbm < 4
 neural_tangents =
     jaxlib <= 0.1.65
-    jax <= 0.2.12
+    jax <= 0.4.8
     neural_tangents < 0.6.0
 gpflow =
     tensorflow >= 2.5, < 2.9
@@ -80,7 +80,7 @@ all =
     matplotlib
     lightgbm < 4
     jaxlib <= 0.1.65
-    jax <= 0.2.12
+    jax <= 0.4.8
     neural_tangents < 0.6.0
     tensorflow >= 2.5, < 2.9
     tensorflow-probability >= 0.12, < 0.13