Caching computed values or memoizing

[ASKabalan]

Hello,

What would be the best practice and most tracer friendly way to cache repeated computations in JAX

In jax_cosmo this is done by recording some values in a dictionary inside a pytree and reusing them .
The dictionary is part of the pytree class but is not a leaf

This is causing tracer leak issues in some cases (calling the same function from two places under the same jit)

any suggestions of a better way to do this?

[jakevdp]

If you want to cache values at runtime, probably io_callback would be the best bet.