diff --git a/pyproject.toml b/pyproject.toml
index 4b1f36da..bb7fc130 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,7 +6,7 @@ readme = "README.md"
 requires-python = ">=3.9,<3.11"
 dependencies = [
   "absl-py>=1.4.0",
-  "jax @ git+https://github.com/google/jax@9a21ff0780d1a16adf35b22dd7521e2804833760",
+  "jax @ git+https://github.com/google/jax@a0c1265bbae2c3ec644d6181f23264b4794e9eac",
   "triton-nightly @ https://aiinfra.pkgs.visualstudio.com/2692857e-05ef-43b4-ba9c-ccf1c22c437c/_packaging/07c94329-d4c3-4ad4-9e6b-f904a60032ec/pypi/download/triton-nightly/2.1.dev20230714011643/triton_nightly-2.1.0.dev20230714011643-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"
 ]