conda install doesn't work with latest jax / jaxlib

[jbloom]

I'm not really sure if this is a evofr issue versus a jax issue, but I'm going to post it here for reference since it certainly affects the ability to install a running version of evofr. (I also discussed this with @huddlej @plsteinberg).

If I just install evofr with conda, I get an error that in the end boils down to jax throwing a lengthy message that ends with:

XlaRuntimeError: INTERNAL: libdevice not found at ./libdevice.10.bc

Looking at the jax CHANGELOG, it says that starting with jax 0.4.31:

libdevice.10.bc is no longer bundled with CUDA wheels. It must be installed either as a part of local CUDA installation, or via NVIDIA's CUDA pip wheels.

I don't really understand what this means, but I simply tried building a new conda environment that included evofr but also pinned jax and jaxlib to 0.4.28 rather than letting them be the latest version on conda, which is 0.4.31 as of now.

That fixed the problem for me.

So it appears to me that the current conda installation of evofr requires pinning jax and jaxlib to 0.4.28 (or at least something before 0.4.31).

[huddlej]

Thank you, @jbloom! I'm pretty sure this is a JAX issue that is specific to the community-maintained Conda package and the fact that you're installing JAX on a system with GPUs (the Hutch's gizmo cluster).

The description of the community-maintained Conda package linked above mentions how to install the GPU utilities you need, if you really want to use a GPU, and it links to tips and tricks for installing tools that depend on JAX in different environments. If you don't plan to use a GPU, you can install the JAX libraries from a Conda package built just for CPUs. I did this like so on gizmo:

mamba create -c conda-forge -c bioconda -n evofr evofr jax jaxlib=*=cpu*

Installing this way, I was able to run the MLR models from evofr with only the following warning:

An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.

I'd leave this to @marlinfiggins to decide how he wants to document installation of evofr. Since this package does depend so heavily on JAX, it might be worth including the above reference links or examples.

[jbloom]

Yes, I think this second solution that @huddlej suggests is better.