Merge pull request #26325 from hawkinsp:tpu2

PiperOrigin-RevId: 723298028

CHANGELOG.md

@@ -26,6 +26,9 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
 * Changes
   * `JAX_CPU_COLLECTIVES_IMPLEMENTATION` and `JAX_NUM_CPU_DEVICES` now work as
     env vars. Before they could only be specified via jax.config or flags.
+  * The `jax[tpu]` TPU extra no longer depends on the `libtpu-nightly` package.
+    This package may safely be removed if it is present on your machine; JAX now
+    uses `libtpu` instead.
 
 ## jax 0.5.0 (Jan 17, 2025)
 

setup.py

@@ -24,7 +24,6 @@
 _latest_jaxlib_version_on_pypi = '0.5.0'
 
 _libtpu_version = '0.0.8'
-_libtpu_nightly_terminal_version = '0.1.dev20241010+nightly.cleanup'
 
 def load_version_module(pkg_path):
   spec = importlib.util.spec_from_file_location(
@@ -77,8 +76,6 @@ def load_version_module(pkg_path):
         # $ pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
         'tpu': [
           f'jaxlib>={_current_jaxlib_version},<={_jax_version}',
-          # TODO(phawkins): remove the libtpu-nightly dependency in Q1 2025.
-          f'libtpu-nightly=={_libtpu_nightly_terminal_version}',
           f'libtpu=={_libtpu_version}',
           'requests',  # necessary for jax.distributed.initialize
         ],