Add compilation cache support #8
Description
It would be extremely useful to add utils + examples for using jax's compilation_cache feature, which could really speed up experiments over time when using the same machine.
{{ message }}
It would be extremely useful to add utils + examples for using jax's compilation_cache feature, which could really speed up experiments over time when using the same machine.