Clarify installation instructions for GPU usage

[basileMarchand]

Hello again !

In the installation section, it is mentioned that running pip install ikpls is sufficient, which is true for CPU usage.
However, for GPU usage, there is additional work required, particularly for the installation of Jax. A user might mistakenly believe that running pip install ikpls on a computer with a GPU would automatically configure everything for GPU usage, which is not the case.

Could you please add a warning in the documentation to inform users that if they want to use the GPU capabilities, they need to install Jax separately and give ? Ideally, provide some guidance on how to perform this installation.

Thank you !

For review openjournals/joss-reviews#6533

[Sm00thix]

Hi @basileMarchand!

I believe this is what I address in the Prerequisites section in the README. There, I have provided a link to the JAX Installation Guide which addresses configuration for usage with CPU, GPU, and TPU. As the JAX installation guide may change, I believe the best I can do is to refer users to the official documentation to ensure that nothing I write will be out of date.

Please let me know if I have understood correctly :-)

[boisgera]

Hi @Sm00thix!

I agree with you that you give the relevant information in the first section.

But I guess that some potential users won't read everything in the README linearly but will instead skip to the first piece of code that looks like an install command (especially if they are not familiar with the JAX ecosystem and just care about having a fast implementation of the algorithm).

Maybe if you add in the prerequisite section something like :

To enable GPU support for IKPLS, you need to install JAX with CUDA before, typically with:

pip3 install -U "jax[cuda12]"

Check the JAX Installation Guide for details. [... then more about TPU, etc. ...]

that would highlight the need for the extra step? (I agree that you should not try to replicate the full JAX installation guide in your README but link to it instead!)

Or do you think that this addition would be counterproductive?

[Sm00thix]

Hi @boisgera and @basileMarchand!

I have taken the suggestion by @boisgera and, additionally, also added clarification for a JAX TPU installation.
I agree that this is the most user-friendly approach. I have clarified that if a user desires a more customized JAX installation, they should follow the JAX Installation Guide. These changes have been implemented in e318e8e.

If you agree with this, then we can consider this issue closed :-)

[basileMarchand]

Thanks for this modification. It works perfectly for me. The user has the minimal info needed to start the installation, so it's good for me.