ADTypes.jl

Documentation for ADTypes.jl.

ADTypes.ADTypes — Module
ADTypes.jl

ADTypes.jl is a multi-valued logic system to choose an automatic differentiation (AD) package and specify its parameters.

Dense AD

Forward mode

Algorithmic differentiation:

ADTypes.AutoForwardDiff — Type
AutoForwardDiff{chunksize,T}

Struct used to select the ForwardDiff.jl backend for automatic differentiation.

Defined by ADTypes.jl.

Constructors

AutoForwardDiff(; chunksize=nothing, tag=nothing)

Type parameters

  • chunksize: the preferred chunk size to evaluate several derivatives at once

Fields

  • tag::T: a custom tag to handle nested differentiation calls (usually not necessary)
ADTypes.AutoPolyesterForwardDiff — Type
AutoPolyesterForwardDiff{chunksize,T}

Struct used to select the PolyesterForwardDiff.jl backend for automatic differentiation.

Defined by ADTypes.jl.

Constructors

AutoPolyesterForwardDiff(; chunksize=nothing, tag=nothing)

Type parameters

  • chunksize: the preferred chunk size to evaluate several derivatives at once

Fields

  • tag::T: a custom tag to handle nested differentiation calls (usually not necessary)

Finite differences:

ADTypes.AutoFiniteDiff — Type
AutoFiniteDiff{T1,T2,T3}

Struct used to select the FiniteDiff.jl backend for automatic differentiation.

Defined by ADTypes.jl.

Constructors

AutoFiniteDiff(; fdtype=Val(:forward), fdjtype=fdtype, fdhtype=Val(:hcentral))

Fields

  • fdtype::T1: finite difference type
  • fdjtype::T2: finite difference type for the Jacobian
  • fdhtype::T3: finite difference type for the Hessian

Reverse mode

ADTypes.AutoReverseDiff — Type
AutoReverseDiff{compile}

Struct used to select the ReverseDiff.jl backend for automatic differentiation.

Defined by ADTypes.jl.

Constructors

AutoReverseDiff(; compile::Union{Val, Bool} = Val(false))

Fields

  • compile::Union{Val, Bool}: whether to compile the tape prior to differentiation (the boolean version is also the type parameter)
ADTypes.AutoTapir — Type
AutoTapir

Struct used to select the Tapir.jl backend for automatic differentiation.

Defined by ADTypes.jl.

Constructors

AutoTapir(; safe_mode=true)

Fields

  • safe_mode::Bool: whether to run additional checks to catch errors early. While this is on by default to ensure that users are aware of this option, you should generally turn it off for actual use, as it has substantial performance implications. If you encounter a problem with using Tapir (it fails to differentiate a function, or something truly nasty like a segfault occurs), then you should try switching safe_mode on and look at what happens. Often errors are caught earlier and the error messages are more useful.

Forward or reverse mode

ADTypes.AutoEnzyme — Type
AutoEnzyme{M,constant_function}

Struct used to select the Enzyme.jl backend for automatic differentiation.

Defined by ADTypes.jl.

Constructors

AutoEnzyme(; mode=nothing, constant_function::Bool=false)

The constant_function keyword argument (and type parameter) determines whether the function object itself should be considered constant or not during differentiation with Enzyme.jl. For simple functions, constant_function should usually be set to true, which leads to increased performance. However, in the case of closures or callable structs which contain differentiated data, constant_function should be set to false to ensure correctness (more details below).

Fields

  • mode::M: can be either

    • an object subtyping EnzymeCore.Mode (like EnzymeCore.Forward or EnzymeCore.Reverse) if a specific mode is required
    • nothing to choose the best mode automatically

Notes

We now give several examples of functions. For each one, we explain how constant_function should be set in order to compute the correct derivative with respect to the input x.

function f1(x)
    return x[1]
end

The function f1 is not a closure, it does not contain any data. Thus f1 can be differentiated with AutoEnzyme(constant_function=true) (although here setting constant_function=false would change neither correctness nor performance).

parameter = [0.0]
function f2(x)
    return parameter[1] + x[1]
end

The function f2 is a closure over parameter, but parameter is never modified based on the input x. Thus, f2 can be differentiated with AutoEnzyme(constant_function=true) (setting constant_function=false would not change correctness but would hinder performance).

cache = [0.0]
function f3(x)
    cache[1] = x[1]
    return cache[1] + x[1]
end

The function f3 is a closure over cache, and cache is modified based on the input x. That means cache cannot be treated as constant, since derivative values must be propagated through it. Thus f3 must be differentiated with AutoEnzyme(constant_function=false) (setting constant_function=true would make the result incorrect).

Symbolic mode

Sparse AD

ADTypes.AutoSparse — Type
AutoSparse{D,S,C}

Wraps an ADTypes.jl object to deal with sparse Jacobians and Hessians.

Fields

Constructors

AutoSparse(
    dense_ad;
    sparsity_detector=ADTypes.NoSparsityDetector(),
    coloring_algorithm=ADTypes.NoColoringAlgorithm()
)
ADTypes.dense_ad — Function
dense_ad(ad::AutoSparse)::AbstractADType

Return the underlying AD package for a sparse AD choice.

See also

Sparsity detector

ADTypes.jacobian_sparsity — Function
jacobian_sparsity(f, x, sd::AbstractSparsityDetector)::AbstractMatrix{Bool}
jacobian_sparsity(f!, y, x, sd::AbstractSparsityDetector)::AbstractMatrix{Bool}

Use detector sd to construct a (typically sparse) matrix S describing the pattern of nonzeroes in the Jacobian of f (resp. f!) applied at x (resp. (y, x)).

ADTypes.hessian_sparsity — Function
hessian_sparsity(f, x, sd::AbstractSparsityDetector)::AbstractMatrix{Bool}

Use detector sd to construct a (typically sparse) matrix S describing the pattern of nonzeroes in the Hessian of f applied at x.

Coloring algorithm

ADTypes.AbstractColoringAlgorithm — Type
AbstractColoringAlgorithm

Abstract supertype for Jacobian/Hessian coloring algorithms, defined for example in SparseDiffTools.jl.

Required methods

Note

The terminology and definitions are taken from the following paper:

"What Color Is Your Jacobian? Graph Coloring for Computing Derivatives"

Assefaw Hadish Gebremedhin, Fredrik Manne, and Alex Pothen (2005)

https://epubs.siam.org/doi/10.1137/S0036144504444711

ADTypes.column_coloring — Function
column_coloring(M::AbstractMatrix, ca::ColoringAlgorithm)::AbstractVector{<:Integer}

Use algorithm ca to construct a structurally orthogonal partition of the columns of M.

The result is a coloring vector c of length size(M, 2) such that for every non-zero coefficient M[i, j], column j is the only column of its color c[j] with a non-zero coefficient in row i.

ADTypes.row_coloring — Function
row_coloring(M::AbstractMatrix, ca::ColoringAlgorithm)::AbstractVector{<:Integer}

Use algorithm ca to construct a structurally orthogonal partition of the rows of M.

The result is a coloring vector c of length size(M, 1) such that for every non-zero coefficient M[i, j], row i is the only row of its color c[i] with a non-zero coefficient in column j.

ADTypes.symmetric_coloring — Function
symmetric_coloring(M::AbstractMatrix, ca::ColoringAlgorithm)::AbstractVector{<:Integer}

Use algorithm ca to construct a symmetrically structurally orthogonal partition of the columns (or rows) of the symmetric matrix M.

The result is a coloring vector c of length size(M, 1) == size(M, 2) such that for every non-zero coefficient M[i, j], at least one of the following conditions holds:

  • column j is the only column of its color c[j] with a non-zero coefficient in row i;
  • column i is the only column of its color c[i] with a non-zero coefficient in row j.

Modes

ADTypes.ForwardMode — Type
ForwardMode

Trait for AD choices that rely on forward mode algorithmic differentiation or finite differences.

These two paradigms are classified together because they can both efficiently compute Jacobian-vector products.

Miscellaneous

ADTypes.Auto — Function
ADTypes.Auto(package::Symbol)

A shortcut that converts an AD package name into an instance of AbstractADType, with all parameters set to their default values.

Warning

This function is type-unstable by design and might lead to suboptimal performance. In most cases, you should never need it: use the individual backend types directly.

Example

import ADTypes
backend = ADTypes.Auto(:Zygote)

# output

ADTypes.AutoZygote()