diff --git a/requirements.txt b/requirements.txt
index 6d5e29024..c545288b2 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,5 @@
 --extra-index-url https://download.pytorch.org/whl/cpu
-jax>=0.6.2
+jax@git+https://github.com/jax-ml/jax.git
 jaxlib>=0.4.30
 grain
 google-cloud-storage>=2.17.0
diff --git a/requirements_with_jax_ai_image.txt b/requirements_with_jax_ai_image.txt
index 955a5e76f..834f8a70a 100644
--- a/requirements_with_jax_ai_image.txt
+++ b/requirements_with_jax_ai_image.txt
@@ -1,7 +1,7 @@
 # Requirements for Building the MaxDifussion Docker Image
 # These requirements are additional to the dependencies present in the JAX AI base image.
 --extra-index-url https://download.pytorch.org/whl/cpu
-jax>=0.6.2
+jax@git+https://github.com/jax-ml/jax.git
 jaxlib>=0.4.30
 grain
 google-cloud-storage>=2.17.0