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

Merge pull request #21897 from matsen:patch-1

Browse Source 
PiperOrigin-RevId: 643735130
This commit is contained in:
jax authors 2024-06-16 02:01:42 -07:00
commit ae0127d696
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
docs/installation.md
View File

@ -9,14 +9,14 @@ different builds for different operating systems and accelerators.
* **CPU-only (Linux/macOS/Windows)**
```
pip install -U jax[cpu]
pip install -U "jax[cpu]"
```
* **GPU (NVIDIA, CUDA 12)**
```
pip install -U "jax[cuda12]"
```
* **TPU (Google Cloud TPU VM) **
* **TPU (Google Cloud TPU VM)**
```
pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
```