RuntimeError: CUDNN_BACKEND_TENSOR_DESCRIPTOR cudnnFinalize failed cudnn_status: CUDNN_STATUS_SUBLIBRARY_LOADING_FAILED

[Plankson]

When I run the following code, the error accur:

import jax
from flax import linen as nn

class ConvTest(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Conv(16, kernel_size=(3, 3))(x)
        return x

key = jax.random.PRNGKey(0)
x = jax.random.normal(key, (1, 32, 32, 3))
model = ConvTest()
params = model.init(key, x)
output = model.apply(params, x)

2025-04-01 17:16:46.124778: W external/xla/xla/service/gpu/nvptx_compiler.cc:765] The NVIDIA driver's CUDA version is 12.2 which is older than the ptxas CUDA version (12.8.93). Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.

Unable to load any of {libcudnn_engines_precompiled.so.9.5.0, libcudnn_engines_precompiled.so.9.5, libcudnn_engines_precompiled.so.9, libcudnn_engines_precompiled.so}
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 4, in __call__
  File "/home/qingyunpeng/miniconda3/envs/jaxmpc/lib/python3.9/site-packages/flax/linen/linear.py", line 663, in __call__
    y = conv_general_dilated(
RuntimeError: CUDNN_BACKEND_TENSOR_DESCRIPTOR cudnnFinalize failed cudnn_status: CUDNN_STATUS_SUBLIBRARY_LOADING_FAILED

I have install jax==0.4.29, jaxlib==0.4.29+cuda12.cudnn91, while the nvidia-smi results is:

+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.129.03             Driver Version: 535.129.03   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA A100-SXM4-80GB          On  | 00000000:1F:00.0 Off |                    0 |
| N/A   26C    P0              94W / 400W |  21642MiB / 81920MiB |    100%      Default |
|                                         |                      |             Disabled |
+-----------------------------------------+----------------------+----------------------+

Meanwhile, I have installed cuda 12.1.0 and cudnn 9.1.1.17 by

conda install nvidia/label/cuda-12.1.0::cuda-toolkit -y
wget https://repo.anaconda.com/pkgs/main/linux-64/cudnn-9.1.1.17-cuda12_1.conda
conda install --use-local cudnn-9.1.1.17-cuda12_.conda

It seems that the cuda/cudnn/jax versions are all corresponding. Meanwhile, I found that the nn.Dense module can work, which makes me so confused. Please help me🙏

>>> import jax
>>> import jax.numpy as jnp
>>> from flax import linen as nn
>>> 
>>> class DenseTest(nn.Module):
...     @nn.compact
...     def __call__(self, x):
...         x = nn.Dense(16)(x)
...         return x
... 
>>> key = jax.random.PRNGKey(0)
>>> x = jax.random.normal(key, (1, 32, 32, 3))
>>> model = DenseTest()
>>> params = model.init(key, x)
>>> output = model.apply(params, x)
>>> print(output.shape) 
(1, 32, 32, 16)