diffrax icon indicating copy to clipboard operation
diffrax copied to clipboard

Support Events for BacksolveAdjoint

Open LuggiStruggi opened this issue 2 months ago • 2 comments

Hi! :)

I just wanted to ask if its possible to also support Events for the BacksolveAdjoint. I know that its advised against using BacksolveAdjoint by this library in general but it would be nontheless interesting for me, regarding the memory efficiency as well as a baseline.

I guess in _adjoint.py, _loop_backsolve_fwd would need to save the event_mask (or just a flag if an event happend in the end or not) as one of the residuals and then in _loop_backsolve_bwd if the event happend add the correction term to the backward dynamics based on gradients of the cond_fn or smth like that?

I might be overlooking something that is obvious, why it is not yet implemented. If you think its possible i could try to implement it.

LuggiStruggi avatar Oct 24 '25 11:10 LuggiStruggi

Hey there! The main reason it's not implemented is the work required to review (not just to create) a PR for adding support; plus the maintenance cost. So I'm afraid that I don't think I can commit to adding this feature.

That said, this shouldn't meaningfully stand in the way of you adding this feature just for you to use yourself: copy-paste BacksolveAdjoint and add your own support. You can then call diffeqsolve(..., adjoint=LuggisBacksolveAdjoint()) and it should just work :)

patrick-kidger avatar Oct 24 '25 20:10 patrick-kidger

I see. Yes that's understandable. In that case I'll try to implement it on my side :) Thanks for always replying so quickly!

LuggiStruggi avatar Oct 25 '25 08:10 LuggiStruggi