t0==t1 branch

Author: sammccallum

diffrax/_adjoint.py

@@ -1101,7 +1101,6 @@ def cond_fun(state):
 
     state = jax.lax.while_loop(cond_fun, grad_step, state)
     _, _, y0, _, grad_y0, grad_state, grad_args, grad_terms = state
-    jax.debug.print("{}", y0)
 
     # Pull solver_state gradients back onto y0, args, terms.
 

diffrax/_integrate.py

@@ -870,6 +870,34 @@ def _make_ys(out, old_outs):
         lambda s: s.save_state, final_state, save_state, is_leaf=_is_none
     )
 
+    # if t0 == t1 and we are using diffrax.ReversibleAdjoint then we need to update the
+    # reversible_ts and reversible_ts_index to get correct gradients
+    def _reversible_info_if_t0_equals_t1(reversible_ts, reversible_save_index):
+        reversible_ts = eqxi.buffer_at_set(final_state.reversible_ts, 1, t0)
+        reversible_save_index += 1
+        return reversible_ts, reversible_save_index
+
+    reversible_ts, reversible_save_index = jax.lax.cond(
+        eqxi.unvmap_any(t0 == t1),
+        lambda __ts, __index: jax.lax.cond(
+            t0 == t1,
+            lambda _ts, _index: _reversible_info_if_t0_equals_t1(_ts, _index),
+            lambda _ts, _index: (_ts, _index),
+            __ts,
+            __index,
+        ),
+        lambda __ts, __index: (__ts, __index),
+        final_state.reversible_ts,
+        final_state.reversible_save_index,
+    )
+
+    final_state = eqx.tree_at(
+        lambda s: (s.reversible_ts, s.reversible_save_index),
+        final_state,
+        (reversible_ts, reversible_save_index),
+        is_leaf=_is_none,
+    )
+
     final_state = _handle_static(final_state)
     result = RESULTS.where(cond_fun(final_state), RESULTS.max_steps_reached, result)
     aux_stats = dict()  # TODO: put something in here?

test/test_reversible.py

@@ -25,19 +25,32 @@ def __call__(self, t, y, args):
 
 
 @eqx.filter_value_and_grad
-def _loss(y0__args__term, solver, saveat, adjoint, stepsize_controller, dual_y0):
+def _loss(
+    y0__args__term,
+    solver,
+    saveat,
+    adjoint,
+    stepsize_controller,
+    dual_y0,
+    t0_equals_t1,
+):
     y0, args, term = y0__args__term
 
     if isinstance(stepsize_controller, diffrax.StepTo):
         dt0 = None
     else:
         dt0 = 0.01
 
+    if t0_equals_t1:
+        t1 = 0
+    else:
+        t1 = 5
+
     sol = diffrax.diffeqsolve(
         term,
         solver,
         t0=0,
-        t1=5,
+        t1=t1,
         dt0=dt0,
         y0=y0,
         args=args,
@@ -54,7 +67,13 @@ def _loss(y0__args__term, solver, saveat, adjoint, stepsize_controller, dual_y0)
 
 
 def _compare_grads(
-    y0__args__term, base_solver, solver, saveat, stepsize_controller, dual_y0
+    y0__args__term,
+    base_solver,
+    solver,
+    saveat,
+    stepsize_controller,
+    dual_y0=False,
+    t0_equals_t1=False,
 ):
     loss, grads_base = _loss(
         y0__args__term,
@@ -63,6 +82,7 @@ def _compare_grads(
         adjoint=diffrax.RecursiveCheckpointAdjoint(),
         stepsize_controller=stepsize_controller,
         dual_y0=dual_y0,
+        t0_equals_t1=t0_equals_t1,
     )
     loss, grads_reversible = _loss(
         y0__args__term,
@@ -71,6 +91,7 @@ def _compare_grads(
         adjoint=diffrax.ReversibleAdjoint(),
         stepsize_controller=stepsize_controller,
         dual_y0=dual_y0,
+        t0_equals_t1=t0_equals_t1,
     )
     assert tree_allclose(grads_base, grads_reversible, atol=1e-5)
 
@@ -130,9 +151,7 @@ def test_reversible_heun_ode(stepsize_controller, saveat):
     args = jnp.array([0.5])
     solver = diffrax.ReversibleHeun()
 
-    _compare_grads(
-        (y0, args, terms), solver, solver, saveat, stepsize_controller, dual_y0=False
-    )
+    _compare_grads((y0, args, terms), solver, solver, saveat, stepsize_controller)
 
 
 @pytest.mark.parametrize(
@@ -161,9 +180,7 @@ def test_reversible_heun_sde(stepsize_controller, saveat):
     args = jnp.array([0.5])
     solver = diffrax.ReversibleHeun()
 
-    _compare_grads(
-        (y0, args, terms), solver, solver, saveat, stepsize_controller, dual_y0=False
-    )
+    _compare_grads((y0, args, terms), solver, solver, saveat, stepsize_controller)
 
 
 @pytest.mark.parametrize(
@@ -189,9 +206,7 @@ def test_leapfrog_midpoint(stepsize_controller, saveat):
     args = jnp.array([0.5])
     solver = diffrax.LeapfrogMidpoint()
 
-    _compare_grads(
-        (y0, args, terms), solver, solver, saveat, stepsize_controller, dual_y0=False
-    )
+    _compare_grads((y0, args, terms), solver, solver, saveat, stepsize_controller)
 
 
 @pytest.mark.parametrize(
@@ -226,14 +241,7 @@ def test_reversible_explicit(stepsize_controller, saveat):
     if saveat.subs.ts is not None:
         base_solver = solver
 
-    _compare_grads(
-        (y0, args, terms),
-        base_solver,
-        solver,
-        saveat,
-        stepsize_controller,
-        dual_y0=False,
-    )
+    _compare_grads((y0, args, terms), base_solver, solver, saveat, stepsize_controller)
 
 
 @pytest.mark.parametrize(
@@ -270,11 +278,33 @@ def test_reversible_sde(stepsize_controller, saveat):
     if saveat.subs.ts is not None:
         base_solver = solver
 
+    _compare_grads((y0, args, terms), base_solver, solver, saveat, stepsize_controller)
+
+
+@pytest.mark.parametrize(
+    "saveat",
+    [
+        diffrax.SaveAt(t0=True),
+        diffrax.SaveAt(t1=True),
+        diffrax.SaveAt(t0=True, t1=True),
+    ],
+)
+def test_reversible_t0_equals_t1(saveat):
+    n = 10
+    y0 = jnp.linspace(1, 10, num=n)
+    key = jr.PRNGKey(10)
+    f = VectorField(n, n, n, depth=4, key=key)
+    terms = diffrax.ODETerm(f)
+    args = jnp.array([0.5])
+    base_solver = diffrax.Tsit5()
+    solver = diffrax.UReversible(base_solver)
+    stepsize_controller = diffrax.ConstantStepSize()
+
     _compare_grads(
         (y0, args, terms),
         base_solver,
         solver,
         saveat,
         stepsize_controller,
-        dual_y0=False,
+        t0_equals_t1=True,
     )