Adding support for Events to ReversibleAdjoint

[jacob-hjortlund]

This might be a bit early, considering #593 hasn't been approved as of yet, but I'd like to discuss how to move forward with adding support for Events to the ReversibleAdjoint.

Currently, I am not sure why _loop_reversible_bwd does not support this. As far as I've understood, we'd need to ensure that the backwards pass integrates from t_event back to t_0. When we pass an event, then the current event-handling in loop ensures that reversible_save_index corresponds to the index of t_event.

I think what I'm missing here is that we would need to calculate gradients through the cond_fn of the Event as well, in the case that it depends on y / the parameters of the vector field. Is this correct, or am I on the wrong track here?

Would a path forward be something along the lines of saving the event_mask as part of the residual passed to _loop_reversible_bwd and jax.vjp through the cond_fn at the event to add the gradient contribution from the cond_fn?

[patrick-kidger]

I think I defer to @sammccallum on this one! WDYT? In terms of technical feasibility and (if it is technically reasonable) time/interest in the problem.

[sammccallum]

I think what I'm missing here is that we would need to calculate gradients through the cond_fn of the Event as well, in the case that it depends on y / the parameters of the vector field. Is this correct, or am I on the wrong track here?

This is probably the way we need to go but is very fiddly! That is, we work out what operations our final event y1 value is dependent on and pull back the gradients appropriately. Looking at the events code, this involves interpolation and root finding.

I imagine we would want to factor out this calculation of y1 into a calculate_event_y1 function which takes as input dense_info, and event. When backpropagating we would have something like eqx.filter_vjp(calculate_event_y1); eqx.filter_vjp(event). This is similar to how we handle interpolation currently:

https://github.com/sammccallum/diffrax/blob/e63dacc293d252d138af3ae0d985e2a5207caa92/diffrax/_adjoint.py#L1020

[patrick-kidger]

Okay!

I think 'very fiddly' might be synonymous with 'difficult to maintain', so on balance I'm inclined to suggest we don't try supporting this combination for now. :)

[jacob-hjortlund]

Thanks for the input, @sammccallum, glad to hear I at least was on the right track. I agree on it looking fiddly. I won't be pushing for this @patrick-kidger, based on our discussion in #602, so I'll go ahead and mark this as closed. :D