operation cuInit(0) failed: CUDA_ERROR_NO_DEVICE

[jiulong1306]

PC, NVIDIA GeForce RTX 4060 Laptop GPU, wsl
When I run

from desc import set_device
set_device('gpu')

from desc.backend import print_backend_info
print_backend_info()

It returns

DESC version=0+unknown.
Using JAX backend: jax version=0.6.2, jaxlib version=0.6.2, dtype=float64.
Using device: NVIDIA GeForce RTX 4060 Laptop GPU (id=0), with 7.01 GB available memory.

But when I run

import desc.io

returns

ERROR:2025-10-29 22:18:42,187:jax._src.xla_bridge:444: Jax plugin configuration error: Exception when calling jax_plugins.xla_cuda12.initialize()
Traceback (most recent call last):
  File "/home/xiao/miniconda3/envs/desc/lib/python3.12/site-packages/jax/_src/xla_bridge.py", line 442, in discover_pjrt_plugins
    plugin_module.initialize()
  File "/home/xiao/miniconda3/envs/desc/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 324, in initialize
    _check_cuda_versions(raise_on_first_error=True)
  File "/home/xiao/miniconda3/envs/desc/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 281, in _check_cuda_versions
    local_device_count = cuda_versions.cuda_device_count()
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: jaxlib/cuda/versions_helpers.cc:113: operation cuInit(0) failed: CUDA_ERROR_NO_DEVICE

I just created a virtual environment Python=3.12 in WSL and followed the installation tutorial in the manual

[f0uriest]

I think this is a Jax issue, it looks like GPU support on WSL is still listed as experimental: https://docs.jax.dev/en/latest/installation.html#supported-platforms

[jiulong1306]

I think this is a Jax issue, it looks like GPU support on WSL is still listed as experimental: https://docs.jax.dev/en/latest/installation.html#supported-platforms

But I tested jax and it works properly on the GPU

import time
import jax
import jax.numpy as jnp
import numpy as np
from jax import random
 
def benchmark_operations():
    key = random.key(42)
    size = 8192
    
    a_jax = random.normal(key, (size, size), dtype=jnp.float32)
    b_jax = random.normal(key, (size, size), dtype=jnp.float32)
    
    a_np = np.array(a_jax)
    b_np = np.array(b_jax)
    
    @jax.jit
    def jax_matmul(a, b):
        return jnp.dot(a, b)

    def numpy_matmul(a, b):
        return np.dot(a, b)
    

    jax_matmul(a_jax, b_jax).block_until_ready()
    

    start = time.time()
    for _ in range(10):
        result_jax = jax_matmul(a_jax, b_jax)
        result_jax.block_until_ready()
    jax_time = time.time() - start
    

    start = time.time()
    for _ in range(10):
        result_np = numpy_matmul(a_np, b_np)
    numpy_time = time.time() - start
    
    return jax_time, numpy_time
 
jax_time, numpy_time = benchmark_operations()
print(f"JAX (GPU): {jax_time:.3f}s")
print(f"NumPy (CPU): {numpy_time:.3f}s")
print(f"rate: {numpy_time/jax_time:.1f}x")

It works

[YigitElma]

Hi @jiulong1306! Can you run these and share the output please?

from desc import set_device
set_device('gpu')

from desc.backend import print_backend_info
print_backend_info()

print(os.environ["CUDA_VISIBLE_DEVICES"])

and

import os
import jax
import jax.numpy as jnp

_ = jnp.arange(5)
print(jax.devices("cpu"), jax.devices("gpu"))

I use JAX CUDA on WSL regularly, but as @f0uriest mentioned, it is technically experimental. I suspect that the way we set the CUDA_VISIBLE_DEVICES is not working properly on your system.

[YigitElma]

@jiulong1306 any updates on this?

[dpanici]

Any updates?