diff --git a/requirements_jax.txt b/requirements_jax.txt
index 76f9015..f07c738 100644
--- a/requirements_jax.txt
+++ b/requirements_jax.txt
@@ -1 +1 @@
-jax==0.4.14
+jax==0.4.21