From e506d80bc83a9e3a1071e1ac5d8bb6dc77aa212d Mon Sep 17 00:00:00 2001
From: Dario Panici <dpanici@princeton.edu>
Date: Fri, 4 Oct 2024 12:42:15 -0400
Subject: [PATCH] pin jax version temporarily

---
 requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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