mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update Cloud TPU install command to be simpler.

Browse Source 
We used to need the extra stuff for a very old Cloud TPU VM image, but we don't anymore.
This commit is contained in:
Skye Wanderman-Milne 2023-02-27 23:15:19 +00:00
parent 055fa6b90f
commit 56b237cfbc
1 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
README.md
View File

@ -413,7 +413,7 @@ pip install --upgrade "jax[cpu]"
```
On Linux, it is often necessary to first update `pip` to a version that supports
`manylinux2014` wheels. Also note that for Linux, we currently release wheels for `x86_64` architectures only, other architectures require building from source. Trying to pip install with other Linux architectures may lead to `jaxlib` not being installed alongside `jax`, although `jax` may successfully install (but fail at runtime).
`manylinux2014` wheels. Also note that for Linux, we currently release wheels for `x86_64` architectures only, other architectures require building from source. Trying to pip install with other Linux architectures may lead to `jaxlib` not being installed alongside `jax`, although `jax` may successfully install (but fail at runtime).
**These `pip` installations do not work with Windows, and may fail silently; see
[above](#installation).**
@ -502,8 +502,7 @@ JAX also provides pre-built wheels for
To install JAX along with appropriate versions of `jaxlib` and `libtpu`, you can run
the following in your cloud TPU VM:
```bash
pip install --upgrade pip
pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
```
### pip installation: Colab TPU