Patrick Kidger

In terms of the implementation, I like this example! It's short and easy to follow. (Nit: `jr.PRNGKey -> jr.key`; the former has been replaced with the latter.) Also tagging @frederic-bender...

So (a) I think you've made a few mistakes in the benchmarking, and (b) most Lineax/Optimistix/Diffrax routines all finish with an option to throw a runtime error if things have...

> I am surprised that the vmap is not triggering multiple internal init? `init` is called only on the non-vmap'd input A, so it won't be vmap'd.

This sounds like an open research question :)

Perhaps have a wrapper at the `jax.grad` level, counting the number of arrays that cross its boundary. This would seem to be the most reliable solution I think.

In theory Diffrax should work under parallelism! In practice it's not a very common combination, so perhaps something has gone wrong! My first usual thing to be suspicious of is...

As for saving results: probably yes by using a `jax.pure_callback`. if you want to do that as you go along then you could pass it to `SaveAt(fn=...)`. You'll need to...

Ah, this is actually an interesting point! I've been aware of this wart but I'm not completely sure how to fix it. Pretty much everywhere else we do actually compute...

I think the bullet is just figuring out how to implement it! I agree that a linear solve probably isn't that expensive a lot of the time. (At the very...

So this is something which [Haliax](https://github.com/stanford-crfm/haliax), a downstream library of Equinox, explicitly adds. I'd be open to open suggestions on how to accomplish something similar within Equinox, although I'm also...