Free JAX/XLA buffers of size exceeded threshold

jax-free.py

import jax


def collect(threshold=256 * 1024 ** 2):
    backend = jax.lib.xla_bridge.get_backend()
    freed = 0
    for buf in backend.live_buffers():
        if buf.nbytes >= threshold:
            buf += buf.nbytes
            buf.delete()
    return freed