AbstractReversibleSolver + ReversibleAdjoint#603
Conversation
|
I've also added the |
patrick-kidger
left a comment
patrick-kidger
left a comment
There was a problem hiding this comment.
Okay, gosh, this one took far too long for me to get around. Thank you for your patience! If I can I'd like this to be the next big thing I focus on getting in to Diffrax.
| reversible_save_index + 1, tprev, reversible_ts | ||
| ) | ||
| reversible_save_index = reversible_save_index + jnp.where(keep_step, 1, 0) | ||
|
|
There was a problem hiding this comment.
A very minor bug here: if it just so happens that we run with t0 == t1 then we'll end up with reversible_ts = [t0 inf inf inf ...], which will not produce desired results in the backward solve.
We have a special branch to handle the saving in the t0 == t1 case, we should add a line handling the state.reversible_ts is not None case there.
| # Reversible info | ||
| if max_steps is None: | ||
| reversible_ts = None | ||
| reversible_save_index = None | ||
| else: | ||
| reversible_ts = jnp.full(max_steps + 1, jnp.inf, dtype=time_dtype) | ||
| reversible_save_index = 0 |
There was a problem hiding this comment.
I've thought of an alternative for this extra buffer, btw: ReversibleAdjoint.loop could intercept saveat and add an SubSaveAt(steps=True, save_fn=lambda t, y, args: None) to record the extra times. Then peel it off again when returning the final state.
I think that (a) might be doable without making any changes to _integrate.py and (b) would allow for also supporting SaveAt(steps=True). (As in that case we can just skip adding the extra SubSaveAt.) And (c) would avoid a few of the subtle issues I've commented on above about exactly which tprev/tnext-like value is actually being saved, because you can trust in the rest of the existing diffeqsolve to do that for you.
It's not a strong suggestion though.
There was a problem hiding this comment.
Yeah, this was the original idea I tried but I couldn't get around a leaked tracer error! I'm willing to give it another go if you start feeling strongly about it though ;)
There was a problem hiding this comment.
Maybe let's nail everything else down and then consider this. Reflecting on this, I do suspect it will make the code much easier to maintain in the long run.
sammccallum
force-pushed
the
AbstractReversibleSolver
branch
from
0cfd4ec to
3a26ac3
Compare
patrick-kidger
left a comment
patrick-kidger
left a comment
There was a problem hiding this comment.
Not sure if you were ready for a review on this yet, but I took a look over anyway 😁 We're making really good progress! In particular now that we're settled on just AbstractERK then I think all our complicated state-reconstruction concerns go away, so the chance of footgunning ourselves has gone way down 😁
|
Heads-up that I've just updated the base branch to (Unrelatedly, lmk when this branch is ready for review.) |
add reversible testing testing AbstractReversibleSolver + ReversibleAdjoint allow arbitrary interpolation unpacking over indexing jax while loop collapse saveat ValueErrors remove statonovich solver condition remove unused returns from AbstractReversibleSolver backward_step add test and remove messy benchmark add wrapped solver + tests made_jump=True for both solver steps improve docstrings AbstractSolver and docstring note about SDEs add AbstractReversibleSolver to public API newline in docstrings return RESULTS from reversible backward_step restrict Reversible to AbstractERK and check result in adjoint correct tprev and tnext of solver init switch to linear interpolation and y0,y1 dense_info name UReversible various doc formatting changes AbstractReversibleSolver check add disable_fsal property to AbstractRungeKutta and use in UReversible allow t0 != 0 Handle StepTo controller t0==t1 branch
sammccallum
force-pushed
the
AbstractReversibleSolver
branch
from
8d59058 to
ae01942
Compare
|
I think I've now addressed all of your comments, so it should be ready for review 👍 Understanding how to rebase through multiple merges was an experience but I believe that is correct now... |
| final_state, | ||
| (reversible_ts, reversible_save_index), | ||
| is_leaf=_is_none, | ||
| ) |
There was a problem hiding this comment.
I've left this t0==t1 branch separate from the above jax.lax.cond for readability but we can obviously combine if this is okay.
|
(No pressure to review anytime soon Patrick, I just marked this as "review requested" so it's easy for you to see across your sprawling jax empire) |
|
Would love to see this implemented, especially in light of https://arxiv.org/abs/2507.21006 |
|
@luke-a-thompson in the meantime if you're interested, Daniil and I implemented the EES methods here: https://github.com/sammccallum/diffrax/tree/EES |
|
This one hasn't been intentionally closed btw! Just autoclosed with the new release. I think this is the second time this has happened to you (#593!) but I still intend to get this in. I've recently found more time to review some of these larger feature PRs. So Optimistix is next on my list and after that I'm aiming for a new Diffrax release, including this PR :) |
3 participants
Re-opening #593.
Implements
AbstractReversibleSolverbase class andReversibleAdjointfor reversible back propagation.This updates
SemiImplicitEuler,LeapfrogMidpointandReversibleHeunto subclassAbstractReversibleSolver.Implementation
AbstractReversibleSolversubclassesAbstractSolverand adds abackward_stepmethod:This method should reconstruct
y0,solver_stateatt0fromy1,solver_stateatt1. See the aforementioned solvers for examples.When backpropagating,
ReversibleAdjointuses thisbackward_stepto reconstruct state. We then take avjpthrough a local forward step and accumulate gradients.ReversibleAdjointnow also pulls back gradients from any interpolated values, so we can useSaveAt(ts=...)!We allow arbitrary
solver_state(provided it can be reconstructed reversibly) and calculate gradients w.r.t.solver_state. Finally, we pull back these gradients ontoy0,args,termsusing thesolver.initmethod.