diff --git a/requirements.txt b/requirements.txt
index 227f95eac..0e1223279 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,7 +2,7 @@ colorama
 diffrax >= 0.4.1
 h5py >= 3.0.0, < 4.0
 interpax >= 0.3.3
-jax[cpu] >= 0.3.2, < 0.5.0
+jax[cpu] >= 0.3.2, < 0.4.34
 matplotlib >= 3.5.0, < 4.0.0
 mpmath >= 1.0.0, < 2.0
 netcdf4 >= 1.5.4, < 2.0