From c785eeb4e5e83fadbdd5107a1d43023d7729d7fa Mon Sep 17 00:00:00 2001 From: Jake VanderPlas Date: Fri, 12 Jan 2024 11:12:39 -0800 Subject: [PATCH] DOC: add a note about installing older GPU wheels --- docs/installation.md | 4 ++++ docs/tutorials/installation.md | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/docs/installation.md b/docs/installation.md index 769fa0732..5dba6757e 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -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 +``` diff --git a/docs/tutorials/installation.md b/docs/tutorials/installation.md index 93f9c6e36..34508c4fd 100644 --- a/docs/tutorials/installation.md +++ b/docs/tutorials/installation.md @@ -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 +``` \ No newline at end of file