Unable to Reuse GPU After JAX Execution

[NeyKoZv]

I'm trying to use both PyCUDA and JAX in the same Python script, but I'm running into a frustrating issue. Everything works fine when I:

-Run PyCUDA code first (works perfectly)
-Then run JAX code (also works fine)

But when I try to go back to PyCUDA after using JAX, it crashes saying the GPU is "busy or unavailable". It's like JAX won't let go of the GPU even after it's done.

Here is a small snippet of code I used to try and understand what was going on :

import numpy as np
from pycuda import driver as cuda, compiler
import jax.numpy as jnp

def run_pycuda():
    cuda.init()
    ctx = cuda.Device(0).make_context()
    
    # Simple CUDA multiplication kernel
    mod = compiler.SourceModule("""
    __global__ void test(float *x) { x[threadIdx.x] *= 2; }
    """)
    kernel = mod.get_function("test")
    
    arr = np.random.randn(100).astype(np.float32)
    kernel(cuda.InOut(arr), block=(100,1,1))
    
    ctx.detach()
    print("First PyCUDA run worked!")
def run_jax():
    x = jnp.array([1.0, 2.0, 3.0])

if __name__ == "__main__":
    run_pycuda()  # Works fine
    run_jax()     # Also works
    run_pycuda()  # Crashes here

jax version : 0.6.1 (Installed with pip install -U "jax[cuda12]")
Cuda version : 12.9 (nvidia-smi (version 575.51.03))

Any idea on how I can make jax release the GPU ?

Thank you in advance !

[dfm]

You might be seeing the fact that JAX uses asynchronous dispatch. Try wrapping run_jax using jax.block_until_ready to make the call blocking?

[NeyKoZv]

Unfortunatly adding .block_until_ready the the jax array or @jax.block_until_ready above the run_jax function does not release the gpu, I tried adding a time.sleep(3) after the function to see if some time would let jax release the gpu but it didn't work !

[NeyKoZv]

using :
XLA_PYTHON_CLIENT_PREALLOCATE=false
XLA_PYTHON_CLIENT_ALLOCATOR=platform
in the .env file made it work