[Python] jax.random.uniform causing segmentation fault when called on GPU but not on CPU, nor...

I ran the following 4 commands at the command line (bash):

JAX_PLATFORM_NAME=cpu python -c "import jax; import jax.numpy as jnp; key = jax.random.PRNGKey(1); print(jax.random.uniform(key, (2, 2),))"
JAX_PLATFORM_NAME=cpu python -c "import jax; import jax.numpy as jnp; key = jax.random.PRNGKey(1); print(jax.random.normal(key, (2, 2),))"
JAX_PLATFORM_NAME=gpu python -c "import jax; import jax.numpy as jnp; key = jax.random.PRNGKey(1); print(jax.random.uniform(key, (2, 2),))"
JAX_PLATFORM_NAME=gpu python -c "import jax; import jax.numpy as jnp; key = jax.random.PRNGKey(1); print(jax.random.normal(key, (2, 2),))"


All of them ran fine, except for the uniform sampling on GPU, which resulted in a Segmentation fault (core dumped) with exit code 139 (and exit code 245 when similar code was run as part of a longer program).

Partial output of nvidia-smi:

NVIDIA-SMI 565.57.01 Driver Version: 565.57.01 CUDA Version: 12.7

NVIDIA GeForce RTX 4090


jax and jaxlib versions:

Name: jax
Version: 0.4.37
Name: jaxlib
Version: 0.4.36


ChatGPT thinks it's a CUDA and jax compatibility issue but the page here seems to suggest CUDA >=12.1 should be fine.

Any ideas?

Update: I upgraded jax and jaxlib to 0.6.2 and now get:

~/anaconda3/envs/unc/lib/python3.10/site-packages/jaxlib/plugin_support.py:71: RuntimeWarning: JAX plugin jax_cuda12_plugin version 0.4.33 is installed, but it is not compatible with the installed jaxlib version 0.6.2, so it will not be used.
warnings.warn(
Segmentation fault (core dumped)

Continue reading...