Improve jax installation instructions

lanl · Feb 6, 2023 · fd46f3d · fd46f3d
1 parent d064e45
commit fd46f3d
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/docs/source/install.rst b/docs/source/install.rst
@@ -53,13 +53,23 @@ The instructions above install a CPU-only version of SCICO. To install a version
 
 2. Install the version of jaxlib with GPU support, as described in the `JAX installation
    instructions  <https://github.com/google/jax#installation>`_.
-   The command will be something like
+   In the simplest case, the appropriate command is
 
    ::
 
       pip install -upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
 
+   but it may be necessary to explicitly specify the ``jaxlib`` version if the most recent release is not yet
+   supported by SCICO (as specified in the ``requirements.txt`` file), or if using a version of CUDA older
+   than 11.4, or CuDNN older than 8.2, in which case the command would be of the form
+   ::
+
+      pip install -upgrade "jaxlib==0.4.2+cuda11.cudnn82" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+
+   with appropriate substitution of ``jaxlib``, CUDA, and CuDNN version numbers.
+
+
 
 Additional Dependencies
 -----------------------