Patrick Kidger

I'd be happy to have a cross-link if think it's worth it! :) (Whilst we're here, I'm also conscious that our current rope implementation is kind of annoying to use,...

Ah awesome! We're making real progress here. Do we specifically need the jvp-of-jvp? That's a weird high-order construct so I'm willing to believe something is up there! (Given that jvp-of-linsolve...

Under-the-hood `eqx.filter_jit` splits its arguments into three pieces, that are then passed to a `jax.jit`-wrapped function. Those pieces are (a) donated arrays; (b) non-donated arrays; (c) static non-arrays. It may...

> (using either eqx.clear_caches() or jax.clear_caches(). The optionality here is a little surprising to me -- these clear totally unrelated caches, I think.

So I can see that you're creating a new `jvp_solve` function on every iteration. I think we keep strong references to the functions in the cache. FWIW https://github.com/jax-ml/jax/issues/16226 is potentially...

Aha awesome! That's excellent progress. I think I'd start by settling for a repro that doesn't use Equinox. Then we'll have something that we can send to the JAX/XLA team,...

Aha, interesting! So either way it looks like it's probably to do with keeping too much around in memory. So I think that gives us two clear action items: (a)...

Nice, I like this in the same way as #164. I think the same comments I have there apply here, otherwise this essentially LGTM. I'm happy to assume the presence...

@Joshuaalbert -- thanks for getting involved! To address your various points: 1. So I believe we actually experimented with this exact idea during development. (Certainly it's a standard one that...

> we could make a list of possible improvements and then do them in one fell PR Multiple smaller PRs are much easier to review :) > step_size isn't unitless...