In progress commit on branch delay.

Author: patrick-kidger

diffrax/__init__.py

@@ -7,6 +7,7 @@
 )
 from .autocitation import citation, citation_rules
 from .brownian import AbstractBrownianPath, UnsafeBrownianPath, VirtualBrownianTree
+from .delays import Delays
 from .event import (
     AbstractDiscreteTerminatingEvent,
     DiscreteTerminatingEvent,

diffrax/adjoint.py

@@ -117,6 +117,7 @@ def loop(
         solver,
         stepsize_controller,
         discrete_terminating_event,
+        delays,
         saveat,
         t0,
         t1,
@@ -522,13 +523,15 @@ def _loop_backsolve_bwd(
     solver,
     stepsize_controller,
     discrete_terminating_event,
+    delays,
     saveat,
     t0,
     t1,
     dt0,
     max_steps,
     throw,
     init_state,
+    y0_history,
 ):
     assert discrete_terminating_event is None
 
@@ -566,6 +569,8 @@ def _loop_backsolve_bwd(
         adjoint=self,
         solver=solver,
         stepsize_controller=stepsize_controller,
+        discrete_terminating_event=discrete_terminating_event,
+        delays=delays,
         terms=adjoint_terms,
         dt0=None if dt0 is None else -dt0,
         max_steps=max_steps,
@@ -745,6 +750,7 @@ def loop(
         passed_solver_state,
         passed_controller_state,
         discrete_terminating_event,
+        delays,
         **kwargs,
     ):
         if jtu.tree_structure(saveat.subs, is_leaf=_is_subsaveat) != jtu.tree_structure(
@@ -790,6 +796,10 @@ def loop(
             raise NotImplementedError(
                 "`diffrax.BacksolveAdjoint` is not compatible with events."
             )
+        if delays is not None:
+            raise NotImplementedError(
+                "Cannot use `delays` with `adjoint=BacksolveAdjoint()`"
+            )
 
         y = init_state.y
         init_state = eqx.tree_at(lambda s: s.y, init_state, object())
@@ -804,6 +814,7 @@ def loop(
             init_state=init_state,
             solver=solver,
             discrete_terminating_event=discrete_terminating_event,
+            delays=delays,
             **kwargs,
         )
         final_state = _only_transpose_ys(final_state)