Clear GPU memory

Dear jax team,

I’d like to use jax alongside other tools running on GPU in the same pipeline. Is there a possibility to “encapsulate” the usage of jax/XLA so that the GPU is freed afterwards? Even if I would have to copy over the DeviceArrays into numpy manually.

Maybe something like:

with jax.Block():
    result = some_jitted_fun(a, b, c)
    result = onp.copy(result)

I can imagine the (design of) handling of objects and their GPU memory is not straightforward, if not practically impossible. Could I at least tell jax to use the GPU only incrementally instead of filling the memory completely on import?

[altosaar]

Any update on this?

[skye]

If you’d like to delete all on-device buffers, you can now reach directly into the underlying C++ internals like this:

backend = jax.lib.xla_bridge.get_backend()
for buf in backend.live_buffers(): buf.delete()

You’ll get errors about trying to use a deleted buffer if you try to use an existing references to these buffers. Note these APIs are still subject to change as well!