diff --git a/docs/requirements.txt b/docs/requirements.txt
index 1cbce0043c..1ec5310527 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -1,7 +1,7 @@
 sphinx>=3.3.1
 sphinx-book-theme
 Pygments>=2.6.1
-jax>=0.3
+jax>=0.4
 jaxlib
 ipykernel
 myst_nb