I want to use a singularity container with jax preinstalled e.g. based on this image. On top of that, I want to install PyMC with numpyro, to get NUTS sampling with GPU acceleration. My attempt is to use the def file Bootstrap: docker From: nvcr.io/nvidia/jax:23.04-py3 %post pip install pymc numpyro %test python -c "import jax;print(jax.devices())" The test fails, since the post-step uninstalls the local jax packages and installs a new version from PyPA. That is not acceptable, since the docker image comes with some very specific locally compiled jaxlib versions and the like. Can I give some instructions to pip to "install these two packages, and select some version that is compatible with what is already installed"? Continue reading...
