From 4e73ff5191f7dfdd3719c10175266ec2504bc1dd Mon Sep 17 00:00:00 2001 From: Yifei Teng Date: Fri, 6 Dec 2024 19:43:18 +0000 Subject: [PATCH] Make it easy to switch to nightly libtpu and jax We can flip USE_NIGHTLY in setup.py --- setup.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/setup.py b/setup.py index 827c31f14c6..0f2025ac4bd 100644 --- a/setup.py +++ b/setup.py @@ -64,11 +64,18 @@ base_dir = os.path.dirname(os.path.abspath(__file__)) -# libtpu 0.0.5 and JAX 0.4.36 are cut on Dec 5, 2024. If we're switching to a nightly build, -# we should not regress to an earlier version. +USE_NIGHTLY = False # whether to use nightly or stable libtpu and jax +_date = '20241205' _libtpu_version = f'0.0.5' -_libtpu_storage_path = f'https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu/libtpu-{_libtpu_version}-py3-none-linux_x86_64.whl' _jax_version = f'0.4.36' +_libtpu_wheel_name = f'libtpu-{_libtpu_version}' + +if USE_NIGHTLY: + _libtpu_version += f".dev{_date}" + _jax_version += f".dev{_date}" + _libtpu_wheel_name += f".dev{_date}+nightly" + +_libtpu_storage_path = f'https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu/{_libtpu_wheel_name}-py3-none-linux_x86_64.whl' def _get_build_mode():