Has anyone successfully used HSSM/Numpyro with GPU on recent JAX versions (Linux/Windows)?

[JamesWeiChen]

Hi all,

I’ve been trying to run HSSM models with GPU acceleration on my local machine (Linux, NVIDIA RTX 5080, CUDA 12.8). However, I ran into compatibility issues between NumPyro and the latest JAX releases:

• JAX ≥ 0.4.31 removed pjit_p, but older NumPyro (0.15.3) still references it.
• Updating to the newest JAX (0.7.x) breaks NumPyro; downgrading JAX to 0.4.x works but requires finding the exact matching jaxlib + CUDA/cuDNN combo.
• I tried combinations like jax==0.4.28 with jaxlib==0.4.28+cuda12.cudnn91 and 0.4.29+cuda12.cudnn91, but GPU initialization still fails (falling back to CPU).

Questions:

1. Has anyone recently managed to get NumPyro + HSSM running on GPU with current JAX versions (0.4.x or newer)
2. Which JAX/jaxlib/numpyro/PyMC version combinations worked for you?
3. Has anyone successfully set this up on Windows (or is Linux the only reliable option right now)?
4. Are there plans to update NumPyro for compatibility with JAX ≥ 0.5?

Any guidance or working environment specs would be greatly appreciated!

[digicosmos86]

Hello @JamesWeiChen,

Thanks for letting us know. We'll try to update numpyro dependency for hssm. We'll let you know when we are done.

Thanks!
Paul

[digicosmos86]

Hi @JamesWeiChen,

I did some digging into the issue. Although we will be updating the numpyro and pymc dependencies in HSSM, it seems that it won't necessarily fix your problem. Both numpyro 0.17.0 and 0.18.0 (last) require Jax 0.4.25, so there is not much we can do on our end to update its JAX compatibility. We can help you find the best compatible version of JAX though. For that purpose, can you let us know how you are installing HSSM, and the cuda and cudnn version in your set up?

Thanks!

[JamesWeiChen]

Hi @JamesWeiChen,

I did some digging into the issue. Although we will be updating the numpyro and pymc dependencies in HSSM, it seems that it won't necessarily fix your problem. Both numpyro 0.17.0 and 0.18.0 (last) require Jax 0.4.25, so there is not much we can do on our end to update its JAX compatibility. We can help you find the best compatible version of JAX though. For that purpose, can you let us know how you are installing HSSM, and the cuda and cudnn version in your set up?

Thanks!

Just FYI, pip install numpyro==0.19.0 can now install the numpyro 0.19.0, and based on the release note, it fix the problem above. I am tesating it now.

[digicosmos86]

Looks like it was just released 4 hours ago. Thanks for letting us know. We'll incorporate this update. In the meantime, let us know if it indeed solve the problem

[JamesWeiChen]

Looks like it was just released 4 hours ago. Thanks for letting us know. We'll incorporate this update. In the meantime, let us know if it indeed solve the problem

I still cannot work with my computer.

Environment

• OS: Pop!_OS 22.04
• GPU: NVIDIA RTX 5080
• Driver: 570.172.08
• System CUDA: 12.9
• cuDNN: 9.11
• Python: 3.11 (conda)
• JAX: 0.7.0 (CUDA 12 wheel)

I tried running JAX on GPU using the CUDA and cuDNN installed on my system, but it failed. I then attempted to use the cudatoolkit package within my conda environment, but that also failed. I am quite new to Linux, which might be contributing to the issue.

[JamesWeiChen]

After few hours of work... it seems like the problem is that Jax does not support 5080...

[digicosmos86]

nVidia does have an NGC for JAX with support for Blackwell GPUs, so 5080 should be supported. Can you provide a bit more details on how you got JAX installed into your conda environment? Was it a pip install or a conda install?

[JamesWeiChen]

nVidia does have an NGC for JAX with support for Blackwell GPUs, so 5080 should be supported. Can you provide a bit more details on how you got JAX installed into your conda environment? Was it a pip install or a conda install?

I tried both methods I think.

[digicosmos86]

Does the GPU version of PyTorch work? Just to see whether it's a jax-specific thing

[JamesWeiChen]

Does the GPU version of PyTorch work? Just to see whether it's a jax-specific thing

Thanks for the suggestion. I am able to run the code with gpu now. But I did too many things at once... I don't know which step fix the problem... I think it is because the JAX was trying to use many different CUDA toolkits.

Anyway... now it works.
Python version: 3.11.13 (main, Jun 5 2025, 13:12:00) [GCC 11.2.0]
JAX version: 0.7.0
jaxlib version: 0.7.0
NumPyro version: 0.19.0
HSSM version: 0.2.8

[JamesWeiChen]

Setting Up JAX + NumPyro + HSSM with CUDA 12.9 and cuDNN 9.11 on Pop!_OS

This is how I successfully configured JAX (GPU), NumPyro, and HSSM on a machine with NVIDIA RTX 5080 using CUDA 12.9 and cuDNN 9.11 on Pop!_OS (Ubuntu 22.04 base).

1. Install CUDA 12.9 (System-wide)

Add the CUDA 12.9 local repository and install the toolkit:

wget https://developer.download.nvidia.com/compute/cuda/12.9.0/local_installers/cuda-repo-ubuntu2204-12-9-local_12.9.0-575.51.03-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2204-12-9-local_12.9.0-575.51.03-1_amd64.deb
sudo cp /var/cuda-repo-ubuntu2204-12-9-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda-toolkit-12-9

2. Configure Environment Variables

Add CUDA paths to .bashrc:

echo 'export PATH=/usr/local/cuda-12.9/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.9/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

Check installation:

nvcc --version

(Should show: Cuda compilation tools, release 12.9)

3. Install cuDNN 9.11 (System-wide)

Add the cuDNN local repository and install:

wget https://developer.download.nvidia.com/compute/cudnn/9.11.0/local_installers/cudnn-local-repo-ubuntu2204-9.11.0_1.0-1_amd64.deb
sudo dpkg -i cudnn-local-repo-ubuntu2204-9.11.0_1.0-1_amd64.deb
sudo cp /var/cudnn-local-repo-ubuntu2204-9.11.0/cudnn-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cudnn

4. Create a Clean Conda Environment

Use a fresh environment to avoid conflicts with system CUDA:

conda create -n hssm python=3.11 -y
conda activate hssm

5. Install JAX (CUDA 12 Support)

Install JAX with the CUDA 12 wheel (no conda CUDA packages):

pip install --upgrade "jax[cuda12_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

6. Install NumPyro and HSSM

Ensure NumPyro is version 0.19.0 (compatible with HSSM):

pip install numpyro==0.19.0 hssm

7. Verify GPU Setup

Run a quick test:

python - <<'EOF'
import jax
import jax.numpy as jnp
print("Devices:", jax.devices())
print("Backend:", jax.default_backend())
print("Computation:", jnp.sum(jnp.array([1.0,2.0,3.0])**2))
EOF

Expected output:

Devices: [CudaDevice(id=0)]
Backend: gpu
Computation: 14.0

[digicosmos86]

Thank you so much for the info, @JamesWeiChen! I'll turn this issue into a discussion thread so more people can see this