Skip to content

Commit 09309e6

Browse files
authored
Update conda-forge installation docs after CUDA 12 upgrade
1 parent 90de28c commit 09309e6

File tree

1 file changed

+4
-8
lines changed

1 file changed

+4
-8
lines changed

docs/installation.md

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -253,18 +253,14 @@ simply run:
253253
conda install jax -c conda-forge
254254
```
255255

256-
To install it on a machine with an NVIDIA GPU, run:
256+
If you run this command on machine with an NVIDIA GPU, this should install a CUDA-enabled package of `jaxlib`.
257+
258+
To ensure that the jax version you are installing is indeed CUDA-enabled, run:
257259

258260
```bash
259-
conda install "jaxlib=*=*cuda*" jax cuda-nvcc -c conda-forge -c nvidia
261+
conda install "jaxlib=*=*cuda*" jax -c conda-forge
260262
```
261263

262-
Note the `cudatoolkit` distributed by `conda-forge` is missing `ptxas`, which
263-
JAX requires. You must therefore either install the `cuda-nvcc` package from
264-
the `nvidia` channel, or install CUDA on your machine separately so that `ptxas`
265-
is in your path. The channel order above is important (`conda-forge` before
266-
`nvidia`).
267-
268264
If you would like to override which release of CUDA is used by JAX, or to
269265
install the CUDA build on a machine without GPUs, follow the instructions in the
270266
[Tips & tricks](https://conda-forge.org/docs/user/tipsandtricks.html#installing-cuda-enabled-packages-like-tensorflow-and-pytorch)

0 commit comments

Comments
 (0)