API reference

Public

ImplicitDifferentiation.ImplicitFunctionType
ImplicitFunction{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 represent A and B with a LinearOperator from 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 computing y(x), does not need to be compatible with automatic differentiation
  • conditions: a callable computing c(x, y), must be compatible with automatic differentiation
  • linear_solver: a callable to solve the linear system
  • conditions_x_backend: how the conditions will be differentiated w.r.t. the first argument x
  • conditions_y_backend: how the conditions will be differentiated w.r.t. the second argument y

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:

standardbyproduct
forward(x, args...; kwargs...) = yconditions(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::AbstractVector such that A * s = b
  • (A, B::AbstractMatrix) -> S::AbstractMatrix such that A * 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 the ImplicitFunction) is used.
  • an object subtyping AbstractADType from ADTypes.jl;

When differentiating with Enzyme as an outer backend, the default setting assumes that conditions does not contain writeable data involved in derivatives.

source
ImplicitDifferentiation.KrylovLinearSolverType
KrylovLinearSolver

Callable 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.

source

Internals