Adding pip installation guide (within conda environment)

[HMUNACHI]

A guide for local end-to-end installation in a conda environment (not the community conda-forge build) would be nice. I would love to add this to the readme.

[jakevdp]

I believe the only conda installation available is the community conda-forge build. Can you say more about why you rule that one out?

[HMUNACHI]

Thanks, so I'm talking about an additional guide for locally installing Jax via pip along with its required CUDA and CuDNN versions, but in a Conda environment. Something like:

conda install -c conda-forge cudatoolkit=11.8.0

python3 -m pip install nvidia-cudnn-cu11==8.6.0.163       
                                            
pip install --upgrade "jax[cuda11_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

mkdir -p $CONDA_PREFIX/etc/conda/activate.d

echo 'CUDNN_PATH=$(dirname $(python -c "import nvidia.cudnn;print(nvidia.cudnn.__file__)"))' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh

echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/:$CUDNN_PATH/lib' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh

source $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh

python3 -c "import jax.numpy as jnp; print(jnp.ones(5))"

Installing Jax-Cuda locally is generally non-trivial for less experienced developers so I felt a more detailed guide like this might make it easier.