jax

7fedc0e5 - Introduce `jax.extend.clear_in_memory_compilation_cache` to allow users to clear the LRU cache for mesh executables during compilation