Recommend the plugin in the CUDA installation instructions.

PiperOrigin-RevId: 631524916
google · May 7, 2024 · d84c7ad · d84c7ad
1 parent 5f70267
commit d84c7ad
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -397,7 +397,8 @@ Some standouts:
 | Hardware   | Instructions                                                                                                    |
 |------------|-----------------------------------------------------------------------------------------------------------------|
 | CPU        | `pip install -U "jax[cpu]"`                                                                                       |
-| NVIDIA GPU on x86_64 | `pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html`        |
+| NVIDIA GPU on x86_64 | `pip install -U "jax[cuda12]"`        |
+| NVIDIA GPU on x86_64 legacy | `pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html`        |
 | Google TPU | `pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html`                 |
 | AMD GPU    | Use [Docker](https://hub.docker.com/r/rocm/jax) or [build from source](https://jax.readthedocs.io/en/latest/developer.html#additional-notes-for-building-a-rocm-jaxlib-for-amd-gpus). |
 | Apple GPU  | Follow [Apple's instructions](https://developer.apple.com/metal/jax/).                                          |

diff --git a/docs/installation.md b/docs/installation.md
@@ -96,6 +96,9 @@ pip install --upgrade pip
 
 # NVIDIA CUDA 12 installation
 # Note: wheels only available on linux.
+pip install --upgrade "jax[cuda12]"
+
+# NVIDIA CUDA 12 installation legacy
 pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
 
@@ -202,6 +205,12 @@ pip install -U libtpu-nightly -f https://storage.googleapis.com/jax-releases/lib
 
 - `jaxlib` NVIDIA GPU (CUDA 12):
 
+```bash
+pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
+pip install -U --pre jax-cuda12-pjrt jax-cuda12-plugin -f https://storage.googleapis.com/jax-releases/jax_cuda_plugin_nightly_releases.html
+
+- `jaxlib` NVIDIA GPU (CUDA 12) legacy:
+
 ```bash
 pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_cuda12_releases.html
 ```