Adjoints

An adjoint selects how derivatives flow through a solve, passed as adjoint= on ds.pdesolve. See sensitivities and autodiff.

Recommended

ForwardAdjoint is the default, and CheckpointAdjoint supports jax.grad and jax.vjp.

diff_soup.ForwardAdjoint

Forward-mode differentiation through a solve, the default of JNDF, JBDF, and JIDA.

It supports jax.jvp, jax.jacfwd, and forward over forward. Reverse mode needs CheckpointAdjoint.

Parameters:

Name	Type	Description	Default
linear_solver		Linear solver passed on to the solve, ignored by the built-in solvers.	required
composable		If True, reverse mode over the forward sensitivities works, as in jax.jacrev(jax.jacfwd(f)), and forward over forward is refused.	required
checkpoints		Number of states stored for reverse mode over the composable forward mode, chosen from max_steps when None.	required

diff_soup.CheckpointAdjoint

Reverse-mode differentiation through a solve by checkpointed replay.

Pass adjoint=CheckpointAdjoint() to use jax.grad or jax.vjp on a solve. It cannot be combined with SaveAt(steps=...).

Parameters:

Name	Type	Description	Default
checkpoints		Number of states the forward pass stores for the backward replay, chosen from max_steps when None.	required

---

Nonlinear solves

diff_soup.nonlinear.ImplicitAdjoint

Reverse-mode differentiation at a converged root, by the implicit function theorem.

It is the default adjoint of the nonlinear solvers, and DiffraxStepper maps it to diffrax.ImplicitAdjoint. JNDF, JBDF, and JIDA refuse it.

Parameters:

Name	Type	Description	Default
linear_solver		Linear solver passed on to the solve, ignored by the built-in solvers.	required

---

diff_soup.nonlinear.AbstractAdjoint

The details matter only when writing a strategy of your own.

diff_soup.nonlinear.AbstractAdjoint

Base class of the strategies that decide how derivatives flow through a solve.

A strategy passed to JNDF, JBDF, or JIDA must report is_forward_only or uses_online_checkpoints as True, or the solve raises TypeError.

Attributes:

Name	Type	Description
ift_mode	Optional[str]	Direction of the implicit-function derivative, "vjp" for reverse or "jvp" for forward, or None when the strategy has none.

ift_mode class-attribute

ift_mode: Optional[str] = None

compositions

compositions(lowered_problem: Any) -> frozenset[str]

Return the derivative compositions this strategy supports on a lowered problem.

The names are "forward", "forward-over-forward", "reverse", and "forward-over-reverse". The base class supports none.

Parameters:

Name	Type	Description	Default
lowered_problem	Any	The lowered problem the strategy will differentiate.	required

Returns:

Type	Description
frozenset[str]	A frozenset of composition names.

apply

apply(primal_fn: Callable[..., PyTree], rewrite_fn: Optional[Callable[..., Any]], inputs: PyTree, tags: frozenset) -> PyTree

Run primal_fn on inputs under this strategy.

The base class calls primal_fn(inputs). ImplicitAdjoint and ForwardAdjoint also pass ift_mode, linear_solver, and tags as keywords.

Parameters:

Name	Type	Description	Default
primal_fn	Callable[..., PyTree]	Function computing the result from inputs.	required
rewrite_fn	Optional[Callable[..., Any]]	Ignored by the built-in strategies.	required
inputs	PyTree	Pytree handed to primal_fn.	required
tags	frozenset	Frozenset of tags that the implicit strategies pass on.	required

Returns:

Type	Description
PyTree	The value primal_fn returns.