avoid jax 0.5.1

Author: gboehl

requirements.txt

@@ -1,6 +1,6 @@
 # ubuntu seems to have issues with jax v0.4.6. Recheck when v0.4.7 is out
-jax>=0.4.13
-jaxlib>=0.4.13
+jax>=0.4.13,<0.5.1
+jaxlib>=0.4.13,<0.5.1
 grgrjax>=0.4.6
 scipy
 pyyaml

setup.py

@@ -34,8 +34,8 @@
         'linear': ['grgrlib>=0.1.22'],
     },
     install_requires=[
-        "jax>=0.4.13",
-        "jaxlib>=0.4.13",
+        "jax>=0.4.13,<0.5.1",
+        "jaxlib>=0.4.13,<0.5.1",
         "grgrjax>=0.4.6",
         "pyyaml",
         "scipy",