Memory growing when integrating lots (different) functions

Hi again 😃

I noticed that memory seems to grow indefinitely when integrating lots of newly defined functions. In the following example, memory grows at each iteration if vector_field is defined inside solve_ode. I expect that compilation is triggered at each iteration in that case, and one could define vector_field globally to avoid that. But I also expected that the compilation cache for the locally defined function would be cleared at the end of solve_ode.

My use case here is benchmarking tons of different ODE models on lots of different data in a single script. While I am trying to take care to free the equinox modules containing the vector_fields, deleting those objects does not seem to clear that compilation cache, hence the system is running out of memory (that's at least my guess on what is happening).

I noticed the clear_cache() snipped in the test suite, but applying it here doesn't seem to free the right resources.

import gc, psutil, sys
import jax.numpy as jnp
import diffrax as dfx

def solve_ode():
  def vector_field(t, x, _): return -0.1*x
  t = jnp.arange(4800)
  sol = dfx.diffeqsolve(
    terms=dfx.ODETerm(vector_field),
    solver=dfx.Dopri5(),
    t0=t[0], t1=t[-1], dt0=t[1], y0=1., max_steps=len(t))
  return sol.ys

def clear_caches():
  process = psutil.Process()
  if process.memory_info().rss > 0.5 * 2**30: # >500MB 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()
    print("Cache cleared")

# loop that grows memory
for i in range(100):
  res = solve_ode()
  clear_caches()
  print(f"Process uses {psutil.Process().memory_info().rss / (1024 * 1024)} MB memory")

Start by running the provided Python reproduction and compare its behavior with the clear_cache() snippet mentioned from the test suite. Trace how locally defined vector_field functions and JAX compilation caches are retained during repeated diffeqsolve calls; done means repeated integrations no longer cause unbounded memory growth or the required cache-clearing behavior is documented.