Add mechanism to clear full compilation cache

There is currently no way to clear the full compilation class in JAX — something like f._clear_cache(), but which applies to all JIT-ted functions that have been compiled so far. This would be useful in situations where multiple high-memory-usage tasks need to be done sequentially. See discussion in https://github.com/google/jax/discussions/10826

[patrick-kidger]

Whilst we’re at it – JAX caches a lot of things (e.g. jaxprs) and I’ve found that this can also contribute to OOM on limited memory machines. (Such as GitHub Actions runners, for running tests.) Ways to clean these up would also be desirable.

FWIW my hugely hacky approach so far has been

def clear_caches():
    process = psutil.Process()
    if process.memory_info().vms > 4 * 2**30:  # >4GB memory usage
        for module_name, module in sys.modules.items():
            if module_name.startswith("jax"):
                for obj_name in dir(module):
                    obj = getattr(module, obj_name)
                    if hasattr(obj, "cache_clear"):
                        obj.cache_clear()
        gc.collect()

(which in the context of tests is wrapped in a @pytest.fixture(autouse=True).)

[simitii]

I think the need for manually cleaning the cache once in a while is not natural. It would be nice to have an internal garbage collector inside Jax that cleans the least recently used objects in the cache (depending on the memory usage by the cache at the time).