API reference
Public
ImplicitDifferentiation.ImplicitDifferentiation — ModuleImplicitDifferentiationA Julia package for automatic differentiation of implicit functions.
Its main export is the type ImplicitFunction.
ImplicitDifferentiation.ImplicitFunction — TypeImplicitFunction{lazy}Wrapper for an implicit function defined by a forward mapping y and a set of conditions c.
An ImplicitFunction object behaves like a function, and every call is differentiable with respect to the first argument x. When a derivative is queried, the Jacobian of y is computed using the implicit function theorem:
∂/∂y c(x, y(x)) * ∂/∂x y(x) = -∂/∂x c(x, y(x))This requires solving a linear system A * J = -B where A = ∂c/∂y, B = ∂c/∂x and J = ∂y/∂x.
Type parameters
lazy::Bool: whether to representAandBwith aLinearOperatorfrom LinearOperators.jl (lazy = true) or a dense Jacobian matrix (lazy = false)
Usually, dense Jacobians are more efficient in small dimension, while lazy operators become necessary in high dimension. The value of lazy must be chosen together with the linear_solver, see below.
Fields
forward: a callable computingy(x), does not need to be compatible with automatic differentiationconditions: a callable computingc(x, y), must be compatible with automatic differentiationlinear_solver: a callable to solve the linear systemconditions_x_backend: how the conditions will be differentiated w.r.t. the first argumentxconditions_y_backend: how the conditions will be differentiated w.r.t. the second argumenty
Constructors
ImplicitFunction(
+ forward, conditions;
+ linear_solver=KrylovLinearSolver(),
+ conditions_x_backend=nothing,
+ conditions_x_backend=nothing,
+)Picks the lazy parameter automatically based on the linear_solver, using the following heuristic: lazy = linear_solver != \.
ImplicitFunction{lazy}(
+ forward, conditions;
+ linear_solver=lazy ? KrylovLinearSolver() : \,
+ conditions_x_backend=nothing,
+ conditions_y_backend=nothing,
+)Picks the linear_solver automatically based on the lazy parameter.
Callable behavior
(implicit::ImplicitFunction)(x::AbstractVector, args...; kwargs...)Return implicit.forward(x, args...; kwargs...), which can be either an AbstractVector y or a tuple (y, z).
This call makes y differentiable with respect to x.
Function signatures
There are two possible signatures for forward and conditions, which must be consistent with one another:
| standard | byproduct |
|---|---|
forward(x, args...; kwargs...) = y | conditions(x, y, args...; kwargs...) = c |
forward(x, args...; kwargs...) = (y, z) | conditions(x, y, z, args...; kwargs...) = c |
In both cases, x, y and c must be AbstractVectors, with length(y) = length(c). In the second case, the byproduct z can be an arbitrary object generated by forward. The positional arguments args... and keyword arguments kwargs... must be the same for both forward and conditions.
The byproduct z and the other positional arguments args... beyond x are considered constant for differentiation purposes.
Linear solver
The provided linear_solver objects needs to be callable, with two methods:
(A, b::AbstractVector) -> s::AbstractVectorsuch thatA * s = b(A, B::AbstractMatrix) -> S::AbstractMatrixsuch thatA * S = B
It can be either a direct solver (like \) or an iterative one (like KrylovLinearSolver). Typically, direct solvers work best with dense Jacobians (lazy = false) while iterative solvers work best with operators (lazy = true).
Condition backends
The provided conditions_x_backend and conditions_y_backend can be either:
nothing(the default), in which case the outer backend (the one differentiating through theImplicitFunction) is used.- an object subtyping
AbstractADTypefrom ADTypes.jl;
When differentiating with Enzyme as an outer backend, the default setting assumes that conditions does not contain writeable data involved in derivatives.
ImplicitDifferentiation.KrylovLinearSolver — TypeKrylovLinearSolverCallable object that can solve linear systems Ax = b and AX = B in the same way as the built-in \. Uses an iterative solver from Krylov.jl under the hood.
Constructor
KrylovLinearSolver(; verbose=true)If verbose is true, the solver logs a warning in case of failure. Otherwise it will fail silently, and may return solutions that do not exactly satisfy the linear system.
Callable behavior
(::KylovLinearSolver)(A, b::AbstractVector)Solve a linear system with a single right-hand side.
(::KrylovLinearSolver)(A, B::AbstractMatrix)Solve a linear system with multiple right-hand sides.
