No registered implementation for custom call to cusolver_gesvdj_ffi for platform CUDA

[ZedongPeng]

I am using jnp.linalg.norm to compute the norm of a matrix on GPU. However, when I set ord='nuc' or ord=2, I encounter an error. Any suggestions on how to resolve this issue?

Here is the code I am running:

>>> x = jnp.array([[1., 2., 3.],
...                [4., 5., 7.]])
>>> jnp.linalg.norm(x)  # Frobenius norm
Array(10.198039, dtype=float32)
>>> jnp.linalg.norm(x, ord='nuc')  # nuclear norm
Array(10.762535, dtype=float32)
>>> jnp.linalg.norm(x, ord=1)  # 1-norm
Array(10., dtype=float32)

Output

10.198039
Traceback (most recent call last):
  File "test6.py", line 8, in <module>
    print(jnp.linalg.norm(x, ord='nuc'))  # nuclear norm
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jaxlib.xla_extension.XlaRuntimeError: UNIMPLEMENTED: No registered implementation for custom call to cusolver_gesvdj_ffi for platform CUDA
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The GPU information

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.54.14              Driver Version: 550.54.14      CUDA Version: 12.4     |
|-----------------------------------------+------------------------+----------------------+
| 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 H100 80GB HBM3          On  |   00000000:55:00.0 Off |                    0 |
| N/A   33C    P0             72W /  700W |       0MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
                                                                                         
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|  No running processes found                                                             |
+-----------------------------------------------------------------------------------------+

Python and JAX versions:

Python 3.11.9 (main, Apr 19 2024, 16:48:06) [GCC 11.2.0] on linux

# Name                    Version                   Build  Channel
jax                       0.4.35                   pypi_0    pypi
jax-cuda12-pjrt           0.4.33                   pypi_0    pypi
jax-cuda12-plugin         0.4.33                   pypi_0    pypi
jaxlib                    0.4.35                   pypi_0    pypi

[dfm]

It looks to me like the issue is that the jax/jaxlib versions are out of sync with jax-cuda12-pjrt/jax-cuda12-plugin. I expect that if you update:

pip install -U "jax[cuda12]"

everything should work!

[ZedongPeng]

Thank you. pip install -U "jax[cuda12]" resolved this issue.