Do not use jax version 0.9.1 for now due to its several regressions

Author: JanLuca

pyproject.toml

@@ -15,7 +15,7 @@ requires-python = ">=3.11"
 dependencies = [
     "numpy >= 1.21.2",
     "scipy >= 1.7.1",
-    "jax[cpu]>=0.8.1",
+    "jax[cpu]>=0.8.1,<0.9.1",
     "h5py >= 3.6.0",
     "tqdm[notebook] >= 4.64.1",
     "tqdm-loggable >= 0.1.3",
@@ -51,7 +51,7 @@ dev = [
 ]
 
 [build-system]
-requires = ["hatchling", "scikit-build-core>=0.11.0", "nanobind>=2.7.0", "jaxlib>=0.8.1"]
+requires = ["hatchling", "scikit-build-core>=0.11.0", "nanobind>=2.7.0", "jaxlib>=0.8.1,<0.9.1"]
 build-backend = "hatchling.build"
 
 [tool.scikit-build]