AutoDiffOperators

Provides Julia operators that act via automatic differentiation.

AutoDiffOperators.FwdRevADSelector{Fwd<:ADSelector,Rev<:ADSelector} <: ADSelector

Represent an automatic differentiation backend that forwards forward-mode and reverse-mode AD to two separate selectors fwd::ADSelector and rev::ADSelector.

User code should not instantiate AutoDiffOperators.FwdRevADSelector directly, but use ADSelector(fwd, rev) or ADSelector(fwd = fwd, rev = rev) instead.

abstract type AutoDiffOperators.WrappedADSelector

Supertype for AD selectors that wrap other AD selectors.

const ADSelector = Union{
    AbstractADType,
    WrappedADSelector
}

Instances specify an automatic differentiation backend.

Either a subtype of ADTypes.AbstractADType, or an AD-selector wrapper like AutoDiffOperators.FwdRevADSelector.

In addition to using instances of AbstractADType directly (e.g. ADTypes.AutoForwardDiff()), ADSelector (specifically AbstractADType) instances for AD backends can be constructed directly from the backend modules (using default backend parameters):

import ForwardDiff

ADTypes.AutoForwardDiff()
ADSelector(ForwardDiff)
convert(ADSelector, ForwardDiff)

all construct an identical AutoForwardDiff object.

Separate AD backends for forward- and reverse-mode AD can be specified via ADSelector(fwd_adtype, rev_adtype), e.g.

import ForwardDiff, Mooncake

ADSelector(ADTypes.AutoForwardDiff(), ADTypes.AutoMooncake())
ADSelector(ADSelector(ForwardDiff), ADSelector(Mooncake))
ADSelector(ForwardDiff, Mooncake)

Implementation

ADSelector instances can also be constructed from module names, though this should be avoided in end-user code:

ADSelector(:ForwardDiff)
ADSelector(Val(:ForwardDiff))
convert(ADSelector, :ForwardDiff)
convert(ADSelector, Val(:ForwardDiff))

End-users should use module objects instead of module name, so that the respective AD backend package must be part of their environment/dependencies.

forward_adtype(ad::ADSelector)::ADTypes.AbstractADType

Returns the forward-mode AD backend selector for ad.

Returns ad itself by default if ad supports forward-mode automatic differentiation, or an instance of ADTypes.NoAutoDiff if it does not.

May be specialized for some AD selector types, see FwdRevADSelector, for example.

gradient_func(f, ad::ADSelector)
gradient_func(f, ad::ADSelector, dummy_x::AbstractVector{<:Number})

Returns a function ∇f that calculates the gradient of f at a given point x, so that ∇f(x) is equivalent to only_gradient(f, x, ad).

Passing a dummy_x that matches later arguments of ∇f in type and size enables backend-specific preparation. ∇f is thread-safe either way.

jvp_func(f, x::AbstractVector{<:Number}, ad::ADSelector)

Returns a function jvp with jvp(z) == J * z.

AutoDiffOperators.mulfunc_operator(
    ::Type{OP},
    ::Type{T}, sz::Dims{2}, ovp, vop,
    ::Val{sym}, ::Val{herm}, ::Val{posdef}
) where {OP, T<:Real, sym, herm, posdef}

Generates a linear operator object of type OP that supports multiplication with (adjoint) vectors based on a multiplication function ovp and an adjoint multiplication function vop.

An operator op = mulfunc_operator(OP, T, sz, ovp, vop, Val(sym), ::Val{herm}, ::Val{posdef}) must show the following behavior:

op isa OP
eltype(op) == T
size(op) == sz
op * x_r == ovp(x_r)
x_l' * op == vop(x_l)
issymmetric(op) == sym
ishermitian(op) == herm
isposdef(op) == posdef

where x_l and x_r are vectors of size sz[1] and sz[2] respectively.

only_gradient(f, x::AbstractVector{<:Number}, ad::ADSelector)

Returns the gradient ∇f(x) of f at x.

See also with_gradient(f, x, ad).

only_gradient!(f, δx, x::AbstractVector{<:Number}, ad::ADSelector)

Fills δx with the gradient ∇f(x) of f at x and returns it.

only_gradient!!(f, δx, x::AbstractVector{<:Number}, ad::ADSelector)

Returns the gradient ∇f(x) of f at x.

δx may or may not be reused/overwritten and returned as ∇f(x).

reverse_adtype(ad::ADSelector)::ADTypes.AbstractADType

Returns the reverse-mode AD backend selector for ad.

Returns ad itself by default if ad supports reverse-mode automatic differentiation, or an instance of ADTypes.NoAutoDiff if it does not.

May be specialized for some AD selector types, see FwdRevADSelector, for example.

AutoDiffOperators.similar_onehot(A::AbstractArray, ::Type{T}, n::Integer, i::Integer)

Return an array similar to A, but with n elements of type T, all set to zero but the i-th element set to one.

valgrad_func(f, ad::ADSelector)
valgrad_func(f, ad::ADSelector, dummy_x::AbstractVector{<:Number})

Returns a function f_∇f that calculates the value and gradient of f at given points, so that f_∇f(x) is equivalent to with_gradient(f, x, ad).

Passing a dummy_x that matches later arguments of f_∇f in type and size enables backend-specific preparation. f_∇f is thread-safe either way.

valid_forward_adtype(ad::ADSelector)::ADTypes.AbstractADType

Similar to forward_adtype, but throws an exception if ad doesn't support forward-mode automatic differentiation instead of returning a NoAutoDiff.

valid_reverse_adtype(ad::ADSelector)::ADTypes.AbstractADType

Similar to reverse_adtype, but throws an exception if ad doesn't support reverse-mode automatic differentiation instead of returning a NoAutoDiff.

AutoDiffOperators.with_floatlike_contents(A::AbstractArray)

If the elements of A are integer-like, convert them using float, otherwise return A unchanged.

The elements must be real-valued numbers, though not necessarily Real (e.g. tracing- and symbolic-number types), so realness is checked via real(T) === T instead of dispatch. Complex and other non-real numbers are rejected, real-AD conventions would silently give incorrect results for them.

with_gradient(f, x::AbstractVector{<:Number}, ad::ADSelector)

Returns a tuple (f(x), ∇f(x)) with the gradient ∇f(x) of f at x.

See also with_gradient!!(f, δx, x, ad) for the "maybe-in-place" variant of this function.

with_gradient!(f, δx, x::AbstractVector{<:Number}, ad::ADSelector)

Fills δx with the gradient ∇f(x) of f at x and returns the tuple (f(x), δx).

with_gradient!!(f, δx, x::AbstractVector{<:Number}, ad::ADSelector)

Returns a tuple (f(x), ∇f(x)) with the gradient ∇f(x) of f at x.

δx may or may not be reused/overwritten and returned as ∇f(x).

with_jacobian(f, x::AbstractVector{<:Number}, OP, ad::ADSelector)

Returns a tuple (f(x), J) with a multiplicative Jacobian operator J of type OP.

Example:

using AutoDiffOperators, LinearMaps
y, J = with_jacobian(f, x, LinearMap, ad)
y == f(x)
_, J_explicit = with_jacobian(f, x, DenseMatrix, ad)
J * z_r ≈ J_explicit * z_r
z_l' * J ≈ z_l' * J_explicit

OP may be LinearMaps.LinearMap (resp. LinearMaps.FunctionMap) or Matrix. Other operator types can be supported by specializing mulfunc_operator for the operator type.

The default implementation of with_jacobian uses jvp_func and with_vjp_func to implement (adjoint) multiplication of J with (adjoint) vectors.

with_jvp(f, x::AbstractVector{<:Number}, z::AbstractVector{<:Number}, ad::ADSelector)

Returns a tuple (f(x), J * z).

with_vjp_func(f, x::AbstractVector{<:Number}, ad::ADSelector)

Returns a tuple (f(x), vjp) with the function vjp(z) ≈ J' * z.