Calling nn.SelfAttention on a GPU run results in segfault

[tajtac]

I am using jax 0.5.3, jaxlib 0.5.3, flax 0.10.6, and cuda/12.6.1.
Below is a minimal working example. It always segfaults at params = model.init(rng, x), but only when using the GPU. If I start the script in the cpu-only mode, no problem occurs.

--------------------------

Minimal reproducible example

--------------------------

import jax
import jax.numpy as jnp
from flax import linen as nn

class SimpleSelfAttention(nn.Module):
d_model: int = 128
n_heads: int = 8

@nn.compact
def __call__(self, x):
    att = nn.SelfAttention(
        num_heads=self.n_heads,
        qkv_features=self.d_model,
        out_features=self.d_model,
    )(x)
    return att

B = 1 # batch size
D = 16 # sequence length
d_model = 128

x = jnp.ones((B, D, d_model), dtype=jnp.float32)

model = SimpleSelfAttention(d_model=d_model, n_heads=8)
print('Created an instance of SimpleSelfAttention. Segfault occurs after this.')
rng = jax.random.PRNGKey(0)

params = model.init(rng, x)
print('Initialized the model')

--------------------------

A couple of notes about the versions that I am using in case it is important:

--------------------------

1. The versions you see above are the result of me trying to get jax to work on my gpu cluster. I simply do 'pip install "jax[cuda12]" flax' and it installs these versions, even though a much higher version of jax is available.
2. If I do something like "pip install flax" or "pip install jax flax", it installs the latest version of jax, jaxlib and flax but then jax doesn't run on my cluster. It complains that the cudnn version used to compile was 9.8 and the current cudnn on my system is 9.4. I currently don't want to get my university IT people to install a later version of cuda/cudnn. I found that the way to get jax to work on my gpu is to install it using "pip install "jax[cuda12]".
3. Flax's website says you should install jaxlib first, then flax. But when installing jaxlib I can't specify "[cuda12]", which leads it to install the latest version, which, again, doesn't seem to work on my system. If I try to install "jax[cuda12]" before installing flax, flax just overwrites it with the latest version of jax anyway, which doesn't work as I said.

[vfdev-5]

@tajtac can you try to install the latest version of jax, jaxlib and flax using cuda12-local option:

pip install -U "jax[cuda12-local]"
pip install -U flax

• https://docs.jax.dev/en/latest/installation.html#pip-installation-nvidia-gpu-cuda-installed-locally-harder

Can you also use nnx instead of linen? I assume that the segfault is due to incompatible cuda versions and not related to linen or nnx...