make_adjoint_layer#

iskra.adjoint.make_adjoint_layer(solver_fn: Callable[[PS], TS], implicit_fn: Callable[[PI], TI], param_args: Sequence[int] | int = 1, sol_args: Sequence[int] | int = 0, zero_args: Sequence[int] | int = 0, bwd_method: Literal['gmres'] = 'gmres', gmres_init: Tensor | None = None, callback_gmres_sol: Callable | None = None, bwd_max_iter: int = 500, bwd_abs_tol: float = 1e-06, bwd_rel_tol: float = 0.001, verbose: bool = False) → Callable[[PS], TS][SOURCE]#

Makes a solver differentiable via the adjoint method using a corresponding implicit function.

Given a non-differentiable solver \(y \leftarrow g(x)\) (solver_fn), and a function representing the implicit relation \(f(x, y(x)) = 0\) (implicit_fn), constructs a differentiable solver via the adjoint method. The inputs and outputs of \(g\) are concatenated and passed forward into \(f\). Specifically, treating \(x\) and \(y\) as tuples, we index into \(x\) using param_args and into \(y\) using sol_args and concatenate those values into a new tuple that will be passed into \(f\). The user must select which outputs of \(f\) at indices zero_args are zero when the implicit relationship is satisfied.

Parameters:

  • solver_fn (Callable[PS, TS]) – The solver function \(g\).

  • implicit_fn (Callable[PI, TI]) – The implicit relation \(f\).

  • param_args (Sequence[int] | int) – Indices of \(g\)’s inputs that correspond to x.

  • sol_args (Sequence[int] | int) – Indices of \(g\)’s outputs that correspond to y.

  • zero_args (Sequence[int] | int) – Indices of \(f\)’s outputs which should equal 0.

  • bwd_method (Literal["gmres"]) – What solver to use for the adjoint equations. Currently only GMRES is provided.

  • gmres_init (Tensor | None) – Initial value for GMRES solver.

  • callback_gmres_sol (Callable | None) – Callback function that exposes the GMRES solution to the user.

  • bwd_max_iter (int) – Number of GMRES iterations.

  • bwd_abs_tol (float) – GMRES absolute tolerance.

  • bwd_rel_tol (float) – GMRES relative tolerance.

  • verbose (bool) – Print verbose statements.

Returns:

​ (Callable[[P], T]) – The differentiable solver function.

On this page