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

Merge pull request #19345 from jakevdp:installation-docs

Browse Source 
PiperOrigin-RevId: 597926417
This commit is contained in:
jax authors 2024-01-12 12:53:09 -08:00
commit 963cd6fd3d
2 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
docs/installation.md
View File

@ -280,3 +280,7 @@ pip install jax[cpu]==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_
# Install the jaxlib 0.3.25 CPU wheel directly
pip install jaxlib==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_releases.html
```
For specific older GPU wheels, be sure to use the `jax_cuda_releases.html` URL; for example
```
pip install jaxlib==0.3.25+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
```

4
docs/tutorials/installation.md
View File

@ -312,3 +312,7 @@ pip install jax[cpu]==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_
# Install the jaxlib 0.3.25 CPU wheel directly
pip install jaxlib==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_releases.html
```
For specific older GPU wheels, be sure to use the `jax_cuda_releases.html` URL; for example
```bash
pip install jaxlib==0.3.25+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
```