diff --git a/setup.cfg b/setup.cfg
index de49569..284f1ad 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -43,14 +43,14 @@ python_requires = >=3.8
 [options.extras_require]
 # JAX CPU
 jax_cpu =
-  jax==0.4.26
-  jaxlib==0.4.26
+  jax==0.4.13
+  jaxlib==0.4.13
   dm-haiku==0.0.12
 
 # JAX GPU
 jax_gpu =
-  jax[cuda]==0.4.26
-  jaxlib==0.4.264+cuda11.cudnn82
+  jax[cuda]==0.4.13
+  jaxlib==0.4.13+cuda11.cudnn82
   dm-haiku==0.0.12
 
 ###############################################################################