Update JAX version to fix TPU CI (#7601)

pytorch · Jul 1, 2024 · 593dce8 · 593dce8
1 parent 755ae84
commit 593dce8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -67,7 +67,7 @@
 _date = '20240628'
 _libtpu_version = f'0.1.dev{_date}'
 _libtpu_storage_path = f'https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/wheels/libtpu-nightly/libtpu_nightly-{_libtpu_version}-py3-none-any.whl'
-_jax_version = f'0.4.30.dev{_date}'
+_jax_version = f'0.4.31.dev{_date}'
 
 
 def _get_build_mode():