First DDE version

diffrax/__init__.py

@@ -33,20 +33,28 @@
 )
 from ._global_interpolation import (
     AbstractGlobalInterpolation as AbstractGlobalInterpolation,
+)
+from ._global_interpolation import CubicInterpolation as CubicInterpolation
+from ._global_interpolation import DenseInterpolation as DenseInterpolation
+from ._global_interpolation import LinearInterpolation as LinearInterpolation
+from ._global_interpolation import (
     backward_hermite_coefficients as backward_hermite_coefficients,
-    CubicInterpolation as CubicInterpolation,
-    DenseInterpolation as DenseInterpolation,
-    linear_interpolation as linear_interpolation,
-    LinearInterpolation as LinearInterpolation,
+)
+from ._global_interpolation import linear_interpolation as linear_interpolation
+from ._global_interpolation import (
     rectilinear_interpolation as rectilinear_interpolation,
 )
 from ._integrate import diffeqsolve as diffeqsolve
 from ._local_interpolation import (
     AbstractLocalInterpolation as AbstractLocalInterpolation,
+)
+from ._local_interpolation import (
     FourthOrderPolynomialInterpolation as FourthOrderPolynomialInterpolation,
-    LocalLinearInterpolation as LocalLinearInterpolation,
-    ThirdOrderHermitePolynomialInterpolation as ThirdOrderHermitePolynomialInterpolation,  # noqa: E501
 )
+from ._local_interpolation import LocalLinearInterpolation as LocalLinearInterpolation
+from ._local_interpolation import (
+    ThirdOrderHermitePolynomialInterpolation as ThirdOrderHermitePolynomialInterpolation,
+)  # noqa: E501
 from ._misc import adjoint_rms_seminorm as adjoint_rms_seminorm
 from ._path import AbstractPath as AbstractPath
 from ._progress_meter import (
@@ -120,6 +128,8 @@
 )
 from ._step_size_controller import (
     AbstractAdaptiveStepSizeController as AbstractAdaptiveStepSizeController,
+)
+from ._step_size_controller import (
     AbstractStepSizeController as AbstractStepSizeController,
     ClipStepSizeController as ClipStepSizeController,
     ConstantStepSize as ConstantStepSize,
@@ -135,6 +145,18 @@
     UnderdampedLangevinDriftTerm as UnderdampedLangevinDriftTerm,
     WeaklyDiagonalControlTerm as WeaklyDiagonalControlTerm,
 )
+from ._step_size_controller import ConstantStepSize as ConstantStepSize
+from ._step_size_controller import PIDController as PIDController
+from ._step_size_controller import StepTo as StepTo
+from ._term import AbstractTerm as AbstractTerm
+from ._term import ControlTerm as ControlTerm
+from ._term import MultiTerm as MultiTerm
+from ._term import ODETerm as ODETerm
+from ._term import WeaklyDiagonalControlTerm as WeaklyDiagonalControlTerm
 
+from ._delays import Delays as Delays
+from ._delays import bind_history as bind_history
+from ._delays import history_extrapolation_implicit as history_extrapolation_implicit
+from ._delays import maybe_find_discontinuity as maybe_find_discontinuity
 
 __version__ = importlib.metadata.version("diffrax")

diffrax/_adjoint.py

@@ -127,6 +127,7 @@ def loop(
         solver,
         stepsize_controller,
         event,
+        delays,
         saveat,
         t0,
         t1,
@@ -137,6 +138,7 @@ def loop(
         passed_solver_state,
         passed_controller_state,
         progress_meter,
+        y0_history,
     ) -> Any:
         """Runs the main solve loop. Subclasses can override this to provide custom
         backpropagation behaviour; see for example the implementation of
@@ -578,6 +580,7 @@ def _loop_backsolve_bwd(
     solver,
     stepsize_controller,
     event,
+    delays,
     saveat,
     t0,
     t1,
@@ -586,6 +589,7 @@ def _loop_backsolve_bwd(
     throw,
     init_state,
     progress_meter,
+    y0_history,
 ):
     assert event is None
 
@@ -623,6 +627,7 @@ def _loop_backsolve_bwd(
         adjoint=self,
         solver=solver,
         stepsize_controller=stepsize_controller,
+        delays=delays,
         terms=adjoint_terms,
         dt0=None if dt0 is None else -dt0,
         max_steps=max_steps,
@@ -804,6 +809,7 @@ def loop(
         passed_solver_state,
         passed_controller_state,
         event,
+        delays,
         **kwargs,
     ):
         if jtu.tree_structure(saveat.subs, is_leaf=_is_subsaveat) != jtu.tree_structure(
@@ -849,6 +855,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())
@@ -863,6 +873,7 @@ def loop(
             init_state=init_state,
             solver=solver,
             event=event,
+            delays=delays,
             **kwargs,
         )
         final_state = _only_transpose_ys(final_state)