diff --git a/requirements_jax.txt b/requirements_jax.txt
index 6da1e323..aa8a8ed0 100644
--- a/requirements_jax.txt
+++ b/requirements_jax.txt
@@ -1 +1 @@
-jax==0.9.0
+jax==0.9.0.1