`jax.clear_caches` after backpropagating through `diffeqsolve` causes segmentation fault

Hi Patrick. As discussed in #142, memory usage may grow as we integrate more functions because of caching on jax's part.

For this reason, I am trying to find a way to reliably clear cache every once in a while using jax.clear_caches. However, this seems to invariably lead to a segmentation fault if done after differentiating through diffeqsolve.

Possibly related to google/jax#14212 and google/jax#15973.

Any idea of what's going on here?

Minimal working example:

import equinox as eqx
import jax
import jax.numpy as jnp
from diffrax import ODETerm, Dopri5, diffeqsolve
from jaxtyping import Array, Float


class Function(eqx.Module):
    a: Float[Array, ""] = eqx.field(converter=jnp.asarray)

    def __call__(self, x: Float[Array, "*batch"]) -> Float[Array, "*batch"]:
        return - self.a * x


def solve(f: Function) -> Float[Array, ""]:
    return diffeqsolve(
        terms=ODETerm(lambda t, x, args: f(x)),
        solver=Dopri5(),
        t0=0.0,
        t1=1.0,
        dt0=0.1,
        y0=1.0,
    ).ys[0]


grad_f = jax.grad(solve)

f = Function(1.0)
grad_f(f)
jax.clear_caches()  # segmentation fault

Start by running the minimal Python example through solve, jax.grad, and jax.clear_caches to reproduce the segmentation fault. Investigate the interaction between diffrax's diffeqsolve differentiation path and JAX cache clearing, using the linked JAX issues as context. Done means cache clearing no longer crashes after backpropagating through diffeqsolve.