patrick-kidger / patrick-kidger/diffrax

Support Events for BacksolveAdjoint

Open
#698 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

question
Dominant language
Python
Stars
2.1k
Forks
189
Avg merge
3d 18h
Merged PRs (30d)
1

Description

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.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reading _adjoint.py, especially _loop_backsolve_fwd and _loop_backsolve_bwd, to understand how residuals and event handling currently work. Determine whether BacksolveAdjoint can support Events and what correction terms are required; done means the behavior is implemented and verified for event-triggered backward solves.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.