class AbstractAdjoint(eqx.Module):
    """Abstract base class for all adjoint methods."""

    @abc.abstractmethod
    def loop(
        self,
        *,
        args,
        terms,
        solver,
        stepsize_controller,
        saveat,
        t0,
        t1,
        dt0,
        max_steps,
        throw,
        init_state,
    ):
        """Runs the main solve loop. Subclasses can override this to provide custom
        backpropagation behaviour; see for example the implementation of
        [`diffrax.BacksolveAdjoint`][].
        """

    # Eurgh, delayed imports to handle circular dependencies.
    #
    # `integrate.py` defines the forward pass. `adjoint.py` defines the backward pass.
    # These pretty much necessarily depend on each other:
    #
    # - diffeqsolve needs to know about AbstractAdjoint, since it's one its arguments.
    # - BacksolveAdjoint needs to know about how to integrate a differential equation,
    #   since that's what it does.
    #
    # As such we get a circular dependency. We resolve it by lazily importing from
    # `integrate.py`. For convenience we make them available as properties here so all
    # adjoint methods can access these.
    @property
    def _loop_fn(self):
        from .integrate import loop

        return loop

    @property
    def _diffeqsolve(self):
        from .integrate import diffeqsolve

        return diffeqsolve class BacksolveAdjoint(AbstractAdjoint):
    """Backpropagate through [`diffrax.diffeqsolve`][] by solving the continuous
    adjoint equations backwards-in-time. This is also sometimes known as
    "optimise-then-discretise", the "continuous adjoint method" or simply the "adjoint
    method".

    This method implies very low memory usage, but is usually relatively slow, and the
    computed gradients will only be approximate. As such other methods are generally
    preferred unless exceeding memory is a concern.

    !!! note

        This was popularised by [this paper](https://arxiv.org/abs/1806.07366). For
        this reason it is sometimes erroneously believed to be a better method for
        backpropagation than the other choices available.

    !!! warning

        Using this method prevents computing forward-mode autoderivatives of
        [`diffrax.diffeqsolve`][]. (That is to say, `jax.jvp` will not work.)
    """

    kwargs: Dict[str, Any]

    def __init__(self, **kwargs):
        """
        **Arguments:**

        - `**kwargs`: The arguments for the [`diffrax.diffeqsolve`][] operations that
            are called on the backward pass. For example use
            ```python
            BacksolveAdjoint(solver=Dopri5())
            ```
            to specify a particular solver to use on the backward pass.
        """
        valid_keys = {
            "dt0",
            "solver",
            "stepsize_controller",
            "adjoint",
            "max_steps",
            "throw",
        }
        given_keys = set(kwargs.keys())
        diff_keys = given_keys - valid_keys
        if len(diff_keys):
            raise ValueError(
                f"The following keys are not valid for `BacksolveAdjoint`: {diff_keys}"
            )
        self.kwargs = kwargs

    def loop(self, *, args, terms, saveat, init_state, **kwargs):
        if saveat.steps or saveat.dense:
            raise NotImplementedError(
                "Cannot use `adjoint=BacksolveAdjoint()` with "
                "`saveat=Steps(steps=True)` or `saveat=Steps(dense=True)`."
            )

        y = init_state.y
        sentinel = object()
        init_state = eqx.tree_at(
            lambda s: jax.tree_leaves(s.y), init_state, replace_fn=lambda _: sentinel
        )

        final_state, aux_stats = _loop_backsolve(
            (y, args, terms), self=self, saveat=saveat, init_state=init_state, **kwargs
        )

        # We only allow backpropagation through `ys`; in particular not through
        # `solver_state` etc.
        ys = final_state.ys
        final_state = jax.tree_map(nondifferentiable_output, final_state)
        final_state = eqx.tree_at(
            lambda s: jax.tree_leaves(s.ys), final_state, jax.tree_leaves(ys)
        )

        return final_state, aux_stats class NoAdjoint(AbstractAdjoint):
    """Disable backpropagation through [`diffrax.diffeqsolve`][].

    Forward-mode autodifferentiation (`jax.jvp`) will continue to work as normal.

    If you do not need to differentiate the results of [`diffrax.diffeqsolve`][] then
    this may sometimes improve the speed at which the differential equation is solved.
    """

    def loop(self, *, throw, **kwargs):
        del throw
        final_state, aux_stats = self._loop_fn(**kwargs, is_bounded=False)
        final_state = jax.tree_map(nondifferentiable_output, final_state)
        return final_state, aux_stats class RecursiveCheckpointAdjoint(AbstractAdjoint):
    """Backpropagate through [`diffrax.diffeqsolve`][] by differentiating the numerical
    solution directly. This is sometimes known as "discretise-then-optimise", or
    described as "backpropagation through the solver".

    For most problems this is the preferred technique for backpropagating through a
    differential equation.

    A binomial checkpointing scheme is used so that memory usage is low.
    """

    def loop(self, *, throw, **kwargs):
        del throw
        return self._loop_fn(**kwargs, is_bounded=True) def _item_to_str(item: Union[str, type, slice]) -> str:
def _maybe_tuple_to_str(
    item: Union[str, type, slice, Tuple[Union[str, type, slice], ...]]
) -> str:
    def __class_getitem__(cls, item):
    def __class_getitem__(cls, item):
    def __class_getitem__(cls, item):
    def __class_getitem__(cls, item):
class _Array(Generic[_Annotation]):
class _PyTree(Generic[_Annotation]):
class Array:
        class X:
class PyTree:
        class X:
class Scalar:
class Array:
class PyTree: class DenseInterpolation(AbstractGlobalInterpolation):
    ts_size: Int
    infos: DenseInfos
    direction: Scalar
    interpolation_cls: Type[AbstractLocalInterpolation] = eqx.static_field()

    def __post_init__(self):
        def _check(_infos):
            assert _infos.shape[0] + 1 == self.ts.shape[0]

        jax.tree_map(_check, self.infos)

    # DenseInterpolations typically get `ts` and `infos` that are way longer than they
    # need to be, and padded with `nan`s. This means the normal way of measuring how
    # many entries we have - ts.shape[0] - won't be correct.
    def _ts_size(self):
        return self.ts_size

    def _get_local_interpolation(self, t: Scalar, left: bool):
        index, _ = self._interpret_t(t, left)
        prev_t = self.ts[index]
        next_t = self.ts[index + 1]
        infos = ω(self.infos)[index].ω
        return self.interpolation_cls(t0=prev_t, t1=next_t, **infos)

    @eqx.filter_jit
    def evaluate(
        self, t0: Scalar, t1: Optional[Scalar] = None, left: bool = True
    ) -> PyTree:
        if t1 is not None:
            return self.evaluate(t1, left=left) - self.evaluate(t0, left=left)
        t0 = t0 * self.direction
        # Passing `left` doesn't matter on a local interpolation, which is globally
        # continuous.
        return self._get_local_interpolation(t0, left).evaluate(t0)

    @eqx.filter_jit
    def derivative(self, t: Scalar, left: bool = True) -> PyTree:
        # Passing `left` doesn't matter on a local interpolation, which is globally
        # continuous.
        t = t * self.direction
        out = self._get_local_interpolation(t, left).derivative(t)
        return (self.direction * out**ω).ω

    @property
    def t0(self):
        return self.ts[0] * self.direction

    @property
    def t1(self):
        return self.ts[-1] * self.direction def is_sde(terms: PyTree[AbstractTerm]) -> bool:
    is_brownian = lambda x: isinstance(x, AbstractBrownianPath)
    leaves, _ = jax.tree_flatten(terms, is_leaf=is_brownian)
    return any(is_brownian(leaf) for leaf in leaves) def is_unsafe_sde(terms: PyTree[AbstractTerm]) -> bool:
    is_brownian = lambda x: isinstance(x, UnsafeBrownianPath)
    leaves, _ = jax.tree_flatten(terms, is_leaf=is_brownian)
    return any(is_brownian(leaf) for leaf in leaves) def bounded_while_loop(cond_fun, body_fun, init_val, max_steps, base=16):
    """Reverse-mode autodifferentiable while loop.

    Mostly as `lax.while_loop`, with a few small changes.

    Arguments:
        cond_fun: function `a -> a`
        body_fun: function `a -> b -> a`, where `b` is a function that should be used
            instead of performing in-place updates with .at[].set() etc; see below.
        init_val: pytree with structure `a`.
        max_steps: integer or `None`.
        base: integer.

    Limitations with in-place updates.:
        The single big limitation is around making in-place updates. Done naively then
        the XLA compiler will fail to treat these as in-place and will make a copy
        every time. (See JAX issue #8192.)

        Working around this is a bit of a hassle -- as follows -- and it is for this
        reason that `body_fun` takes a second argument.

        If you ever have:
        - an inplace update...
        - ...made to the input to the body_fun...
        - ...whose result is returned from the body_fun...
        ...then you should use

        ```python
        x = inplace(x).at[i].set(u)
        x = HadInplaceUpdate(x)
        ```

        in place of

        ```python
        x = x.at[i].set(u)
        ```

        where `inplace` is the second argument to `body_fun`, and `HadInplaceUpdate` is
        available at `diffrax.misc.HadInplaceUpdate`.

        Internally, `bounded_while_loop` will treat things so as to work around this
        limitation of XLA.

        !!! faq

            `HadInplaceUpdate` is available separately (instead of being returned
            automatically from `inplace().at[].set()`) in case the in-place update
            takes place inside e.g. a `lax.scan` or similar, and you need to maintain
            PyTree structures. Just place the `HadInplaceUpdate` at the very end of
            `body_fun`. (And applied only to those array(s) that actually had in-place
            update(s), if the state is a PyTree.)

        !!! note

            If you need to nest `bounded_while_loop`s, then the two `inplace` functions
            can be merged:

            ```python
            def body_fun(val, inplace):
                ...  # stuff (use inplace)

                def inner_body_fun(_val, _inplace):
                    _inplace = _inplace.merge(inplace)
                    ...  # stuff (use _inplace)

                bounded_while_loop(body_fun=inner_body_fun, ...)

                ... # stuff (use inplace)

            bounded_while_loop(body_fun=body_fun, ...)
            ```

        !!! note

            In-place updates to arrays that are _created_ inside of `body_fun` can be
            made as normal. It's just those arrays that are part of the state (that is
            passed in and out) that need to be treated specially.

    Note the extra `max_steps` argument. If this is `None` then `bounded_while_loop`
    will fall back to `lax.while_loop` (which is not reverse-mode autodifferentiable).
    If it is a non-negative integer then this is the maximum number of steps which may
    be taken in the loop, after which the loop will exit unconditionally.

    Note the extra `base` argument.
    - Run time will increase slightly as `base` increases.
    - Compilation time will decrease substantially as
      `math.ceil(math.log(max_steps, base))` decreases. (Which happens as `base`
      increases.)
    """

    init_val = jax.tree_map(jnp.asarray, init_val)

    if max_steps is None:

        def _make_update(_new_val):
            if isinstance(_new_val, HadInplaceUpdate):
                return _new_val.val
            else:
                return _new_val

        def _body_fun(_val):
            inplace = lambda x: x
            inplace.pred = True
            _new_val = body_fun(_val, inplace)
            return jax.tree_map(
                _make_update,
                _new_val,
                is_leaf=lambda x: isinstance(x, HadInplaceUpdate),
            )

        return lax.while_loop(cond_fun, _body_fun, init_val)

    if not isinstance(max_steps, int) or max_steps < 0:
        raise ValueError("max_steps must be a non-negative integer")
    if max_steps == 0:
        return init_val

    def _cond_fun(val, step):
        return cond_fun(val) & (step < max_steps)

    init_data = (cond_fun(init_val), init_val, 0)
    rounded_max_steps = base ** int(math.ceil(math.log(max_steps, base)))
    _, val, _ = _while_loop(_cond_fun, body_fun, init_data, rounded_max_steps, base)
    return val class HadInplaceUpdate(eqx.Module):
    val: Array def branched_error_if(
    pred: Union[_Bool, Callable[[], _Bool]],
    index: Int,
    msgs: Sequence[str],
    error_cls: Type[Exception] = ValueError,
) -> bool:
    def raises(_arg):
        _pred, _index = _arg
        if _pred:
            if isinstance(_index, jnp.ndarray):
                _index = _index.item()
            raise error_cls(msgs[_index])

    if callable(pred):
        with jax.ensure_compile_time_eval():
            pred = pred()

    if isinstance(pred, jnp.ndarray):
        with jax.ensure_compile_time_eval():
            pred = unvmap_any(pred)

    if isinstance(pred, jax.core.Tracer):
        hcb.call(raises, (pred, index))
    elif isinstance(pred, (bool, np.ndarray, jnp.ndarray)):
        raises((pred, index))
    else:
        msg = (
            "`pred` must either be a `bool`, a JAX array, or a zero-argument callable "
            "that returns a `bool` or JAX array, instead we got "
            f"value: {pred} of type: {type(pred)}."
        )
        assert False, msg def error_if(
    pred: Union[_Bool, Callable[[], _Bool]],
    msg: str,
    error_cls: Type[Exception] = ValueError,
) -> bool:
    """For use as part of validating inputs.
    Works even under JIT.

    Example:
        @jax.jit
        def f(x):
            error_if(x < 0, "x must be >= 0")

        f(jax.numpy.array(-1))
    """
    branched_error_if(pred, 0, [msg], error_cls) def unvmap_all(x):
    return _unvmap_all_p.bind(x) def unvmap_max(x):
    return _unvmap_max_p.bind(x) class SaveAt(eqx.Module):
    """Determines what to save as output from the differential equation solve.

    Instances of this class should be passed as the `saveat` argument of
    [`diffrax.diffeqsolve`][].
    """

    t0: bool = False
    t1: bool = False
    ts: Optional[Union[Sequence[Scalar], Array["times"]]] = None  # noqa: F821
    steps: bool = False
    dense: bool = False
    solver_state: bool = False
    controller_state: bool = False
    made_jump: bool = False

    def __post_init__(self):
        with jax.ensure_compile_time_eval():
            ts = None if self.ts is None else jnp.asarray(self.ts)
        object.__setattr__(self, "ts", ts)
        if (
            not self.t0
            and not self.t1
            and self.ts is None
            and not self.steps
            and not self.dense
        ):
            raise ValueError("Empty saveat -- nothing will be saved.") class RESULTS(metaclass=ContainerMeta):
    successful = ""
    max_steps_reached = (
        "The maximum number of solver steps was reached. Try increasing `max_steps`."
    )
    dt_min_reached = "The minimum step size was reached."
    nan_time = "NaN time encountered during timestepping."
    implicit_divergence = "Implicit method diverged."
    implicit_nonconvergence = (
        "Implicit method did not converge within the required number of iterations."
    ) class Solution(AbstractPath):
    """The solution to a differential equation.

    **Attributes:**

    - `t0`: The start of the interval that the differential equation was solved over.
    - `t1`: The end of the interval that the differential equation was solved over.
    - `ts`: Some ordered collection of times. Might be `None` if no values were saved.
        (i.e. just `diffeqsolve(..., saveat=SaveAt(dense=True))` is used.)
    - `ys`: The value of the solution at each of the times in `ts`. Might `None` if no
        values were saved.
    - `stats`: Statistics for the solve (number of steps etc.).
    - `result`: Integer specifying the success or cause of failure of the solve. A
        value of `0` corresponds to a successful solve. Any other value is a failure.
        A human-readable message can be obtained by looking up messages via
        `diffrax.RESULTS[<integer>]`.
    - `solver_state`: If saved, the final internal state of the numerical solver.
    - `controller_state`: If saved, the final internal state for the step size
        controller.
    - `made_jump`: If saved, the final internal state for the jump tracker.

    !!! note

        If `diffeqsolve(..., saveat=SaveAt(steps=True))` is set, then the `ts` and `ys`
        in the solution object will be padded with `NaN`s, out to the value of
        `max_steps` passed to [`diffrax.diffeqsolve`][].

        This is because JAX demands that shapes be known statically ahead-of-time. As
        we do not know how many steps we will take until the solve is performed, we
        must allocate enough space for the maximum possible number of steps.
    """

    t0: Scalar = field(init=True)
    t1: Scalar = field(init=True)  # override init=False in AbstractPath
    ts: Optional[Array["times"]]  # noqa: F821
    ys: Optional[PyTree["times", ...]]  # noqa: F821
    interpolation: Optional[DenseInterpolation]
    stats: Dict[str, Any]
    result: RESULTS
    solver_state: Optional[PyTree]
    controller_state: Optional[PyTree]
    made_jump: Optional[Bool]

    def evaluate(
        self, t0: Scalar, t1: Optional[Scalar] = None, left: bool = True
    ) -> PyTree:
        """If dense output was saved, then evaluate the solution at any point in the
        region of integration `self.t0` to `self.t1`.

        **Arguments:**

        - `t0`: The point to evaluate the solution at.
        - `t1`: If passed, then the increment from `t0` to `t1` is returned.
            (`=evaluate(t1) - evaluate(t0)`)
        - `left`: When evaluating at a jump in the solution, whether to return the
            left-limit or the right-limit at that point.
        """
        if self.interpolation is None:
            raise ValueError(
                "Dense solution has not been saved; pass saveat.dense=True."
            )
        return self.interpolation.evaluate(t0, t1, left)

    def derivative(self, t: Scalar, left: bool = True) -> PyTree:
        r"""If dense output was saved, then calculate an **approximation** to the
        derivative of the solution at any point in the region of integration `self.t0`
        to `self.t1`.

        That is, letting $y$ denote the solution over the interval `[t0, t1]`, then
        this calculates an approximation to $\frac{\mathrm{d}y}{\mathrm{d}t}$.

        (This is *not* backpropagating through the differential equation -- that
        typically corresponds to e.g. $\frac{\mathrm{d}y(t_1)}{\mathrm{d}y(t_0)}$.)

        !!! example

            For an ODE satisfying

            $\frac{\mathrm{d}y}{\mathrm{d}t} = f(t, y(t))$

            then this value is approximately equal to $f(t, y(t))$.

        !!! warning

            This value is generally not very accurate. Differential equation solvers
            are usually designed to produce splines whose value is close to the true
            solution; not to produce splines whose derivative is close to the
            derivative of the true solution.

            If you need accurate derivatives for the solution of an ODE, it is usually
            best to calculate `vector_field(t, sol.evaluate(t), args)`. That is, to
            pay the extra computational cost of another vector field evaluation, in
            order to get a more accurate value.

            Put precisely: this `derivative` method returns the *derivative of the
            numerical solution*, and *not* an approximation to the derivative of the
            true solution.

        **Arguments:**

        - `t`: The point to calculate the derivative of the solution at.
        - `left`: When evaluating at a jump in the solution, whether to return the
            left-limit or the right-limit at that point.
        """
        if self.interpolation is None:
            raise ValueError(
                "Dense solution has not been saved; pass saveat.dense=True."
            )
        return self.interpolation.derivative(t, left) class AbstractAdaptiveSDESolver(AbstractAdaptiveSolver):
    pass class AbstractItoSolver(AbstractSolver):
    pass class AbstractSolver(eqx.Module, metaclass=_MetaAbstractSolver):
    """Abstract base class for all differential equation solvers."""

    @property
    @abc.abstractmethod
    def term_structure(self) -> PyTreeDef:
        """What PyTree structure `terms` should have when used with this solver."""

    # On the type: frequently just Type[AbstractLocalInterpolation]
    @property
    @abc.abstractmethod
    def interpolation_cls(self) -> Callable[..., AbstractLocalInterpolation]:
        """How to interpolate the solution in between steps."""

    def order(self, terms: PyTree[AbstractTerm]) -> Optional[int]:
        """Order of the solver for solving ODEs."""
        return None

    def strong_order(self, terms: PyTree[AbstractTerm]) -> Optional[Scalar]:
        """Strong order of the solver for solving SDEs."""
        return None

    def error_order(self, terms: PyTree[AbstractTerm]) -> Optional[Scalar]:
        """Order of the error estimate used for adaptive stepping.

        The default (slightly heuristic) implementation is as follows.

        The error estimate is assumed to come from the difference of two methods. If
        these two methods have orders `p` and `q` then the local order of the error
        estimate is `min(p, q) + 1` for an ODE and `min(p, q) + 0.5` for an SDE.

        - In the SDE case then we assume `p == q == solver.strong_order()`.
        - In the ODE case then we assume `p == q + 1 == solver.order()`.
        - We assume that non-SDE/ODE cases do not arise.

        This is imperfect as these assumptions may not be true. In addition in the SDE
        case, then solvers will sometimes exhibit higher orders of convergence for
        specific noise types (see issue #47).
        """
        if is_sde(terms):
            order = self.strong_order(terms)
            if order is not None:
                order = order + 0.5
            return order
        else:
            return self.order(terms)

    def init(
        self,
        terms: PyTree[AbstractTerm],
        t0: Scalar,
        t1: Scalar,
        y0: PyTree,
        args: PyTree,
    ) -> _SolverState:
        """Initialises any hidden state for the solver.

        **Arguments** as [`diffrax.diffeqsolve`][].

        **Returns:**

        The initial solver state, which should be used the first time `step` is called.
        """
        return None

    @abc.abstractmethod
    def step(
        self,
        terms: PyTree[AbstractTerm],
        t0: Scalar,
        t1: Scalar,
        y0: PyTree,
        args: PyTree,
        solver_state: _SolverState,
        made_jump: Bool,
    ) -> Tuple[PyTree, Optional[PyTree], DenseInfo, _SolverState, RESULTS]:
        """Make a single step of the solver.

        Each step is made over the specified interval $[t_0, t_1]$.

        **Arguments:**

        - `terms`: The PyTree of terms representing the vector fields and controls.
        - `t0`: The start of the interval that the step is made over.
        - `t1`: The end of the interval that the step is made over.
        - `y0`: The current value of the solution at `t0`.
        - `args`: Any extra arguments passed to the vector field.
        - `solver_state`: Any evolving state for the solver itself, at `t0`.
        - `made_jump`: Whether there was a discontinuity in the vector field at `t0`.
            Some solvers (notably FSAL Runge--Kutta solvers) usually assume that there
            are no jumps and for efficiency re-use information between steps; this
            indicates that a jump has just occurred and this assumption is not true.

        **Returns:**

        A tuple of several objects:

        - The value of the solution at `t1`.
        - A local error estimate made during the step. (Used by adaptive step size
            controllers to change the step size.) May be `None` if no estimate was
            made.
        - Some dictionary of information that is passed to the solver's interpolation
            routine to calculate dense output. (Used with `SaveAt(ts=...)` or
            `SaveAt(dense=...)`.)
        - The value of the solver state at `t1`.
        - An integer (corresponding to `diffrax.RESULTS`) indicating whether the step
            happened successfully, or if (unusually) it failed for some reason.
        """

    def func_for_init(
        self, terms: PyTree[AbstractTerm], t0: Scalar, y0: PyTree, args: PyTree
    ) -> PyTree:
        """Provides vector field evaluations to select the initial step size.

        This is used to make a point evaluation. This is unlike
        [`diffrax.AbstractSolver.step`][], which operates over an interval.

        In general differential equation solvers are interval-based. There is precisely
        one place where point evaluations are needed: selecting the initial step size
        automatically in an ODE solve. And that is what this function is for.

        **Arguments:** As [`diffrax.diffeqsolve`][]

        **Returns:**

        The evaluation of the vector field at `t0`.
        """

        raise ValueError(
            "An initial step size cannot be selected automatically. The most common "
            "scenario for this error to occur is when trying to use adaptive step "
            "size solvers with SDEs. Please specify an initial `dt0` instead."
        ) class AbstractStratonovichSolver(AbstractSolver):
    pass class Euler(AbstractItoSolver, AbstractSolver):
    """Euler's method.

    1st order explicit Runge--Kutta method. Does not support adaptive step sizing.

    When used to solve SDEs, converges to the Itô solution.
    """

    term_structure = jax.tree_structure(0)
    interpolation_cls = LocalLinearInterpolation

    def order(self, terms):
        return 1

    def strong_order(self, terms):
        return 0.5

    def step(
        self,
        terms: AbstractTerm,
        t0: Scalar,
        t1: Scalar,
        y0: PyTree,
        args: PyTree,
        solver_state: _SolverState,
        made_jump: Bool,
    ) -> Tuple[PyTree, _ErrorEstimate, DenseInfo, _SolverState, RESULTS]:
        del solver_state, made_jump
        control = terms.contr(t0, t1)
        y1 = (y0**ω + terms.vf_prod(t0, y0, args, control) ** ω).ω
        dense_info = dict(y0=y0, y1=y1)
        return y1, None, dense_info, None, RESULTS.successful

    def func_for_init(
        self,
        terms: AbstractTerm,
        t0: Scalar,
        y0: PyTree,
        args: PyTree,
    ) -> PyTree:
        return terms.func_for_init(t0, y0, args) class AbstractAdaptiveStepSizeController(AbstractStepSizeController):
    # Default tolerances taken from scipy.integrate.solve_ivp
    rtol: Scalar = 1e-3
    atol: Scalar = 1e-6

    def wrap_solver(self, solver: AbstractSolver) -> AbstractSolver:
        # Poor man's multiple dispatch
        if isinstance(solver, AbstractImplicitSolver):
            if solver.nonlinear_solver.rtol is None:
                solver = eqx.tree_at(
                    lambda s: s.nonlinear_solver.rtol,
                    solver,
                    self.rtol,
                    is_leaf=lambda x: x is None,
                )
            if solver.nonlinear_solver.atol is None:
                solver = eqx.tree_at(
                    lambda s: s.nonlinear_solver.atol,
                    solver,
                    self.atol,
                    is_leaf=lambda x: x is None,
                )
        return solver class AbstractStepSizeController(eqx.Module):
    """Abstract base class for all step size controllers."""

    @abc.abstractmethod
    def wrap(self, direction: Scalar) -> "AbstractStepSizeController":
        """Remakes this step size controller, adding additional information.

        Most step size controllers can't be used without first calling `wrap` to give
        them the extra information they need.

        **Arguments:**

        - `direction`: Either 1 or -1, indicating whether the integration is going to
            be performed forwards-in-time or backwards-in-time respectively.

        **Returns:**

        A copy of the the step size controller, updated to reflect the additional
        information.
        """

    def wrap_solver(self, solver: AbstractSolver) -> AbstractSolver:
        """Remakes the solver, adding additional information.

        Some step size controllers need to modify the solver slightly. For example,
        adaptive step size controllers can automatically set the tolerances used in
        implicit solvers.

        **Arguments:**

        - `solver`: The solver to modify.

        **Returns:**

        The modified solver.
        """
        return solver

    @abc.abstractmethod
    def init(
        self,
        terms: PyTree[AbstractTerm],
        t0: Scalar,
        t1: Scalar,
        y0: PyTree,
        dt0: Optional[Scalar],
        args: PyTree,
        func_for_init: Callable[[Scalar, PyTree, PyTree], PyTree],
        error_order: Optional[Scalar],
    ) -> Tuple[Scalar, _ControllerState]:
        r"""Determines the size of the first step, and initialise any hidden state for
        the step size controller.

        **Arguments:** As `diffeqsolve`.

        - `func_for_init`: The value of `solver.func_for_init`.
        - `error_order`: The order of the error estimate. If solving an ODE this will
            typically be `solver.order()`. If solving an SDE this will typically be
            `solver.strong_order() + 0.5`.

        **Returns:**

        A 2-tuple of:

        - The endpoint $\tau$ for the initial first step: the first step will be made
            over the interval $[t_0, \tau]$. If `dt0` is specified (not `None`) then
            this is typically `t0 + dt0`. (Although in principle the step size
            controller doesn't have to respect this if it doesn't want to.)
        - The initial hidden state for the step size controller, which is used the
            first time `adapt_step_size` is called.
        """

    @abc.abstractmethod
    def adapt_step_size(
        self,
        t0: Scalar,
        t1: Scalar,
        y0: PyTree,
        y1_candidate: PyTree,
        args: PyTree,
        y_error: Optional[PyTree],
        error_order: Scalar,
        controller_state: _ControllerState,
    ) -> Tuple[Bool, Scalar, Scalar, Bool, _ControllerState, RESULTS]:
        """Determines whether to accept or reject the current step, and determines the
        step size to use on the next step.

        **Arguments:**

        - `t0`: The start of the interval that the current step was just made over.
        - `t1`: The end of the interval that the current step was just made over.
        - `y0`: The value of the solution at `t0`.
        - `y1_candidate`: The value of the solution at `t1`, as estimated by the main
            solver. Only a "candidate" as it is now up to the step size controller to
            accept or reject it.
        - `args`: Any extra arguments passed to the vector field; as
            [`diffrax.diffeqsolve`][].
        - `y_error`: An estimate of the local truncation error, as calculated by the
            main solver.
        - `error_order`: The order of `y_error`. For an ODE this is typically equal to
            `solver.order()`; for an SDE this is typically equal to
            `solver.strong_order() + 0.5`.
        - `controller_state`: Any evolving state for the step size controller itself,
            at `t0`.

        **Returns:**

        A tuple of several objects:

        - A boolean indicating whether the step was accepted/rejected.
        - The time at which the next step is to be started. (Typically equal to the
            argument `t1`, but not always -- if there was a vector field discontinuity
            there then it may be `nextafter(t1)` instead.)
        - The time at which the next step is to finish.
        - A boolean indicating whether a discontinuity in the vector field has just
            been passed. (Which for example some solvers use to recalculate their
            hidden state; in particular the FSAL property of some Runge--Kutta
            methods.)
        - The value of the step size controller state at `t1`.
        - An integer (corresponding to `diffrax.RESULTS`) indicating whether the step
            happened successfully, or if it failed for some reason. (e.g. hitting a
            minimum allowed step size in the solver.)
        """ class ConstantStepSize(AbstractStepSizeController):
    """Use a constant step size, equal to the `dt0` argument of
    [`diffrax.diffeqsolve`][].
    """

    compile_steps: Optional[bool] = False

    def wrap(self, direction: Scalar):
        return self

    def init(
        self,
        terms: PyTree[AbstractTerm],
        t0: Scalar,
        t1: Scalar,
        y0: PyTree,
        dt0: Optional[Scalar],
        args: PyTree,
        func_for_init: Callable[[Scalar, PyTree, PyTree], PyTree],
        error_order: Optional[Scalar],
    ) -> Tuple[Scalar, Scalar]:
        del terms, t1, y0, args, func_for_init, error_order
        if dt0 is None:
            raise ValueError(
                "Constant step size solvers cannot select step size automatically; "
                "please pass a value for `dt0`."
            )
        return t0 + dt0, dt0

    def adapt_step_size(
        self,
        t0: Scalar,
        t1: Scalar,
        y0: PyTree,
        y1_candidate: PyTree,
        args: PyTree,
        y_error: PyTree,
        error_order: Scalar,
        controller_state: Scalar,
    ) -> Tuple[bool, Scalar, Scalar, bool, Scalar, RESULTS]:
        del t0, y0, y1_candidate, args, y_error, error_order
        return (
            True,
            t1,
            t1 + controller_state,
            False,
            controller_state,
            RESULTS.successful,
        ) class StepTo(AbstractStepSizeController):
    """Make steps to just prespecified times."""

    ts: Union[Sequence[Scalar], Array["times"]]  # noqa: F821
    compile_steps: Optional[bool] = False

    def __post_init__(self):
        with jax.ensure_compile_time_eval():
            object.__setattr__(self, "ts", jnp.asarray(self.ts))
        if self.ts.ndim != 1:
            raise ValueError("`ts` must be one-dimensional.")
        if len(self.ts) < 2:
            raise ValueError("`ts` must have length at least 2.")

    def wrap(self, direction: Scalar):
        ts = self.ts * direction
        # Only tested after we've set the direction.
        error_if(
            ts[1:] <= ts[:-1],
            "`StepTo(ts=...)` must be strictly increasing (or strictly decreasing if "
            "t0 > t1).",
        )
        return type(self)(ts=ts, compile_steps=self.compile_steps)

    def init(
        self,
        terms: PyTree[AbstractTerm],
        t0: Scalar,
        t1: Scalar,
        y0: PyTree,
        dt0: None,
        args: PyTree,
        func_for_init: Callable[[Scalar, PyTree, PyTree], PyTree],
        error_order: Optional[Scalar],
    ) -> Tuple[Scalar, int]:
        del y0, args, func_for_init, error_order
        if dt0 is not None:
            raise ValueError(
                "`dt0` should be `None`. Step location is already determined "
                f"by {type(self).__name__}(ts=...).",
            )
        error_if(
            (t0 != self.ts[0]) | (t1 != self.ts[-1]),
            "Must have `t0==ts[0]` and `t1==ts[-1]`.",
        )
        return self.ts[1], 2

    def adapt_step_size(
        self,
        t0: Scalar,
        t1: Scalar,
        y0: Array["state"],  # noqa: F821
        y1_candidate: Array["state"],  # noqa: F821
        args: PyTree,
        y_error: Array["state"],  # noqa: F821
        error_order: Scalar,
        controller_state: int,
    ) -> Tuple[bool, Scalar, Scalar, bool, Int, RESULTS]:
        del t0, y0, y1_candidate, args, y_error, error_order
        return (
            True,
            t1,
            self.ts[controller_state],
            False,
            controller_state + 1,
            RESULTS.successful,
        ) class AbstractTerm(eqx.Module):
    r"""Abstract base class for all terms.

    Let $y$ solve some differential equation with vector field $f$ and control $x$.

    Let $y$ have PyTree structure $T$, let the output of the vector field have
    PyTree structure $S$, and let $x$ have PyTree structure $U$, Then
    $f : T \to S$ whilst the interaction $(f, x) \mapsto f \mathrm{d}x$ is a function
    $(S, U) \to T$.
    """

    @abc.abstractmethod
    def vf(self, t: Scalar, y: PyTree, args: PyTree) -> PyTree:
        """The vector field.

        Represents a function $f(t, y(t), args)$.

        **Arguments:**

        - `t`: the integration time.
        - `y`: the evolving state; a PyTree of structure $T$.
        - `args`: any static arguments as passed to [`diffrax.diffeqsolve`][].

        **Returns:**

        A PyTree of structure $S$.
        """
        pass

    @abc.abstractmethod
    def contr(self, t0: Scalar, t1: Scalar) -> PyTree:
        r"""The control.

        Represents the $\mathrm{d}t$ in an ODE, or the $\mathrm{d}w(t)$ in an SDE, etc.

        Most numerical ODE solvers work by making a step of length
        $\Delta t = t_1 - t_0$. Likewise most numerical SDE solvers work by sampling
        some Brownian motion $\Delta w \sim \mathcal{N}(0, t_1 - t_0)$.

        Correspondingly a control is *not* defined at a point. Instead it is defined
        over an interval $[t_0, t_1]$.

        **Arguments:**

        - `t0`: the start of the interval.
        - `t1`: the end of the interval.

        **Returns:**

        A PyTree of structure $U$. For a control $x$ then the result should
        represent $x(t_1) - x(t_0)$.
        """
        pass

    @abc.abstractmethod
    def prod(self, vf: PyTree, control: PyTree) -> PyTree:
        r"""Determines the interaction between vector field and control.

        With a solution $y$ to a differential equation with vector field $f$ and
        control $x$, this computes $f(t, y(t), args) \Delta x(t)$ given
        $f(t, y(t), args)$ and $\Delta x(t)$.

        **Arguments:**

        - `vf`: The vector field evaluation; a PyTree of structure $S$.
        - `control`: The control evaluated over an interval; a PyTree of structure $U$.

        **Returns:**

        The interaction between the vector field and control; a PyTree of structure
        $T$.

        !!! note

            This function must be bilinear.
        """
        pass

    def vf_prod(self, t: Scalar, y: PyTree, args: PyTree, control: PyTree) -> PyTree:
        r"""The composition of [`diffrax.AbstractTerm.vf`][] and
        [`diffrax.AbstractTerm.prod`][].

        With a solution $y$ to a differential equation with vector field $f$ and
        control $x$, this computes $f(t, y(t), args) \Delta x(t)$ given $t$, $y(t)$,
        $args$, and $\Delta x(t)$.

        Its default implementation is simply
        ```python
        self.prod(self.vf(t, y, args), control)
        ```

        This is offered as a special case that can be overridden when it is more
        efficient to do so.

        !!! example

            Consider when `vf` computes a matrix-matrix product, and `prod` computes a
            matrix-vector product. Then doing a naive composition corresponds to a
            (matrix-matrix)-vector product, which is less efficient than the
            corresponding matrix-(matrix-vector) product. Overriding this method offers
            a way to reclaim that efficiency.

        !!! example

            This is used extensively for efficiency when backpropagating via
            [`diffrax.BacksolveAdjoint`][].

        **Arguments:**

        - `t`: the integration time.
        - `y`: the evolving state; a PyTree of structure $T$.
        - `args`: any static arguments as passed to [`diffrax.diffeqsolve`][].
        - `control`: The control evaluated over an interval; a PyTree of structure $U$.

        **Returns:**

        A PyTree of structure $T$.

        !!! note

            This function must be linear in `control`.
        """
        return self.prod(self.vf(t, y, args), control)

    # This is a pinhole break in our vector-field/control abstraction.
    # Everywhere else we get to evaluate over some interval, which allows us to
    # evaluate our control over that interval. However to select the initial point in
    # an adaptive step size scheme, the standard heuristic is to start by making
    # evaluations at just the initial point -- no intervals involved.
    def func_for_init(self, t: Scalar, y: PyTree, args: PyTree) -> PyTree:
        """This is a special-cased version of [`diffrax.AbstractTerm.vf`][].

        If it so happens that the PyTree structures $T$ and $S$ are the same, then a
        subclass of `AbstractTerm` shoud set `func_for_init = vf`.

        This case is used when selecting the initial step size of an ODE solve
        automatically.

        See [`diffrax.AbstractSolver.func_for_init`][].
        """

        # Heuristic for whether it's safe to select an initial step automatically.
        vf = self.vf(t, y, args)
        flat_vf, tree_vf = jax.tree_flatten(vf)
        flat_y, tree_y = jax.tree_flatten(y)
        if tree_vf != tree_y or any(
            jnp.shape(x) != jnp.shape(y) for x, y in zip(flat_vf, flat_y)
        ):
            raise ValueError(
                "An initial step size cannot be selected automatically. The most "
                "common scenario for this error to occur is when trying to use "
                "adaptive step size solvers with SDEs, or with CDEs without "
                "`ControlTerm(...).to_ode()`. Please specify an initial `dt0` instead."
            )
        else:
            return vf

    def is_vf_expensive(
        self,
        t0: Scalar,
        t1: Scalar,
        y: Tuple[PyTree, PyTree, PyTree, PyTree],
        args: PyTree,
    ) -> bool:
        """Specifies whether evaluating the vector field is "expensive", in the
        specific sense that it is cheaper to evaluate `vf_prod` twice than `vf` once.

        Some solvers use this to change their behaviour, so as to act more efficiently.
        """
        return False class WrapTerm(AbstractTerm):
    term: AbstractTerm
    direction: Scalar

    def vf(self, t: Scalar, y: PyTree, args: PyTree) -> PyTree:
        t = t * self.direction
        return self.term.vf(t, y, args)

    def contr(self, t0: Scalar, t1: Scalar) -> PyTree:
        _t0 = jnp.where(self.direction == 1, t0, -t1)
        _t1 = jnp.where(self.direction == 1, t1, -t0)
        return (self.direction * self.term.contr(_t0, _t1) ** ω).ω

    def prod(self, vf: PyTree, control: PyTree) -> PyTree:
        return self.term.prod(vf, control)
import functools as ft
import warnings
import equinox as eqx
import jax
import jax.lax as lax
import jax.numpy as jnp
from typing import Optional
from .adjoint import (
    AbstractAdjoint,
    BacksolveAdjoint,
    NoAdjoint,
    RecursiveCheckpointAdjoint,
)
from .custom_types import Array, Bool, Int, PyTree, Scalar
from .global_interpolation import DenseInterpolation
from .heuristics import is_sde, is_unsafe_sde
from .misc import (
    bounded_while_loop,
    branched_error_if,
    error_if,
    HadInplaceUpdate,
    unvmap_all,
    unvmap_max,
)
from .saveat import SaveAt
from .solution import RESULTS, Solution
from .solver import (
    AbstractAdaptiveSDESolver,
    AbstractItoSolver,
    AbstractSolver,
    AbstractStratonovichSolver,
    Euler,
)
from .step_size_controller import (
    AbstractAdaptiveStepSizeController,
    AbstractStepSizeController,
    ConstantStepSize,
    StepTo,
)
from .term import AbstractTerm, WrapTerm

    These arguments are infrequently used, and for most purposes you shouldn't need to
    understand these. All of these are keyword-only arguments.

    - `adjoint`: How to backpropagate (and compute forward-mode autoderivatives) of
        `diffeqsolve`. Defaults to discretise-then-optimise with recursive
        checkpointing, which is usually the best option for most problems. See the page
        on [Adjoints](./adjoints.md) for more information.

    - `max_steps`: The maximum number of steps to take before quitting the computation
        unconditionally.

        Can also be set to `None` to allow an arbitrary number of steps, although this
        will disable backpropagation via discretise-then-optimise (backpropagation via
        optimise-then-discretise will still work), and also disables
        `saveat.steps=True` and `saveat.dense=True`.

        Note that (a) compile times; and (b) backpropagation run times; will increase
        as `max_steps` increases. (Specifically, each time `max_steps` passes a power
        of 16.) You can reduce these times by using the smallest value of `max_steps`
        that is reasonable for your problem.

    - `throw`: Whether to raise an exception if the integration fails for any reason.

        If `True` then an integration failure will either raise a `ValueError` (when
        not using `jax.jit`) or print a warning message (when using `jax.jit`).

        If `False` then the returned solution object will have a `result` field
        indicating whether any failures occurred.

        Possible failures include for example hitting `max_steps`, or the problem
        becoming too stiff to integrate. (For most purposes these failures are
        unusual.)

        !!! note

            Note that when `jax.vmap`-ing a differential equation solve, then
            `throw=True` means that an exception will be raised if any batch element
            fails. You may prefer to set `throw=False` and inspect the `result` field
            of the returned solution object, to determine which batch elements
            succeeded and which failed.

    - `solver_state`: Some initial state for the solver. Can be useful when for example
        using a reversible solver to recompute a solution. Generally obtained by
        `SaveAt(solver_state=True)`. It is unlikely you will need to use this option.

    - `controller_state`: Some initial state for the step size controller. Generally
        obtained by `SaveAt(controller_state=True)`. It is unlikely you will need to
        use this option.

    - `made_jump`: Whether a jump has just been made at `t0`. Used to update
        `solver_state` (if passed). It is unlikely you will need to use this option.

    **Returns:**

    Returns a [`diffrax.Solution`][] object specifying the solution to the differential
    equation.

    **Raises:**

    - `ValueError` for bad inputs.
    - `RuntimeError` if `throw=True` and the integration fails (e.g. hitting the
        maximum number of steps).

    !!! note

        It is possible to have `t1 < t0`, in which case integration proceeds backwards
        in time.
    """

    #
    # Initial set-up
    #

    # Error checking
    if dt0 is not None:
        msg = (
            "Must have (t1 - t0) * dt0 > 0, we instead got: "
            f"t1, with value: {t1} and type {type(t1)}, "
            f"t0, with value: {t0} and type {type(t0)} and, "
            f"dt0, with value: {dt0} and type {type(dt0)}"
        )
        error_if(lambda: (t1 - t0) * dt0 <= 0, msg)

    # Error checking
    term_leaves, term_structure = jax.tree_flatten(
        terms, is_leaf=lambda x: isinstance(x, AbstractTerm)
    )
    raises = False
    for leaf in term_leaves:
        if not isinstance(leaf, AbstractTerm):
            raises = True
        del leaf
    if term_structure != solver.term_structure:
        raises = True
    if raises:
        raise ValueError(
            "`terms` must be a PyTree of `AbstractTerms` (such as `ODETerm`), with "
            f"structure {solver.term_structure}"
        )
    del term_leaves, term_structure, raises

    if is_sde(terms):
        if not isinstance(solver, (AbstractItoSolver, AbstractStratonovichSolver)):
            warnings.warn(
                f"`{type(solver).__name__}` is not marked as converging to either the "
                "Itô or the Stratonovich solution."
            )
        if isinstance(adjoint, BacksolveAdjoint):
            if isinstance(solver, AbstractItoSolver):
                raise NotImplementedError(
                    f"`{solver.__name__}` converges to the Itô solution. However "
                    "`BacksolveAdjoint` currently only supports Stratonovich SDEs."
                )
            elif not isinstance(solver, AbstractStratonovichSolver):
                warnings.warn(
                    f"{solver.__name__} is not marked as converging to either the Itô "
                    "or the Stratonovich solution. Note that BacksolveAdjoint will "
                    "only produce the correct solution for Stratonovich SDEs."
                )

        if isinstance(stepsize_controller, AbstractAdaptiveStepSizeController):