Solvers

The solver API computes an equilibrium state

\[ z^\star = f_\theta(z^\star) \]

by iterating on the residual

\[ r(z)=f_\theta(z)-z. \]

The package exposes one configuration object:

from silva_networks import SolverConfig

config = SolverConfig(
    solver="anderson",
    max_iter=25,
    tol=1e-6,
    alpha=0.5,
    history=5,
    stop_mode="relative",
    anderson_batch_dims=1,
    return_best=True,
    indexing=(10, 20),
    backward_mode="implicit",
    backward_solver="gmres",
    backward_stop_mode="relative",
)

fixed_point is the numerical forward-solver dispatcher. Package layers and presets call solve_equilibrium, which uses the same forward solver and then chooses the training rule from SolverConfig.backward_mode.

Picard Iteration

The damped Picard update, interpreted through the classical contraction result [41], starts with an initial state \(z_0\). At iteration \(k\), evaluate the transition

\[ \tilde z_{k+1}=f_\theta(z_k). \]

The residual is

\[ r_k=\tilde z_{k+1}-z_k. \]

Damping blends the old state with the proposed state:

\[ z_{k+1} = (1-\alpha)z_k+\alpha \tilde z_{k+1} = z_k+\alpha r_k. \]

When \(\alpha=1\), this is the classical Picard update. Smaller \(\alpha\) can help when the map is nearly non-contractive.

Anderson Acceleration

Anderson acceleration [10] [11] stores recent residuals

\[ r_i=f_\theta(z_i)-z_i. \]

For the last \(m\) states, form

\[ G=\begin{bmatrix} r_{k-m+1} & \cdots & r_k \end{bmatrix}. \]

The coefficients solve a constrained least-squares problem:

\[ \min_c \|Gc\|_2^2+\lambda\|c\|_2^2 \quad\text{subject to}\quad \mathbf 1^\top c=1. \]

The KKT system used in the implementation is

\[ \begin{bmatrix} G^\top G+\lambda I & \mathbf 1 \\ \mathbf 1^\top & 0 \end{bmatrix} \begin{bmatrix} c \\ \nu \end{bmatrix} = \begin{bmatrix} 0 \\ 1 \end{bmatrix}. \]

The next state is

\[ z_{k+1} = \beta\sum_i c_i f(z_i) +(1-\beta)\sum_i c_i z_i. \]

anderson_batch_dims=1 solves the Anderson coefficient system independently for each leading batch sample, and convergence uses the worst sample residual. Packed coupled states use anderson_batch_dims=0.

For trainable modules, Anderson history is kept detached to control memory. SolverConfig(reengage=True) lets package layers evaluate one final differentiable transition after the accelerated numerical solve, so solver="anderson" remains a usable training option.

Broyden

Broyden's method [12] treats the equilibrium condition as a root-finding problem:

\[ F(z)=f_\theta(z)-z=0. \]

If \(B_k\) approximates the inverse Jacobian \(J_F(z_k)^{-1}\), the Newton-like step is

\[ s_k=-\alpha B_k F(z_k), \qquad z_{k+1}=z_k+s_k. \]

Let

\[ y_k=F(z_{k+1})-F(z_k). \]

The good-Broyden inverse update is

\[ B_{k+1} =B_k+\frac{(s_k-B_k y_k)(s_k^\top B_k)}{s_k^\top B_k y_k}. \]

The implementation stores low-rank inverse updates rather than an (n\times n) matrix. history bounds the retained rank, so Broyden can be selected for sequence, image, and coupled states without allocating a dense Jacobian-sized matrix. The approximation restarts from the initial inverse when that rank is full, then incorporates the newest secant pair.

The returned SolverResult.inverse_estimate is a BroydenInverseEstimate. Its rank, left_factors, and right_factors expose the retained numerical state, while apply_residual_inverse, apply_residual_inverse_transpose, and apply_fixed_point_adjoint_inverse apply it without materializing a dense matrix. This is the shared object used by SHINE.

GMRES for Adjoint Systems

The SILVA study's implicit-gradient diagnostic uses the matrix-free GMRES method [13] for a linear adjoint solve. Around an equilibrium, let

\[ J_f(z^\star)=\frac{\partial f}{\partial z}(z^\star). \]

The standard DEQ adjoint vector \(u\) solves

\[ (I-J_f(z^\star)^\top)u=g, \]

where \(g=\partial \mathcal L/\partial z^\star\). The package exposes a matrix-free GMRES helper:

from silva_networks import gmres

result = gmres(lambda v: A(v), b, max_iter=40, tol=1e-6)
u = result.x

For damped update diagnostics,

\[ T_\alpha(z)=(1-\alpha)z+\alpha f(z), \]

the helper implicit_adjoint_solve solves

\[ (I-J_{T_\alpha}(z^\star)^\top)u=g. \]

This is useful for reproducing the local linear analysis and GMRES-style diagnostic experiments.

Backward Modes

SolverConfig(backward_mode="unrolled") is the default. It differentiates through the finite solver computation, with reengage=True giving Anderson training a final differentiable transition after the detached accelerated history.

SolverConfig(backward_mode="implicit") runs the forward fixed-point solve without recording the solver history, then uses GMRES in the backward pass:

from silva_networks import SolverConfig

config = SolverConfig(
    solver="anderson",
    alpha=0.5,
    max_iter=40,
    backward_mode="implicit",
    backward_solver="gmres",
    backward_max_iter=40,
    backward_tol=1e-6,
    backward_stop_mode="relative",
    backward_relative_eps=1e-8,
)

This is the package-level DEQ/SILVA adjoint path. It is useful when reproducing paper setups that train equilibria with implicit differentiation rather than finite unrolling. The transition should be deterministic during the implicit backward solve; set stochastic layers such as dropout to zero or use a stable masking strategy when exact reproducibility matters.

backward_mode="phantom" starts from the detached numerical state and records phantom_steps damped transitions with phantom_tau. This includes the common one-step approximation and longer phantom-gradient trajectories.

backward_mode="jfb" implements Jacobian-Free Backpropagation [88]. It treats the converged state as a constant and differentiates through exactly one final transition:

\[ (I-J_f^\top)^{-1}\approx I. \]

config = SolverConfig(
    solver="anderson",
    max_iter=40,
    tol=1e-6,
    backward_mode="jfb",
)

backward_mode="shine" implements inverse-estimate sharing [89]. It requires a Broyden forward solve. If (B\approx(J_f-I)^{-1}), the initial adjoint estimate is

\[ u_0=-B^\top g\approx(I-J_f^\top)^{-1}g. \]

shine_refine_steps applies additional good-Broyden updates to the exact adjoint residual. Zero steps uses the raw shared estimate.

config = SolverConfig(
    solver="broyden",
    history=10,
    backward_mode="shine",
    shine_refine_steps=2,
    backward_tol=1e-6,
)

The public helper shine_adjoint_solve accepts an equilibrium, output gradient, and BroydenInverseEstimate for direct numerical comparisons.

The implicit adjoint may use gmres, picard, anderson, or broyden through backward_solver. backward_stop_mode and backward_relative_eps select its criterion independently of the forward solve. indexing retains selected one-based forward iterations for trajectory supervision, and return_best=True returns the lowest-residual observed state when convergence is nonmonotone.

For a derivation and controlled comparison of all backward paths, see Learned Solvers and Backward Approximations.

Output Contract

fixed_point returns SolverResult:

Field	Meaning
z	selected final or best equilibrium state
states	intermediate states requested by indexing
residuals	absolute or relative residual trace
iterations, converged, solver	numerical termination diagnostics
info	nonfinite termination and implicit backward diagnostics

Tensor device and dtype follow the initial state z0.

solve_equilibrium returns the same SolverResult contract and records result.info["backward_mode"] as "unrolled", "implicit", or "phantom".

gmres and implicit_adjoint_solve return LinearSolveResult, with the same fields except that the solution field is named x.

Fixed-point and matrix-free linear solvers.

The fixed-point API follows the DEQ formulation of Bai, Kolter, and Koltun (2019): a layer returns an equilibrium state z_star = f(z_star). Picard iteration is the baseline fixed-point method, Anderson acceleration follows Anderson (1965) and Walker and Ni (2011), Broyden follows Broyden's inverse quasi-Newton update, and GMRES follows Saad and Schultz (1986) for the matrix-free adjoint systems used in implicit-gradient diagnostics.

BroydenInverseEstimate dataclass

Limited-memory approximation of the forward residual inverse.

Broyden solves g(z)=f(z)-z=0 and stores an approximation

\[ B_k \approx J_g(z_k)^{-1}=-\left(I-J_f(z_k)\right)^{-1}. \]

The factors are detached numerical quantities. They can therefore be inspected, serialized, or reused as an adjoint preconditioner without retaining the forward autograd graph.

Source code in src/silva_networks/solvers.py

@dataclass(frozen=True)
class BroydenInverseEstimate:
    r"""Limited-memory approximation of the forward residual inverse.

    Broyden solves ``g(z)=f(z)-z=0`` and stores an approximation

    $$
    B_k \approx J_g(z_k)^{-1}=-\left(I-J_f(z_k)\right)^{-1}.
    $$

    The factors are detached numerical quantities. They can therefore be
    inspected, serialized, or reused as an adjoint preconditioner without
    retaining the forward autograd graph.
    """

    shape: torch.Size
    left_factors: tuple[Tensor, ...] = ()
    right_factors: tuple[Tensor, ...] = ()

    @property
    def rank(self) -> int:
        """Number of retained rank-one inverse updates."""

        return len(self.left_factors)

    def _validate_vector(self, vector: Tensor) -> Tensor:
        if vector.shape != self.shape:
            raise ValueError(
                f"inverse estimate expects shape {tuple(self.shape)}, got {tuple(vector.shape)}"
            )
        return vector.reshape(-1)

    def apply_residual_inverse(self, vector: Tensor) -> Tensor:
        """Apply the estimated inverse of ``J_f-I`` to ``vector``."""

        flat = self._validate_vector(vector)
        value = -flat
        for left, right in zip(self.left_factors, self.right_factors, strict=True):
            value = value + left.to(flat) * torch.dot(right.to(flat), flat)
        return value.reshape(self.shape)

    def apply_residual_inverse_transpose(self, vector: Tensor) -> Tensor:
        """Apply the transpose of the estimated inverse of ``J_f-I``."""

        flat = self._validate_vector(vector)
        value = -flat
        for left, right in zip(self.left_factors, self.right_factors, strict=True):
            value = value + right.to(flat) * torch.dot(left.to(flat), flat)
        return value.reshape(self.shape)

    def apply_fixed_point_adjoint_inverse(self, vector: Tensor) -> Tensor:
        r"""Approximate ``(I-J_f^T)^{-1} vector`` from the forward solve."""

        return -self.apply_residual_inverse_transpose(vector)

rank property

rank

Number of retained rank-one inverse updates.

apply_fixed_point_adjoint_inverse

apply_fixed_point_adjoint_inverse(vector)

Approximate (I-J_f^T)^{-1} vector from the forward solve.

Source code in src/silva_networks/solvers.py

def apply_fixed_point_adjoint_inverse(self, vector: Tensor) -> Tensor:
    r"""Approximate ``(I-J_f^T)^{-1} vector`` from the forward solve."""

    return -self.apply_residual_inverse_transpose(vector)

apply_residual_inverse

apply_residual_inverse(vector)

Apply the estimated inverse of J_f-I to vector.

Source code in src/silva_networks/solvers.py

def apply_residual_inverse(self, vector: Tensor) -> Tensor:
    """Apply the estimated inverse of ``J_f-I`` to ``vector``."""

    flat = self._validate_vector(vector)
    value = -flat
    for left, right in zip(self.left_factors, self.right_factors, strict=True):
        value = value + left.to(flat) * torch.dot(right.to(flat), flat)
    return value.reshape(self.shape)

apply_residual_inverse_transpose

apply_residual_inverse_transpose(vector)

Apply the transpose of the estimated inverse of J_f-I.

Source code in src/silva_networks/solvers.py

def apply_residual_inverse_transpose(self, vector: Tensor) -> Tensor:
    """Apply the transpose of the estimated inverse of ``J_f-I``."""

    flat = self._validate_vector(vector)
    value = -flat
    for left, right in zip(self.left_factors, self.right_factors, strict=True):
        value = value + right.to(flat) * torch.dot(left.to(flat), flat)
    return value.reshape(self.shape)

LinearSolveResult dataclass

Output of a matrix-free linear solve.

Attributes:

Name	Type	Description
x	Tensor	Linear-system solution tensor.
residuals	list[float]	Linear residual norms collected during iteration.
iterations	int	Number of Krylov iterations performed.
converged	bool	Whether the tolerance criterion was met.
solver	str	Solver name.

Source code in src/silva_networks/solvers.py

@dataclass
class LinearSolveResult:
    """Output of a matrix-free linear solve.

    Attributes:
        x: Linear-system solution tensor.
        residuals: Linear residual norms collected during iteration.
        iterations: Number of Krylov iterations performed.
        converged: Whether the tolerance criterion was met.
        solver: Solver name.
    """

    x: Tensor
    residuals: list[float]
    iterations: int
    converged: bool
    solver: str

    @property
    def residual(self) -> float:
        return self.residuals[-1] if self.residuals else float("nan")

SolverConfig dataclass

Configuration for matrix-free fixed-point solvers.

Parameters:

Name	Type	Description	Default
solver	SolverName	Fixed-point method: picard, anderson, or broyden.	'picard'
max_iter	int	Maximum number of iterations.	50
tol	float	Residual tolerance for convergence.	1e-06
alpha	float	Damping factor used by Picard-style updates.	1.0
history	int	Number of previous states used by Anderson acceleration or inverse updates retained by limited-memory Broyden.	5
ridge	float	Ridge term in the Anderson least-squares system.	0.0001
beta	float	Anderson mixing parameter.	1.0
stop_mode	StopMode	Use an absolute residual or normalize by ||f(z)||.	'absolute'
relative_eps	float	Positive stabilizer in the relative residual denominator.	1e-08
anderson_batch_dims	int	Number of leading state dimensions that represent independent Anderson solves. Use 1 for states shaped (batch, features) and 0 for one coupled tensor state.	0
track_residuals	bool	If true, store residual norms in the result.	True
reengage	bool	If true, trainable modules may apply one differentiable transition after a detached accelerated solve.	True
backward_mode	BackwardMode	unrolled differentiates through finite solver steps. implicit uses a detached forward solve and an implicit adjoint; phantom differentiates through a short damped trajectory started from the detached numerical equilibrium; neumann applies a truncated Neumann series to the equilibrium adjoint; jfb differentiates through one final transition while treating the equilibrium as a constant; and shine reuses the limited-memory Broyden inverse in the adjoint calculation.	'unrolled'
backward_solver	BackwardSolverName	Matrix-free linear solver for implicit adjoints.	'gmres'
backward_max_iter	int	Maximum number of backward linear-solver iterations.	50
backward_tol	float	Residual tolerance for the backward linear solve.	1e-06
backward_stop_mode	StopMode	Absolute or relative backward residual criterion.	'absolute'
backward_relative_eps	float	Positive stabilizer for relative backward residuals.	1e-08
phantom_steps	int	Number of differentiable refinement steps used by backward_mode="phantom". One step gives the common one-step gradient approximation.	1
phantom_tau	float	Damping used by phantom-gradient refinement steps.	1.0
neumann_terms	int	Number of terms retained in the truncated Neumann approximation of the implicit adjoint.	5
shine_refine_steps	int	Number of quasi-Newton refinement steps applied to the forward Broyden inverse estimate in backward_mode="shine".	0
indexing	tuple[int, ...]	One-based solver iteration numbers whose states should be retained in SolverResult.states for trajectory supervision.	()
return_best	bool	Return the state with the lowest observed residual instead of the final iterate when the solver does not converge monotonically.	False

Source code in src/silva_networks/solvers.py

@dataclass(frozen=True)
class SolverConfig:
    """Configuration for matrix-free fixed-point solvers.

    Args:
        solver: Fixed-point method: `picard`, `anderson`, or `broyden`.
        max_iter: Maximum number of iterations.
        tol: Residual tolerance for convergence.
        alpha: Damping factor used by Picard-style updates.
        history: Number of previous states used by Anderson acceleration or
            inverse updates retained by limited-memory Broyden.
        ridge: Ridge term in the Anderson least-squares system.
        beta: Anderson mixing parameter.
        stop_mode: Use an absolute residual or normalize by ``||f(z)||``.
        relative_eps: Positive stabilizer in the relative residual denominator.
        anderson_batch_dims: Number of leading state dimensions that represent
            independent Anderson solves. Use ``1`` for states shaped
            ``(batch, features)`` and ``0`` for one coupled tensor state.
        track_residuals: If true, store residual norms in the result.
        reengage: If true, trainable modules may apply one differentiable
            transition after a detached accelerated solve.
        backward_mode: `unrolled` differentiates through finite solver steps.
            `implicit` uses a detached forward solve and an implicit adjoint;
            `phantom` differentiates through a short damped trajectory started
            from the detached numerical equilibrium; `neumann` applies a
            truncated Neumann series to the equilibrium adjoint; `jfb` differentiates
            through one final transition while treating the equilibrium as a
            constant; and `shine` reuses the limited-memory Broyden inverse in
            the adjoint calculation.
        backward_solver: Matrix-free linear solver for implicit adjoints.
        backward_max_iter: Maximum number of backward linear-solver iterations.
        backward_tol: Residual tolerance for the backward linear solve.
        backward_stop_mode: Absolute or relative backward residual criterion.
        backward_relative_eps: Positive stabilizer for relative backward
            residuals.
        phantom_steps: Number of differentiable refinement steps used by
            `backward_mode="phantom"`. One step gives the common one-step
            gradient approximation.
        phantom_tau: Damping used by phantom-gradient refinement steps.
        neumann_terms: Number of terms retained in the truncated Neumann
            approximation of the implicit adjoint.
        shine_refine_steps: Number of quasi-Newton refinement steps applied to
            the forward Broyden inverse estimate in `backward_mode="shine"`.
        indexing: One-based solver iteration numbers whose states should be
            retained in `SolverResult.states` for trajectory supervision.
        return_best: Return the state with the lowest observed residual instead
            of the final iterate when the solver does not converge monotonically.
    """

    solver: SolverName = "picard"
    max_iter: int = 50
    tol: float = 1e-6
    alpha: float = 1.0
    history: int = 5
    ridge: float = 1e-4
    beta: float = 1.0
    stop_mode: StopMode = "absolute"
    relative_eps: float = 1e-8
    anderson_batch_dims: int = 0
    track_residuals: bool = True
    reengage: bool = True
    backward_mode: BackwardMode = "unrolled"
    backward_solver: BackwardSolverName = "gmres"
    backward_max_iter: int = 50
    backward_tol: float = 1e-6
    backward_stop_mode: StopMode = "absolute"
    backward_relative_eps: float = 1e-8
    phantom_steps: int = 1
    phantom_tau: float = 1.0
    neumann_terms: int = 5
    shine_refine_steps: int = 0
    indexing: tuple[int, ...] = ()
    return_best: bool = False

    def __post_init__(self) -> None:
        if self.solver not in {"picard", "anderson", "broyden"}:
            raise ValueError(f"Unknown solver: {self.solver}")
        if self.max_iter < 1:
            raise ValueError("max_iter must be positive")
        if self.tol <= 0:
            raise ValueError("tol must be positive")
        if self.alpha <= 0:
            raise ValueError("alpha must be positive")
        if self.history < 1:
            raise ValueError("history must be positive")
        if self.ridge < 0:
            raise ValueError("ridge must be nonnegative")
        if not 0.0 <= self.beta <= 1.0:
            raise ValueError("beta must satisfy 0 <= beta <= 1")
        if self.stop_mode not in {"absolute", "relative"}:
            raise ValueError(f"Unknown stop_mode: {self.stop_mode}")
        if self.relative_eps <= 0:
            raise ValueError("relative_eps must be positive")
        if self.anderson_batch_dims < 0:
            raise ValueError("anderson_batch_dims must be nonnegative")
        if self.backward_mode not in {
            "unrolled",
            "implicit",
            "phantom",
            "neumann",
            "jfb",
            "shine",
        }:
            raise ValueError(f"Unknown backward_mode: {self.backward_mode}")
        if self.backward_solver not in {"gmres", "picard", "anderson", "broyden"}:
            raise ValueError(f"Unknown backward_solver: {self.backward_solver}")
        if self.backward_max_iter < 1:
            raise ValueError("backward_max_iter must be positive")
        if self.backward_tol <= 0:
            raise ValueError("backward_tol must be positive")
        if self.backward_stop_mode not in {"absolute", "relative"}:
            raise ValueError(f"Unknown backward_stop_mode: {self.backward_stop_mode}")
        if self.backward_relative_eps <= 0:
            raise ValueError("backward_relative_eps must be positive")
        if self.phantom_steps < 1:
            raise ValueError("phantom_steps must be positive")
        if self.phantom_tau <= 0:
            raise ValueError("phantom_tau must be positive")
        if self.neumann_terms < 1:
            raise ValueError("neumann_terms must be positive")
        if self.shine_refine_steps < 0:
            raise ValueError("shine_refine_steps must be nonnegative")
        if any(index < 1 or index > self.max_iter for index in self.indexing):
            raise ValueError("indexing entries must be between 1 and max_iter")
        if len(set(self.indexing)) != len(self.indexing):
            raise ValueError("indexing entries must be unique")

SolverResult dataclass

Output of a fixed-point solve.

Attributes:

Name	Type	Description
z	Tensor	Final state tensor.
residuals	list[float]	Residual norms collected during iteration.
iterations	int	Number of iterations performed.
converged	bool	Whether the tolerance criterion was met.
solver	str	Solver name.
info	dict[str, float | int | str]	Optional extra scalar or string diagnostics.
states	list[Tensor]	Requested intermediate states, in SolverConfig.indexing order.
inverse_estimate	BroydenInverseEstimate | None	Limited-memory inverse retained by Broyden, when available.

Source code in src/silva_networks/solvers.py

@dataclass
class SolverResult:
    """Output of a fixed-point solve.

    Attributes:
        z: Final state tensor.
        residuals: Residual norms collected during iteration.
        iterations: Number of iterations performed.
        converged: Whether the tolerance criterion was met.
        solver: Solver name.
        info: Optional extra scalar or string diagnostics.
        states: Requested intermediate states, in `SolverConfig.indexing` order.
        inverse_estimate: Limited-memory inverse retained by Broyden, when
            available.
    """

    z: Tensor
    residuals: list[float]
    iterations: int
    converged: bool
    solver: str
    info: dict[str, float | int | str] = field(default_factory=dict)
    states: list[Tensor] = field(default_factory=list)
    inverse_estimate: BroydenInverseEstimate | None = None

    @property
    def residual(self) -> float:
        return self.residuals[-1] if self.residuals else float("nan")

anderson

anderson(f, z0, config=None)

Anderson acceleration for vector-shaped states.

Parameters:

Name	Type	Description	Default
f	Callable[[Tensor], Tensor]	Transition map that accepts and returns tensors shaped like z0.	required
z0	Tensor	Initial state.	required
config	SolverConfig | None	Optional solver configuration; history, ridge, and beta control the Anderson least-squares step.	None

Returns:

Type	Description
SolverResult	SolverResult containing the final state and residual history.

Source code in src/silva_networks/solvers.py

def anderson(
    f: Callable[[Tensor], Tensor],
    z0: Tensor,
    config: SolverConfig | None = None,
) -> SolverResult:
    """Anderson acceleration for vector-shaped states.

    Args:
        f: Transition map that accepts and returns tensors shaped like `z0`.
        z0: Initial state.
        config: Optional solver configuration; `history`, `ridge`, and `beta`
            control the Anderson least-squares step.

    Returns:
        `SolverResult` containing the final state and residual history.
    """

    cfg = config or SolverConfig(solver="anderson")
    _validate_state(z0, name="z0")
    if cfg.anderson_batch_dims > z0.dim():
        raise ValueError("anderson_batch_dims cannot exceed z0.dim()")
    z = z0
    xs: list[Tensor] = []
    fs: list[Tensor] = []
    residuals: list[float] = []
    states_by_iteration: dict[int, Tensor] = {}
    converged = False
    best_state = z
    best_residual = float("inf")
    for iteration in range(1, cfg.max_iter + 1):
        fz = f(z)
        _validate_transition_output(fz, z)
        norm = _residual_norm(fz, z, cfg)
        if cfg.track_residuals:
            residuals.append(norm)
        if norm < best_residual:
            best_residual = norm
            best_state = z
        if not torch.isfinite(fz).all():
            return SolverResult(
                _select_best(z, best_state, cfg),
                residuals,
                iteration,
                False,
                "anderson",
                {"termination": "nonfinite_transition"},
                _ordered_states(states_by_iteration, cfg),
            )
        damped_fz = (1.0 - cfg.alpha) * z + cfg.alpha * fz
        if norm < cfg.tol:
            z = damped_fz
            _record_state(states_by_iteration, iteration, z, cfg)
            converged = True
            break

        xs.append(z.detach())
        fs.append(damped_fz.detach())
        if len(xs) > cfg.history:
            xs.pop(0)
            fs.pop(0)

        m = len(xs)
        if m == 1:
            z = damped_fz
            _record_state(states_by_iteration, iteration, z, cfg)
            continue

        coeff = _anderson_coefficients(xs, fs, cfg, m)
        f_mix = _anderson_mix(fs, coeff, cfg.anderson_batch_dims)
        x_mix = _anderson_mix(xs, coeff, cfg.anderson_batch_dims)
        z = cfg.beta * f_mix + (1.0 - cfg.beta) * x_mix
        _record_state(states_by_iteration, iteration, z, cfg)
    termination = "converged" if converged else "max_iter"
    return SolverResult(
        _select_best(z, best_state, cfg),
        residuals,
        iteration,
        converged,
        "anderson",
        {"termination": termination, "best_residual": best_residual},
        _ordered_states(states_by_iteration, cfg),
    )

broyden

broyden(f, z0, config=None)

Limited-memory good-Broyden inverse update for fixed-point solves.

Parameters:

Name	Type	Description	Default
f	Callable[[Tensor], Tensor]	Transition map that accepts and returns tensors shaped like z0.	required
z0	Tensor	Initial state.	required
config	SolverConfig | None	Optional solver configuration.	None

Returns:

Type	Description
SolverResult	SolverResult containing the final state and residual history.

Source code in src/silva_networks/solvers.py

def broyden(
    f: Callable[[Tensor], Tensor],
    z0: Tensor,
    config: SolverConfig | None = None,
) -> SolverResult:
    """Limited-memory good-Broyden inverse update for fixed-point solves.

    Args:
        f: Transition map that accepts and returns tensors shaped like `z0`.
        z0: Initial state.
        config: Optional solver configuration.

    Returns:
        `SolverResult` containing the final state and residual history.
    """

    cfg = config or SolverConfig(solver="broyden")
    _validate_state(z0, name="z0")
    if cfg.anderson_batch_dims != 0:
        raise ValueError("broyden solves one coupled state and requires anderson_batch_dims=0")
    shape = z0.shape
    z = z0.reshape(-1)

    def flat_residual(flat_z: Tensor) -> Tensor:
        zz = flat_z.reshape(shape)
        fz = f(zz)
        _validate_transition_output(fz, zz)
        return (fz - zz).reshape(-1)

    left_factors: list[Tensor] = []
    right_factors: list[Tensor] = []

    def apply_inverse(vector: Tensor) -> Tensor:
        value = -vector
        for left, right in zip(left_factors, right_factors, strict=True):
            value = value + left * torch.dot(right, vector)
        return value

    def apply_inverse_transpose(vector: Tensor) -> Tensor:
        value = -vector
        for left, right in zip(left_factors, right_factors, strict=True):
            value = value + right * torch.dot(left, vector)
        return value

    r = flat_residual(z)
    residuals: list[float] = []
    states_by_iteration: dict[int, Tensor] = {}
    converged = False
    best_state = z
    best_residual = float("inf")
    for iteration in range(1, cfg.max_iter + 1):
        residual = torch.linalg.norm(r)
        if cfg.stop_mode == "relative":
            fz_flat = z + r
            residual = residual / (torch.linalg.norm(fz_flat) + cfg.relative_eps)
        norm = float(residual.detach().cpu())
        if cfg.track_residuals:
            residuals.append(norm)
        if norm < best_residual:
            best_residual = norm
            best_state = z
        if not torch.isfinite(r).all():
            break
        if norm < cfg.tol:
            converged = True
            break
        step = -cfg.alpha * apply_inverse(r)
        z_next = z + step
        _record_state(states_by_iteration, iteration, z_next.reshape(shape), cfg)
        r_next = flat_residual(z_next)
        y = r_next - r
        if len(left_factors) >= cfg.history:
            left_factors.clear()
            right_factors.clear()
        By = apply_inverse(y)
        denom = torch.dot(step, By)
        if torch.isfinite(denom) and torch.abs(denom) > torch.finfo(z.dtype).eps:
            right = apply_inverse_transpose(step)
            left_factors.append((step - By) / denom)
            right_factors.append(right)
        z, r = z_next, r_next
    termination = (
        "converged"
        if converged
        else ("nonfinite_residual" if not torch.isfinite(r).all() else "max_iter")
    )
    return SolverResult(
        _select_best(z, best_state, cfg).reshape(shape),
        residuals,
        iteration,
        converged,
        "broyden",
        {
            "termination": termination,
            "best_residual": best_residual,
            "inverse_rank": len(left_factors),
        },
        _ordered_states(states_by_iteration, cfg),
        BroydenInverseEstimate(
            shape=torch.Size(shape),
            left_factors=tuple(factor.detach().clone() for factor in left_factors),
            right_factors=tuple(factor.detach().clone() for factor in right_factors),
        ),
    )

fixed_point

fixed_point(f, z0, config=None)

Dispatch to the configured fixed-point solver.

Parameters:

Name	Type	Description	Default
f	Callable[[Tensor], Tensor]	Transition map.	required
z0	Tensor	Initial state.	required
config	SolverConfig | None	Solver configuration. Defaults to SolverConfig().	None

Returns:

Type	Description
SolverResult	SolverResult from the selected method.

Source code in src/silva_networks/solvers.py

def fixed_point(
    f: Callable[[Tensor], Tensor],
    z0: Tensor,
    config: SolverConfig | None = None,
) -> SolverResult:
    """Dispatch to the configured fixed-point solver.

    Args:
        f: Transition map.
        z0: Initial state.
        config: Solver configuration. Defaults to `SolverConfig()`.

    Returns:
        `SolverResult` from the selected method.
    """

    cfg = config or SolverConfig()
    if cfg.solver == "picard":
        return picard(f, z0, cfg)
    if cfg.solver == "anderson":
        return anderson(f, z0, cfg)
    if cfg.solver == "broyden":
        return broyden(f, z0, cfg)
    raise ValueError(f"Unknown solver: {cfg.solver}")

gmres

gmres(matvec, b, *, max_iter=50, tol=1e-06, stop_mode='absolute', relative_eps=1e-08)

Matrix-free GMRES for A x = b.

matvec must return A @ v with the same shape as v. The solver materializes only the Arnoldi basis for the requested iteration budget, so it is useful for small and medium implicit-adjoint diagnostics.

Source code in src/silva_networks/solvers.py

def gmres(
    matvec: Callable[[Tensor], Tensor],
    b: Tensor,
    *,
    max_iter: int = 50,
    tol: float = 1e-6,
    stop_mode: StopMode = "absolute",
    relative_eps: float = 1e-8,
) -> LinearSolveResult:
    """Matrix-free GMRES for ``A x = b``.

    ``matvec`` must return ``A @ v`` with the same shape as ``v``. The solver
    materializes only the Arnoldi basis for the requested iteration budget, so
    it is useful for small and medium implicit-adjoint diagnostics.
    """

    if max_iter < 1:
        raise ValueError("max_iter must be positive")
    if tol <= 0:
        raise ValueError("tol must be positive")
    if stop_mode not in {"absolute", "relative"}:
        raise ValueError(f"Unknown stop_mode: {stop_mode}")
    if relative_eps <= 0:
        raise ValueError("relative_eps must be positive")
    _validate_state(b, name="b")
    shape = b.shape
    b_flat = b.reshape(-1)
    x0 = torch.zeros_like(b_flat)
    initial_matvec = matvec(x0.reshape(shape))
    _validate_transition_output(initial_matvec, b)
    r0 = b_flat - initial_matvec.reshape(-1)
    beta = torch.linalg.norm(r0)
    denominator = (
        torch.linalg.norm(b_flat) + relative_eps
        if stop_mode == "relative"
        else b_flat.new_tensor(1.0)
    )
    initial_residual = beta / denominator
    residuals: list[float] = [float(initial_residual.detach().cpu())]
    if residuals[-1] < tol:
        return LinearSolveResult(x0.reshape(shape), residuals, 0, True, "gmres")

    basis: list[Tensor] = [r0 / beta.clamp_min(torch.finfo(b.dtype).tiny)]
    hessenberg_columns: list[Tensor] = []
    x_flat = x0
    converged = False
    iteration = 0

    for iteration in range(1, max_iter + 1):
        j = iteration - 1
        matvec_out = matvec(basis[j].reshape(shape))
        _validate_transition_output(matvec_out, b)
        w = matvec_out.reshape(-1)
        column_entries: list[Tensor] = []
        for i in range(iteration):
            coefficient = torch.vdot(basis[i], w)
            column_entries.append(coefficient)
            w = w - coefficient * basis[i]
        arnoldi_norm = torch.linalg.norm(w)
        column_entries.append(arnoldi_norm)
        hessenberg_columns.append(torch.stack(column_entries))
        if arnoldi_norm > torch.finfo(b.dtype).eps and iteration < max_iter:
            basis.append(w / arnoldi_norm)

        rows = iteration + 1
        padded_columns = [
            torch.cat(
                [
                    column,
                    column.new_zeros(rows - column.numel()),
                ]
            )
            for column in hessenberg_columns
        ]
        hessenberg = torch.stack(padded_columns, dim=1)
        rhs = torch.cat([beta.reshape(1), beta.new_zeros(iteration)])
        y = torch.linalg.lstsq(hessenberg, rhs).solution
        v_mat = torch.stack(basis[:iteration], dim=1)
        x_flat = v_mat @ y
        matvec_out = matvec(x_flat.reshape(shape))
        _validate_transition_output(matvec_out, b)
        residual = torch.linalg.norm(matvec_out.reshape(-1) - b_flat)
        residual_value = residual / denominator
        residuals.append(float(residual_value.detach().cpu()))
        if residuals[-1] < tol:
            converged = True
            break
        if arnoldi_norm <= torch.finfo(b.dtype).eps:
            numerical_floor = 100.0 * torch.finfo(b.dtype).eps
            if stop_mode == "absolute":
                numerical_floor *= 1.0 + float(torch.linalg.norm(b_flat).detach().cpu())
            converged = residuals[-1] <= max(tol, numerical_floor)
            break

    return LinearSolveResult(x_flat.reshape(shape), residuals, iteration, converged, "gmres")

implicit_adjoint_solve

implicit_adjoint_solve(f, z_star, grad_output, *, alpha=1.0, max_iter=50, tol=1e-06, solver='gmres', stop_mode='absolute', relative_eps=1e-08)

Solve the DEQ adjoint system with VJP-backed GMRES.

For the damped update T_alpha(z)=(1-alpha)z+alpha f(z), the adjoint vector u solves

\[ (I - J_{T_\alpha}(z^\star)^T)u = g. \]

The returned u can be used with torch.autograd.grad to obtain parameter sensitivities of the equilibrium map.

Source code in src/silva_networks/solvers.py

def implicit_adjoint_solve(
    f: Callable[[Tensor], Tensor],
    z_star: Tensor,
    grad_output: Tensor,
    *,
    alpha: float = 1.0,
    max_iter: int = 50,
    tol: float = 1e-6,
    solver: BackwardSolverName = "gmres",
    stop_mode: StopMode = "absolute",
    relative_eps: float = 1e-8,
) -> LinearSolveResult:
    r"""Solve the DEQ adjoint system with VJP-backed GMRES.

    For the damped update ``T_alpha(z)=(1-alpha)z+alpha f(z)``, the adjoint
    vector ``u`` solves

    $$
    (I - J_{T_\alpha}(z^\star)^T)u = g.
    $$

    The returned ``u`` can be used with ``torch.autograd.grad`` to obtain
    parameter sensitivities of the equilibrium map.
    """

    if solver not in {"gmres", "picard", "anderson", "broyden"}:
        raise ValueError(f"Unknown backward solver: {solver}")
    z_req = z_star.detach().requires_grad_(True)
    y = f(z_req)
    _validate_transition_output(y, z_req)

    def matvec(v: Tensor) -> Tensor:
        jtv = torch.zeros_like(v)
        if y.requires_grad:
            (maybe_jtv,) = torch.autograd.grad(
                y,
                z_req,
                v,
                retain_graph=True,
                create_graph=False,
                allow_unused=True,
            )
            if maybe_jtv is not None:
                jtv = maybe_jtv
        damped_jtv = (1.0 - alpha) * v + alpha * jtv
        return v - damped_jtv

    if solver == "gmres":
        return gmres(
            matvec,
            grad_output,
            max_iter=max_iter,
            tol=tol,
            stop_mode=stop_mode,
            relative_eps=relative_eps,
        )

    def adjoint_map(v: Tensor) -> Tensor:
        return grad_output + v - matvec(v)

    result = fixed_point(
        adjoint_map,
        torch.zeros_like(grad_output),
        SolverConfig(
            solver=solver,
            max_iter=max_iter,
            tol=tol,
            stop_mode=stop_mode,
            relative_eps=relative_eps,
        ),
    )
    return LinearSolveResult(
        x=result.z,
        residuals=result.residuals,
        iterations=result.iterations,
        converged=result.converged,
        solver=solver,
    )

picard

picard(f, z0, config=None)

Damped Picard iteration for z = f(z).

Parameters:

Name	Type	Description	Default
f	Callable[[Tensor], Tensor]	Transition map that accepts and returns tensors shaped like z0.	required
z0	Tensor	Initial state.	required
config	SolverConfig | None	Optional solver configuration.	None

Returns:

Type	Description
SolverResult	SolverResult containing the final state and residual history.

Source code in src/silva_networks/solvers.py

def picard(
    f: Callable[[Tensor], Tensor], z0: Tensor, config: SolverConfig | None = None
) -> SolverResult:
    """Damped Picard iteration for `z = f(z)`.

    Args:
        f: Transition map that accepts and returns tensors shaped like `z0`.
        z0: Initial state.
        config: Optional solver configuration.

    Returns:
        `SolverResult` containing the final state and residual history.
    """

    cfg = config or SolverConfig()
    _validate_state(z0, name="z0")
    z = z0
    residuals: list[float] = []
    states_by_iteration: dict[int, Tensor] = {}
    converged = False
    best_state = z
    best_residual = float("inf")
    for iteration in range(1, cfg.max_iter + 1):
        fz = f(z)
        _validate_transition_output(fz, z)
        r = _residual_norm(fz, z, cfg)
        if cfg.track_residuals:
            residuals.append(r)
        if r < best_residual:
            best_residual = r
            best_state = z
        if not torch.isfinite(fz).all():
            return SolverResult(
                _select_best(z, best_state, cfg),
                residuals,
                iteration,
                False,
                "picard",
                {"termination": "nonfinite_transition"},
                _ordered_states(states_by_iteration, cfg),
            )
        z = (1.0 - cfg.alpha) * z + cfg.alpha * fz
        _record_state(states_by_iteration, iteration, z, cfg)
        if r < cfg.tol:
            converged = True
            break
    termination = "converged" if converged else "max_iter"
    return SolverResult(
        _select_best(z, best_state, cfg),
        residuals,
        iteration,
        converged,
        "picard",
        {"termination": termination, "best_residual": best_residual},
        _ordered_states(states_by_iteration, cfg),
    )

reengage_result

reengage_result(result, f, config=None, *, force=False)

Reconnect a numerical fixed-point result to autograd when needed.

Anderson acceleration keeps its history detached for numerical stability and memory control. Trainable modules can call this helper after fixed_point so the returned state participates in ordinary PyTorch gradients without making Picard or Broyden runs take an extra step.

Source code in src/silva_networks/solvers.py

def reengage_result(
    result: SolverResult,
    f: Callable[[Tensor], Tensor],
    config: SolverConfig | None = None,
    *,
    force: bool = False,
) -> SolverResult:
    """Reconnect a numerical fixed-point result to autograd when needed.

    Anderson acceleration keeps its history detached for numerical stability and
    memory control. Trainable modules can call this helper after `fixed_point` so
    the returned state participates in ordinary PyTorch gradients without making
    Picard or Broyden runs take an extra step.
    """

    cfg = config or SolverConfig()
    if force or (cfg.reengage and cfg.solver == "anderson"):
        fz = f(result.z)
        _validate_transition_output(fz, result.z)
        result.z = (1.0 - cfg.alpha) * result.z + cfg.alpha * fz
    return result

shine_adjoint_solve

shine_adjoint_solve(f, z_star, grad_output, inverse_estimate, *, refine_steps=0, tol=1e-06, stop_mode='absolute', relative_eps=1e-08)

Reuse a forward Broyden inverse to approximate the DEQ adjoint.

The forward estimate approximates (J_f-I)^{-1}, so its negative transpose approximates the inverse of the adjoint operator A=I-J_f^T. Optional good-Broyden updates refine that shared estimate on the linear residual A u-g while retaining only refine_steps rank-one corrections.

Source code in src/silva_networks/solvers.py

def shine_adjoint_solve(
    f: Callable[[Tensor], Tensor],
    z_star: Tensor,
    grad_output: Tensor,
    inverse_estimate: BroydenInverseEstimate,
    *,
    refine_steps: int = 0,
    tol: float = 1e-6,
    stop_mode: StopMode = "absolute",
    relative_eps: float = 1e-8,
) -> LinearSolveResult:
    r"""Reuse a forward Broyden inverse to approximate the DEQ adjoint.

    The forward estimate approximates ``(J_f-I)^{-1}``, so its negative
    transpose approximates the inverse of the adjoint operator
    ``A=I-J_f^T``. Optional good-Broyden updates refine that shared estimate on
    the linear residual ``A u-g`` while retaining only `refine_steps` rank-one
    corrections.
    """

    if refine_steps < 0:
        raise ValueError("refine_steps must be nonnegative")
    if tol <= 0:
        raise ValueError("tol must be positive")
    if stop_mode not in {"absolute", "relative"}:
        raise ValueError(f"Unknown stop_mode: {stop_mode}")
    if relative_eps <= 0:
        raise ValueError("relative_eps must be positive")
    if z_star.shape != grad_output.shape:
        raise ValueError("z_star and grad_output must have the same shape")
    if inverse_estimate.shape != z_star.shape:
        raise ValueError("inverse_estimate and z_star must have the same shape")

    z_req = z_star.detach().requires_grad_(True)
    with torch.enable_grad():
        y = f(z_req)
    _validate_transition_output(y, z_req)

    def matvec(vector: Tensor) -> Tensor:
        jtv = torch.zeros_like(vector)
        if y.requires_grad:
            (maybe_jtv,) = torch.autograd.grad(
                y,
                z_req,
                vector,
                retain_graph=True,
                create_graph=False,
                allow_unused=True,
            )
            if maybe_jtv is not None:
                jtv = maybe_jtv
        return vector - jtv

    shape = grad_output.shape
    right_hand_side = grad_output.reshape(-1)
    base_inverse = inverse_estimate.apply_fixed_point_adjoint_inverse
    correction_left: list[Tensor] = []
    correction_right: list[Tensor] = []

    def apply_inverse(vector: Tensor) -> Tensor:
        flat = vector.reshape(-1)
        value = base_inverse(vector).reshape(-1)
        for left, right in zip(correction_left, correction_right, strict=True):
            value = value + left * torch.dot(right, flat)
        return value.reshape(shape)

    def apply_inverse_transpose(vector: Tensor) -> Tensor:
        flat = vector.reshape(-1)
        value = -inverse_estimate.apply_residual_inverse(vector).reshape(-1)
        for left, right in zip(correction_left, correction_right, strict=True):
            value = value + right * torch.dot(left, flat)
        return value.reshape(shape)

    solution = apply_inverse(grad_output)
    residual = matvec(solution) - grad_output
    denominator = (
        torch.linalg.norm(right_hand_side) + relative_eps
        if stop_mode == "relative"
        else right_hand_side.new_tensor(1.0)
    )

    def residual_value(value: Tensor) -> float:
        return float((torch.linalg.norm(value.reshape(-1)) / denominator).detach().cpu())

    residuals = [residual_value(residual)]
    converged = residuals[-1] < tol
    performed = 0
    for _ in range(refine_steps):
        if converged:
            break
        step = -apply_inverse(residual)
        next_solution = solution + step
        next_residual = matvec(next_solution) - grad_output
        delta_residual = next_residual - residual
        inverse_delta = apply_inverse(delta_residual)
        denominator_update = torch.dot(step.reshape(-1), inverse_delta.reshape(-1))
        if (
            torch.isfinite(denominator_update)
            and torch.abs(denominator_update) > torch.finfo(step.dtype).eps
        ):
            transpose_step = apply_inverse_transpose(step)
            correction_left.append(((step - inverse_delta) / denominator_update).reshape(-1))
            correction_right.append(transpose_step.reshape(-1))
        solution, residual = next_solution, next_residual
        performed += 1
        residuals.append(residual_value(residual))
        converged = residuals[-1] < tol

    return LinearSolveResult(
        x=solution,
        residuals=residuals,
        iterations=performed,
        converged=converged,
        solver="shine",
    )

solve_equilibrium

solve_equilibrium(f, z0, config=None, *, params=(), tensors=(), backward_map=None)

Solve an equilibrium with the configured forward and backward mode.

backward_mode="unrolled" keeps the ordinary PyTorch finite-solver graph. backward_mode="implicit" runs the forward fixed-point solve detached and reconnects trainable sensitivities through the DEQ adjoint system

\[ (I - J_{T_\alpha}(z^\star)^T)u = \partial \mathcal L / \partial z^\star, \]

where T_alpha(z)=(1-alpha)z+alpha f(z). Pass module parameters through params and differentiable non-state inputs through tensors when using the implicit mode. backward_mode="phantom" instead performs a detached solve followed by phantom_steps differentiable refinements with damping phantom_tau; one step is the one-step-gradient special case. In backward_mode="neumann", the forward root is detached and the adjoint inverse is approximated by a finite Neumann series. In backward_mode="jfb", the converged state is treated as a constant and one final transition supplies the parameter gradient. In backward_mode="shine", a Broyden forward solve shares its inverse estimate with the adjoint and may refine it for shine_refine_steps iterations.

backward_map optionally separates the numerical forward approximation from the equilibrium map used for implicit or phantom differentiation. It is useful when the forward solve uses a source-compatible acceleration, such as thresholded delta updates, while the derivative is defined by the original equilibrium equation. Unrolled differentiation always follows f directly.

Source code in src/silva_networks/solvers.py

def solve_equilibrium(
    f: Callable[[Tensor], Tensor],
    z0: Tensor,
    config: SolverConfig | None = None,
    *,
    params: Iterable[Tensor] = (),
    tensors: Iterable[Tensor] = (),
    backward_map: Callable[[Tensor], Tensor] | None = None,
) -> SolverResult:
    r"""Solve an equilibrium with the configured forward and backward mode.

    `backward_mode="unrolled"` keeps the ordinary PyTorch finite-solver graph.
    `backward_mode="implicit"` runs the forward fixed-point solve detached and
    reconnects trainable sensitivities through the DEQ adjoint system

    $$
    (I - J_{T_\alpha}(z^\star)^T)u = \partial \mathcal L / \partial z^\star,
    $$

    where ``T_alpha(z)=(1-alpha)z+alpha f(z)``. Pass module parameters through
    `params` and differentiable non-state inputs through `tensors` when using
    the implicit mode. `backward_mode="phantom"` instead performs a detached
    solve followed by `phantom_steps` differentiable refinements with damping
    `phantom_tau`; one step is the one-step-gradient special case. In
    `backward_mode="neumann"`, the forward root is detached and the adjoint
    inverse is approximated by a finite Neumann series. In `backward_mode="jfb"`,
    the converged state is treated as a constant and one
    final transition supplies the parameter gradient. In
    `backward_mode="shine"`, a Broyden forward solve shares its inverse estimate
    with the adjoint and may refine it for `shine_refine_steps` iterations.

    ``backward_map`` optionally separates the numerical forward approximation
    from the equilibrium map used for implicit or phantom differentiation. It
    is useful when the forward solve uses a source-compatible acceleration,
    such as thresholded delta updates, while the derivative is defined by the
    original equilibrium equation. Unrolled differentiation always follows
    ``f`` directly.
    """

    cfg = config or SolverConfig()
    if cfg.backward_mode == "unrolled":
        result = fixed_point(f, z0, cfg)
        reengage_result(result, f, cfg)
        result.info.setdefault("backward_mode", "unrolled")
        return result
    sensitivity_map = f if backward_map is None else backward_map
    with torch.no_grad():
        result = fixed_point(f, z0.detach(), cfg)

    if cfg.backward_mode == "phantom":
        z = result.z.detach()
        for _ in range(cfg.phantom_steps):
            fz = sensitivity_map(z)
            _validate_transition_output(fz, z)
            z = (1.0 - cfg.phantom_tau) * z + cfg.phantom_tau * fz
        result.z = z
        result.info.setdefault("backward_mode", "phantom")
        result.info.setdefault("phantom_steps", cfg.phantom_steps)
        result.info.setdefault("phantom_tau", cfg.phantom_tau)
        return result

    if cfg.backward_mode == "jfb":
        z = result.z.detach()
        result.z = sensitivity_map(z)
        _validate_transition_output(result.z, z)
        result.info.setdefault("backward_mode", "jfb")
        result.info.setdefault("backward_solver", "identity")
        return result

    backward_tensors = _unique_trainable_tensors(params, tensors)
    if cfg.backward_mode == "neumann":
        neumann_context = _NeumannBackwardContext(
            f=sensitivity_map,
            alpha=cfg.alpha,
            terms=cfg.neumann_terms,
            info=result.info,
        )
        result.z = _NeumannEquilibriumFunction.apply(
            result.z.detach(), *backward_tensors, neumann_context
        )
        result.info.setdefault("backward_mode", "neumann")
        result.info.setdefault("backward_solver", "truncated_neumann")
        result.info.setdefault("neumann_terms", cfg.neumann_terms)
        return result

    if cfg.backward_mode == "shine":
        if cfg.solver != "broyden" or result.inverse_estimate is None:
            raise ValueError('backward_mode="shine" requires solver="broyden"')
        shine_context = _SHINEBackwardContext(
            f=sensitivity_map,
            inverse_estimate=result.inverse_estimate,
            refine_steps=cfg.shine_refine_steps,
            tol=cfg.backward_tol,
            stop_mode=cfg.backward_stop_mode,
            relative_eps=cfg.backward_relative_eps,
            info=result.info,
        )
        result.z = _SHINEEquilibriumFunction.apply(
            result.z.detach(), *backward_tensors, shine_context
        )
        result.info.setdefault("backward_mode", "shine")
        result.info.setdefault("backward_solver", "shared_broyden_inverse")
        result.info.setdefault("shine_inverse_rank", result.inverse_estimate.rank)
        return result

    context = _ImplicitBackwardContext(
        f=sensitivity_map,
        alpha=cfg.alpha,
        max_iter=cfg.backward_max_iter,
        tol=cfg.backward_tol,
        stop_mode=cfg.backward_stop_mode,
        relative_eps=cfg.backward_relative_eps,
        solver=cfg.backward_solver,
        history=cfg.history,
        ridge=cfg.ridge,
        beta=cfg.beta,
        info=result.info,
    )
    result.z = _ImplicitEquilibriumFunction.apply(result.z.detach(), *backward_tensors, context)
    result.info.setdefault("backward_mode", "implicit")
    result.info.setdefault("backward_solver", cfg.backward_solver)
    return result

Where to Go Next

Question	Page
What mathematical problem do these solvers address?	Fixed Points
How is each update derived?	Solver Derivation Lab
Where is a solver checked against a closed form?	Scalar Equilibrium Example