make_adjoint_vjps#
- iskra.adjoint.make_adjoint_vjps(fn: Callable[[P], T], param_args: Sequence[int] | int = 1, sol_args: Sequence[int] | int = 0, zero_args: Sequence[int] | int = 0, *args: P, **kwargs: P) tuple[Callable[[Tensor], Tensor], Callable[[Tensor], tuple[Tensor, ...]]][SOURCE]#
Constructs the VJP functions needed for adjoint computation.
Given a function representing the implicit relation \(f(x, y(x)) = 0\), constructs functions for computing vector-Jacobian products with the Jacobians \(\frac{\partial f}{\partial y}\) and \(\frac{\partial f}{\partial x}\), needed to compute the adjoint \(\frac{\partial y}{\partial x}\).
- Parameters:
fn (Callable[P, T]) – The implicit relation \(f\).
param_args (Sequence[int] | int) – Indices of \(f\)’s inputs that correspond to x.
sol_args (Sequence[int] | int) – Indices of \(f\)’s inputs that correspond to y.
zero_args (Sequence[int] | int) – Outputs of \(f\) which should equal 0.
*args – (P.args): Arguments to be passed into fn.
**kwargs – (P.kwargs): Keyword arguments to be passed into fn.
- Returns: