From 85add667ede328023a14dedab65780e29b7140c2 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Mon, 24 Feb 2025 17:50:29 -0500
Subject: [PATCH] Change documentation to recommend libtpu from pypi instead of
GCS.
---
README.md | 4 ++--
docs/installation.md | 4 ++--
setup.py | 2 +-
3 files changed, 5 insertions(+), 5 deletions(-)
diff --git a/README.md b/README.md
index 53b1979d8..0aca7cf58 100644
--- a/README.md
+++ b/README.md
@@ -398,8 +398,8 @@ Some standouts:
|-----------------|-----------------------------------------------------------------------------------------------------------------|
| CPU | `pip install -U jax` |
| NVIDIA GPU | `pip install -U "jax[cuda12]"` |
-| Google TPU | `pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html` |
-| AMD GPU (Linux) | Follow [AMD's instructions](https://github.com/jax-ml/jax/blob/main/build/rocm/README.md). |
+| Google TPU | `pip install -U "jax[tpu]"` |
+| AMD GPU (Linux) | Follow [AMD's instructions](https://github.com/jax-ml/jax/blob/main/build/rocm/README.md). |
| Mac GPU | Follow [Apple's instructions](https://developer.apple.com/metal/jax/). |
| Intel GPU | Follow [Intel's instructions](https://github.com/intel/intel-extension-for-openxla/blob/main/docs/acc_jax.md). |
diff --git a/docs/installation.md b/docs/installation.md
index b3afad269..ee675dd1e 100644
--- a/docs/installation.md
+++ b/docs/installation.md
@@ -20,7 +20,7 @@ different builds for different operating systems and accelerators.
* **TPU (Google Cloud TPU VM)**
```
- pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+ pip install -U "jax[tpu]"
```
(install-supported-platforms)=
@@ -199,7 +199,7 @@ To install JAX along with appropriate versions of `jaxlib` and `libtpu`, you can
the following in your cloud TPU VM:
```bash
-pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+pip install "jax[tpu]"
```
For users of Colab (https://colab.research.google.com/), be sure you are
diff --git a/setup.py b/setup.py
index 08696753c..20d6652b7 100644
--- a/setup.py
+++ b/setup.py
@@ -73,7 +73,7 @@ setup(
'ci': [f'jaxlib=={_latest_jaxlib_version_on_pypi}'],
# Cloud TPU VM jaxlib can be installed via:
- # $ pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+ # $ pip install "jax[tpu]"
'tpu': [
f'jaxlib>={_current_jaxlib_version},<={_jax_version}',
f'libtpu=={_libtpu_version}',