odeint_adjoint(func: Callable[..., object], y0: State, t: Tensor | Sequence[float], rtol: float = 1e-07, atol: float = 1e-09, method: str | ButcherTableau | None = None, options: dict[str, object] | None = None, event_fn: Callable[..., Tensor] | None = None, adjoint_rtol: float | None = None, adjoint_atol: float | None = None, adjoint_method: str | ButcherTableau | None = None, adjoint_options: dict[str, object] | None = None, adjoint_params: Sequence[Tensor] | None = None)Integrate dy/dt = f(t, y) and differentiate it at constant memory.
Same result as odeint, different gradient strategy. odeint
differentiates the discretisation — exact for what it computed, but every
stage of every step stays alive until backward. This instead discards the
forward graph and reconstructs the gradient by integrating the adjoint
system backwards, so memory no longer grows with the number of steps.
Use it when the solve is long or the tolerance tight enough that the
retained stages dominate memory. For a short solve, odeint is
both faster and more accurate.
Parameters
funccallablef(t, y) -> dy/dt.y0Tensort[0].tTensor or sequence of floatrtolfloat= 1e-07atolfloat= 1e-07None selects "dopri5".optionsdict or None= Noneevent_fncallable or None= Noneadjoint_rtolfloat or None= NoneNone inherits the forward
values.adjoint_atolfloat or None= NoneNone inherits the forward
values.None inherits method.adjoint_optionsdict or None= NoneNone inherits options.
{"norm": "seminorm"} drops the parameter-gradient block from the
backward solve's error norm, which is the one norm setting that has an
effect here; any other value, including a callable, is refused.None uses
func.parameters() when func exposes it, otherwise nothing.Returns
TensorShape (len(t), *y0.shape); index 0 is y0 itself.
Raises
NotImplementedErrorevent_fn is given.TypeErroradjoint_params is not a tensor.ValueErrorodeint rejects, on either solve.Notes
The gradient is an approximation. It solves the continuous adjoint
equation numerically, so it converges to the true derivative as the
tolerances tighten rather than matching odeint's gradient
exactly — tightening adjoint_rtol / adjoint_atol is what closes
the gap.
Cost is a second integration whose right-hand side performs one
vector-Jacobian product through func, so expect roughly a doubling
of solve time on top of the backward passes.
Gradients with respect to t are not produced. t is read to the
host once as an integration grid throughout lucid.diffeq, so it is
not differentiable in odeint either.
Examples
>>> import lucid, lucid.diffeq as diffeq
>>> k = lucid.tensor([0.5], dtype=lucid.float64, requires_grad=True)
>>> y0 = lucid.tensor([1.0], dtype=lucid.float64)
>>> ys = diffeq.odeint_adjoint(
... lambda t, y: -k * y, y0, [0.0, 1.0], adjoint_params=[k]
... )
>>> ys[-1].sum().backward()
>>> bool(abs(float(k.grad.item()) + 0.60653066) < 1e-5)
TrueSee Also
- lucid.diffeq.odeint—Direct differentiation; exact but memory-hungry.