Patrick Kidger

'define large batch size as 2' sounds reasonable to me! I think we should probably avoid fully unrolling, for specifically the reasons of compile time. I try to do a...

I don't think there should be any complications. (It's totally fine to have 'primitives inside primitives' like this.) If anything it should be easier, as there is no need to...

I agree! This sounds like the performance gain here is in the margins + would increase maintenace cost + would open up a potential footgun for users. I'm inclined not...

Lineax's JVP rule is defined [here](https://github.com/patrick-kidger/lineax/blob/52f9d3f78e153125e76e35f508d9ad97e3013e70/lineax/_solve.py#L131). We actually do define a custom primitive, rather than using `jax.custom_jvp`, because we need to define a custom transposition rule. :) Fortunately for your...

So the only thing that can be passed around inside of JAX, at runtime, are JAX arrays. After all, the whole thing gets lowered to an XLA computation graph! So...

I think handling batching is a separate concern. It looks like you need to both (a) have a way for multiple `pure_callback`s to communicate with each other (as in my...

Ah, I've been down this road before (on other projects). The problem with this is when a newly-released version of (in this case) JAX really is compatible, that this new...

Haha, I mean specifically in issues that *only* occur downstream -- from some intersection of cases not covered by existing tests. This is the same relationship Diffrax etc have with...

Hmm, throwing away the preconditioner is probably a mistake. Happy to take a PR on that! (And ideally a test 😄) As for your second comment, I'm away from my...

> So GMRES takes 1 addiional iteration, which I think is because of this where the first iteration of GMRES is actually a no-op, not sure why this is? I...