Open YigitElma opened 4 weeks ago
Response from NERSC support about installing the latest versions of JAX on Perlmutter:
I would highly recommend that you use containers to run JAX at NERSC, this make your life significantly easier down the line (especially as you start distributing your computation over several nodes): https://docs.nersc.gov/development/languages/python/using-python-perlmutter/#nvidia-containers
Otherwise, to use a nightly release of JAX (versions above 0.4.29) following our documented approach, you can use this page: https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
I have also been told that, with the latest releases, a simple
pip install --upgrade "jax[cuda12]"
might work on Perlmutter but I have yet to test it (at which point I will update the documentation).
@daniel-dudt Were you able to install a recent version on jax + DESC on Perlmutter? If so, can you post your instructions here so I can double-check.
Originally posted by @YigitElma in https://github.com/PlasmaControl/DESC/issues/1293#issuecomment-2436078532
Since they use an older JAX version 0.4.23, current instructions require to change scipy version from requirements file for proper installation. If newer versions of JAX are supported for these clusters, the scipy version change is redundant. JAX should be greater than 0.4.24 ( also this is the new requirement for DESC).