GitHub

API

DifferentiationInterface — Module
DifferentiationInterface

An interface to various automatic differentiation backends in Julia.

Exports

source

First order

Pushforward

DifferentiationInterface.prepare_pushforward — Function
prepare_pushforward(f,     backend, x, dx) -> extras
prepare_pushforward(f!, y, backend, x, dx) -> extras

Create an extras object that can be given to pushforward and its variants.

Warning

If the function changes in any way, the result of preparation will be invalidated, and you will need to run it again. In the two-argument case, y is mutated by f! during preparation.

source
DifferentiationInterface.prepare_pushforward_same_point — Function
prepare_pushforward_same_point(f,     backend, x, dx) -> extras_same
prepare_pushforward_same_point(f!, y, backend, x, dx) -> extras_same

Create an extras_same object that can be given to pushforward and its variants if they are applied at the same point x.

Warning

If the function or the point changes in any way, the result of preparation will be invalidated, and you will need to run it again. In the two-argument case, y is mutated by f! during preparation.

source
DifferentiationInterface.pushforward — Function
pushforward(f,     backend, x, dx, [extras]) -> dy
pushforward(f!, y, backend, x, dx, [extras]) -> dy

Compute the pushforward of the function f at point x with seed dx.

source
DifferentiationInterface.pushforward! — Function
pushforward!(f,     dy, backend, x, dx, [extras]) -> dy
pushforward!(f!, y, dy, backend, x, dx, [extras]) -> dy

Compute the pushforward of the function f at point x with seed dx, overwriting dy.

source
DifferentiationInterface.value_and_pushforward — Function
value_and_pushforward(f,     backend, x, dx, [extras]) -> (y, dy)
value_and_pushforward(f!, y, backend, x, dx, [extras]) -> (y, dy)

Compute the value and the pushforward of the function f at point x with seed dx.

Info

Required primitive for forward mode backends.

source
DifferentiationInterface.value_and_pushforward! — Function
value_and_pushforward!(f,     dy, backend, x, dx, [extras]) -> (y, dy)
value_and_pushforward!(f!, y, dy, backend, x, dx, [extras]) -> (y, dy)

Compute the value and the pushforward of the function f at point x with seed dx, overwriting dy.

source

Pullback

DifferentiationInterface.prepare_pullback — Function
prepare_pullback(f,     backend, x, dy) -> extras
prepare_pullback(f!, y, backend, x, dy) -> extras

Create an extras object that can be given to pullback and its variants.

Warning

If the function changes in any way, the result of preparation will be invalidated, and you will need to run it again. In the two-argument case, y is mutated by f! during preparation.

source
DifferentiationInterface.prepare_pullback_same_point — Function
prepare_pullback_same_point(f,     backend, x, dy) -> extras_same
prepare_pullback_same_point(f!, y, backend, x, dy) -> extras_same

Create an extras_same object that can be given to pullback and its variants if they are applied at the same point x.

Warning

If the function or the point changes in any way, the result of preparation will be invalidated, and you will need to run it again. In the two-argument case, y is mutated by f! during preparation.

source
DifferentiationInterface.pullback — Function
pullback(f,     backend, x, dy, [extras]) -> dx
pullback(f!, y, backend, x, dy, [extras]) -> dx

Compute the pullback of the function f at point x with seed dy.

source
DifferentiationInterface.pullback! — Function
pullback!(f,     dx, backend, x, dy, [extras]) -> dx
pullback!(f!, y, dx, backend, x, dy, [extras]) -> dx

Compute the pullback of the function f at point x with seed dy, overwriting dx.

source
DifferentiationInterface.value_and_pullback — Function
value_and_pullback(f,     backend, x, dy, [extras]) -> (y, dx)
value_and_pullback(f!, y, backend, x, dy, [extras]) -> (y, dx)

Compute the value and the pullback of the function f at point x with seed dy.

Info

Required primitive for reverse mode backends.

source
DifferentiationInterface.value_and_pullback! — Function
value_and_pullback!(f,     dx, backend, x, dy, [extras]) -> (y, dx)
value_and_pullback!(f!, y, dx, backend, x, dy, [extras]) -> (y, dx)

Compute the value and the pullback of the function f at point x with seed dy, overwriting dx.

source

Derivative

DifferentiationInterface.prepare_derivative — Function
prepare_derivative(f,     backend, x) -> extras
prepare_derivative(f!, y, backend, x) -> extras

Create an extras object that can be given to derivative and its variants.

Warning

If the function changes in any way, the result of preparation will be invalidated, and you will need to run it again. In the two-argument case, y is mutated by f! during preparation.

source
DifferentiationInterface.derivative! — Function
derivative!(f,     der, backend, x, [extras]) -> der
derivative!(f!, y, der, backend, x, [extras]) -> der

Compute the derivative of the function f at point x, overwriting der.

source
DifferentiationInterface.value_and_derivative! — Function
value_and_derivative!(f,     der, backend, x, [extras]) -> (y, der)
value_and_derivative!(f!, y, der, backend, x, [extras]) -> (y, der)

Compute the value and the derivative of the function f at point x, overwriting der.

source

Gradient

DifferentiationInterface.prepare_gradient — Function
prepare_gradient(f, backend, x) -> extras

Create an extras object that can be given to gradient and its variants.

Warning

If the function changes in any way, the result of preparation will be invalidated, and you will need to run it again.

source

Jacobian

DifferentiationInterface.prepare_jacobian — Function
prepare_jacobian(f,     backend, x) -> extras
prepare_jacobian(f!, y, backend, x) -> extras

Create an extras object that can be given to jacobian and its variants.

Warning

If the function changes in any way, the result of preparation will be invalidated, and you will need to run it again. In the two-argument case, y is mutated by f! during preparation.

source
DifferentiationInterface.jacobian! — Function
jacobian!(f,     jac, backend, x, [extras]) -> jac
jacobian!(f!, y, jac, backend, x, [extras]) -> jac

Compute the Jacobian matrix of the function f at point x, overwriting jac.

source
DifferentiationInterface.value_and_jacobian — Function
value_and_jacobian(f,     backend, x, [extras]) -> (y, jac)
value_and_jacobian(f!, y, backend, x, [extras]) -> (y, jac)

Compute the value and the Jacobian matrix of the function f at point x.

source
DifferentiationInterface.value_and_jacobian! — Function
value_and_jacobian!(f,     jac, backend, x, [extras]) -> (y, jac)
value_and_jacobian!(f!, y, jac, backend, x, [extras]) -> (y, jac)

Compute the value and the Jacobian matrix of the function f at point x, overwriting jac.

source

Second order

DifferentiationInterface.SecondOrder — Type
SecondOrder

Combination of two backends for second-order differentiation.

Danger

SecondOrder backends do not support first-order operators.

Constructor

SecondOrder(outer_backend, inner_backend)

Fields

  • outer::ADTypes.AbstractADType: backend for the outer differentiation

  • inner::ADTypes.AbstractADType: backend for the inner differentiation

source

Second derivative

Hessian-vector product

DifferentiationInterface.prepare_hvp — Function
prepare_hvp(f, backend, x, dx) -> extras

Create an extras object that can be given to hvp and its variants.

Warning

If the function changes in any way, the result of preparation will be invalidated, and you will need to run it again.

source
DifferentiationInterface.prepare_hvp_same_point — Function
prepare_hvp_same_point(f, backend, x, dx) -> extras_same

Create an extras_same object that can be given to hvp and its variants if they are applied at the same point x.

Warning

If the function or the point changes in any way, the result of preparation will be invalidated, and you will need to run it again.

source

Hessian

DifferentiationInterface.prepare_hessian — Function
prepare_hessian(f, backend, x) -> extras

Create an extras object that can be given to hessian and its variants.

Warning

If the function changes in any way, the result of preparation will be invalidated, and you will need to run it again.

source

Utilities

Backend queries

DifferentiationInterface.check_hessian — Function
check_hessian(backend)

Check whether backend supports second order differentiation by trying to compute a hessian.

Warning

Might take a while due to compilation time.

source

Backend switch

DifferentiationInterface.DifferentiateWith — Type
DifferentiateWith

Callable function wrapper that enforces differentiation with a specified (inner) backend.

This works by defining new rules overriding the behavior of the outer backend that would normally be used.

Warning

This is an experimental functionality, whose API cannot yet be considered stable. At the moment, it only supports one-argument functions, and rules are only defined for ChainRules.jl-compatible outer backends.

Fields

  • f: the function in question
  • backend::AbstractADType: the inner backend to use for differentiation

Constructor

DifferentiateWith(f, backend)

Example

using DifferentiationInterface
import ForwardDiff, Zygote

function f(x)
    a = Vector{eltype(x)}(undef, 1)
    a[1] = sum(x)  # mutation that breaks Zygote
    return a[1]
end

dw = DifferentiateWith(f, AutoForwardDiff());

gradient(dw, AutoZygote(), [2.0])  # calls ForwardDiff instead

# output

1-element Vector{Float64}:
 1.0
source

Sparsity detection

DifferentiationInterface.DenseSparsityDetector — Type
DenseSparsityDetector

Sparsity pattern detector satisfying the detection API of ADTypes.jl.

The nonzeros in a Jacobian or Hessian are detected by computing the relevant matrix with dense AD, and thresholding the entries with a given tolerance (which can be numerically inaccurate).

Warning

This detector can be very slow, and should only be used if its output can be exploited multiple times to compute many sparse matrices.

Danger

In general, the sparsity pattern you obtain can depend on the provided input x. If you want to reuse the pattern, make sure that it is input-agnostic.

Fields

  • backend::AbstractADType is the dense AD backend used under the hood
  • atol::Float64 is the minimum magnitude of a matrix entry to be considered nonzero

Constructor

DenseSparsityDetector(backend; atol, method=:iterative)

The keyword argument method::Symbol can be either:

  • :iterative: compute the matrix in a sequence of matrix-vector products (memory-efficient)
  • :direct: compute the matrix all at once (memory-hungry but sometimes faster).

Note that the constructor is type-unstable because method ends up being a type parameter of the DenseSparsityDetector object (this is not part of the API and might change).

Examples

using ADTypes, DifferentiationInterface, SparseArrays
import ForwardDiff

detector = DenseSparsityDetector(AutoForwardDiff(); atol=1e-5, method=:direct)

ADTypes.jacobian_sparsity(diff, rand(5), detector)

# output

4×5 SparseMatrixCSC{Bool, Int64} with 8 stored entries:
 1  1  ⋅  ⋅  ⋅
 ⋅  1  1  ⋅  ⋅
 ⋅  ⋅  1  1  ⋅
 ⋅  ⋅  ⋅  1  1

Sometimes the sparsity pattern is input-dependent:

ADTypes.jacobian_sparsity(x -> [prod(x)], rand(2), detector)

# output

1×2 SparseMatrixCSC{Bool, Int64} with 2 stored entries:
 1  1
ADTypes.jacobian_sparsity(x -> [prod(x)], [0, 1], detector)

# output

1×2 SparseMatrixCSC{Bool, Int64} with 1 stored entry:
 1  ⋅
source

Internals

The following is not part of the public API.

DifferentiationInterface.Batch — Type
Batch{B,T}

Efficient storage for B elements of type T (NTuple wrapper).

A Batch can be used as seed to trigger batched-mode pushforward, pullback and hvp.

Fields

  • elements::NTuple{B,T}
source
ADTypes.mode — Method
mode(backend::SecondOrder)

Return the outer mode of the second-order backend.

source
DifferentiationInterface.basis — Method
basis(backend, a::AbstractArray, i::CartesianIndex)

Construct the i-th stardard basis array in the vector space of a with element type eltype(a).

Note

If an AD backend benefits from a more specialized basis array implementation, this function can be extended on the backend type.

source
DifferentiationInterface.nested — Method
nested(backend)

Return a possibly modified backend that can work while nested inside another differentiation procedure.

At the moment, this is only useful for Enzyme, which needs autodiff_deferred to be compatible with higher-order differentiation.

source
DifferentiationInterface.pick_batchsize — Method
pick_batchsize(backend::AbstractADType, dimension::Integer)

Pick a reasonable batch size for batched derivative evaluation with a given total dimension.

Returns 1 for backends which have not overloaded it.

source