diff --git a/pyproject.toml b/pyproject.toml
index c69f2f4..4090ea2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -21,7 +21,7 @@ maintainers = [
 
 # TODO add gpu channels
 dependencies = [
-  "jax <= 0.4.35",
+  "jax <= 0.4.36",
   "jaxlib",
   "numpy",
 ]