API

DifferentiationInterface

An interface to various automatic differentiation backends in Julia.

Context

Abstract supertype for additional context arguments, which can be passed to differentiation operators after the active input x but are not differentiated.

Subtypes

• Constant
• Cache
• ConstantOrCache

Constant

Concrete type of Context argument which is kept constant during differentiation.

Note that an operator can be prepared with an arbitrary value of the constant. However, same-point preparation must occur with the exact value that will be reused later.

Example

julia> using DifferentiationInterface

julia> import ForwardDiff

julia> f(x, c) = c * sum(abs2, x);

julia> gradient(f, AutoForwardDiff(), [1.0, 2.0], Constant(10))
2-element Vector{Float64}:
 20.0
 40.0

julia> gradient(f, AutoForwardDiff(), [1.0, 2.0], Constant(100))
2-element Vector{Float64}:
 200.0
 400.0

Cache

Concrete type of Context argument which can be mutated with active values during differentiation.

The initial values present inside the cache do not matter.

For some backends, preparation allocates the required memory for Cache contexts with the right element type, similar to PreallocationTools.jl.

Example

julia> using DifferentiationInterface

julia> import ForwardDiff

julia> f(x, c) = sum(copyto!(c, x));

julia> prep = prepare_gradient(f, AutoForwardDiff(), [1.0, 2.0], Cache(zeros(2)));

julia> gradient(f, prep, AutoForwardDiff(), [3.0, 4.0], Cache(zeros(2)))
2-element Vector{Float64}:
 1.0
 1.0

ConstantOrCache

Concrete type of Context argument which can contain a mixture of constants and caches, passed along to the backend without modification.

Unlike for Cache, it is up to the user to ensure that the internal storage can adapt to the required element types, for instance by using PreallocationTools.jl directly.

prepare_pushforward(f,     backend, x, tx, [contexts...]; strict=Val(false)) -> prep
prepare_pushforward(f!, y, backend, x, tx, [contexts...]; strict=Val(false)) -> prep

Create a prep object that can be given to pushforward and its variants to speed them up.

Depending on the backend, this can have several effects (preallocating memory, recording an execution trace) which are transparent to the user.

For in-place functions, y is mutated by f! during preparation.

When strict=Val(true), type checking is enforced between preparation and execution (but size checking is left to the user).

prepare_pushforward_same_point(f,     backend, x, tx, [contexts...]; strict=Val(false)) -> prep_same
prepare_pushforward_same_point(f!, y, backend, x, tx, [contexts...]; strict=Val(false)) -> prep_same

Create a prep object that can be given to pushforward and its variants to speed them up, if they are applied at the same point x and with the same contexts.

Depending on the backend, this can have several effects (preallocating memory, recording an execution trace) which are transparent to the user.

For in-place functions, y is mutated by f! during preparation.

When strict=Val(true), type checking is enforced between preparation and execution (but size checking is left to the user).

pushforward(f,     [prep,] backend, x, tx, [contexts...]) -> ty
pushforward(f!, y, [prep,] backend, x, tx, [contexts...]) -> ty

Compute the pushforward of the function f at point x with a tuple of tangents tx.

To improve performance via operator preparation, refer to prepare_pushforward and prepare_pushforward_same_point.

pushforward!(f,     dy, [prep,] backend, x, tx, [contexts...]) -> ty
pushforward!(f!, y, dy, [prep,] backend, x, tx, [contexts...]) -> ty

Compute the pushforward of the function f at point x with a tuple of tangents tx, overwriting ty.

To improve performance via operator preparation, refer to prepare_pushforward and prepare_pushforward_same_point.

value_and_pushforward(f,     [prep,] backend, x, tx, [contexts...]) -> (y, ty)
value_and_pushforward(f!, y, [prep,] backend, x, tx, [contexts...]) -> (y, ty)

Compute the value and the pushforward of the function f at point x with a tuple of tangents tx.

To improve performance via operator preparation, refer to prepare_pushforward and prepare_pushforward_same_point.

value_and_pushforward!(f,     dy, [prep,] backend, x, tx, [contexts...]) -> (y, ty)
value_and_pushforward!(f!, y, dy, [prep,] backend, x, tx, [contexts...]) -> (y, ty)

Compute the value and the pushforward of the function f at point x with a tuple of tangents tx, overwriting ty.

To improve performance via operator preparation, refer to prepare_pushforward and prepare_pushforward_same_point.

prepare_pullback(f,     backend, x, ty, [contexts...]; strict=Val(false)) -> prep
prepare_pullback(f!, y, backend, x, ty, [contexts...]; strict=Val(false)) -> prep

Create a prep object that can be given to pullback and its variants to speed them up.

Depending on the backend, this can have several effects (preallocating memory, recording an execution trace) which are transparent to the user.

For in-place functions, y is mutated by f! during preparation.

When strict=Val(true), type checking is enforced between preparation and execution (but size checking is left to the user).

prepare_pullback_same_point(f,     backend, x, ty, [contexts...]; strict=Val(false)) -> prep_same
prepare_pullback_same_point(f!, y, backend, x, ty, [contexts...]; strict=Val(false)) -> prep_same

Create a prep object that can be given to pullback and its variants to speed them up, if they are applied at the same point x and with the same contexts.

Depending on the backend, this can have several effects (preallocating memory, recording an execution trace) which are transparent to the user.

For in-place functions, y is mutated by f! during preparation.

When strict=Val(true), type checking is enforced between preparation and execution (but size checking is left to the user).

pullback(f,     [prep,] backend, x, ty, [contexts...]) -> tx
pullback(f!, y, [prep,] backend, x, ty, [contexts...]) -> tx

Compute the pullback of the function f at point x with a tuple of tangents ty.

To improve performance via operator preparation, refer to prepare_pullback and prepare_pullback_same_point.

pullback!(f,     dx, [prep,] backend, x, ty, [contexts...]) -> tx
pullback!(f!, y, dx, [prep,] backend, x, ty, [contexts...]) -> tx

Compute the pullback of the function f at point x with a tuple of tangents ty, overwriting dx.

To improve performance via operator preparation, refer to prepare_pullback and prepare_pullback_same_point.

value_and_pullback(f,     [prep,] backend, x, ty, [contexts...]) -> (y, tx)
value_and_pullback(f!, y, [prep,] backend, x, ty, [contexts...]) -> (y, tx)

Compute the value and the pullback of the function f at point x with a tuple of tangents ty.

To improve performance via operator preparation, refer to prepare_pullback and prepare_pullback_same_point.

value_and_pullback!(f,     dx, [prep,] backend, x, ty, [contexts...]) -> (y, tx)
value_and_pullback!(f!, y, dx, [prep,] backend, x, ty, [contexts...]) -> (y, tx)

Compute the value and the pullback of the function f at point x with a tuple of tangents ty, overwriting dx.

To improve performance via operator preparation, refer to prepare_pullback and prepare_pullback_same_point.

prepare_derivative(f,     backend, x, [contexts...]; strict=Val(false)) -> prep
prepare_derivative(f!, y, backend, x, [contexts...]; strict=Val(false)) -> prep

Create a prep object that can be given to derivative and its variants to speed them up.

Depending on the backend, this can have several effects (preallocating memory, recording an execution trace) which are transparent to the user.

For in-place functions, y is mutated by f! during preparation.

When strict=Val(true), type checking is enforced between preparation and execution (but size checking is left to the user).

derivative(f,     [prep,] backend, x, [contexts...]) -> der
derivative(f!, y, [prep,] backend, x, [contexts...]) -> der

Compute the derivative of the function f at point x.

To improve performance via operator preparation, refer to prepare_derivative.

derivative!(f,     der, [prep,] backend, x, [contexts...]) -> der
derivative!(f!, y, der, [prep,] backend, x, [contexts...]) -> der

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

To improve performance via operator preparation, refer to prepare_derivative.

value_and_derivative(f,     [prep,] backend, x, [contexts...]) -> (y, der)
value_and_derivative(f!, y, [prep,] backend, x, [contexts...]) -> (y, der)

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

To improve performance via operator preparation, refer to prepare_derivative.

value_and_derivative!(f,     der, [prep,] backend, x, [contexts...]) -> (y, der)
value_and_derivative!(f!, y, der, [prep,] backend, x, [contexts...]) -> (y, der)

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

To improve performance via operator preparation, refer to prepare_derivative.

prepare_gradient(f, backend, x, [contexts...]; strict=Val(false)) -> prep

Create a prep object that can be given to gradient and its variants to speed them up.

Depending on the backend, this can have several effects (preallocating memory, recording an execution trace) which are transparent to the user.

When strict=Val(true), type checking is enforced between preparation and execution (but size checking is left to the user).

gradient(f, [prep,] backend, x, [contexts...]) -> grad

Compute the gradient of the function f at point x.

To improve performance via operator preparation, refer to prepare_gradient.

gradient!(f, grad, [prep,] backend, x, [contexts...]) -> grad

Compute the gradient of the function f at point x, overwriting grad.

To improve performance via operator preparation, refer to prepare_gradient.

value_and_gradient(f, [prep,] backend, x, [contexts...]) -> (y, grad)

Compute the value and the gradient of the function f at point x.

To improve performance via operator preparation, refer to prepare_gradient.

value_and_gradient!(f, grad, [prep,] backend, x, [contexts...]) -> (y, grad)

Compute the value and the gradient of the function f at point x, overwriting grad.

To improve performance via operator preparation, refer to prepare_gradient.

prepare_jacobian(f,     backend, x, [contexts...]; strict=Val(false)) -> prep
prepare_jacobian(f!, y, backend, x, [contexts...]; strict=Val(false)) -> prep

Create a prep object that can be given to jacobian and its variants to speed them up.

Depending on the backend, this can have several effects (preallocating memory, recording an execution trace) which are transparent to the user.

For in-place functions, y is mutated by f! during preparation.

When strict=Val(true), type checking is enforced between preparation and execution (but size checking is left to the user).

jacobian(f,     [prep,] backend, x, [contexts...]) -> jac
jacobian(f!, y, [prep,] backend, x, [contexts...]) -> jac

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

To improve performance via operator preparation, refer to prepare_jacobian.

jacobian!(f,     jac, [prep,] backend, x, [contexts...]) -> jac
jacobian!(f!, y, jac, [prep,] backend, x, [contexts...]) -> jac

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

To improve performance via operator preparation, refer to prepare_jacobian.

value_and_jacobian(f,     [prep,] backend, x, [contexts...]) -> (y, jac)
value_and_jacobian(f!, y, [prep,] backend, x, [contexts...]) -> (y, jac)

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

To improve performance via operator preparation, refer to prepare_jacobian.

value_and_jacobian!(f,     jac, [prep,] backend, x, [contexts...]) -> (y, jac)
value_and_jacobian!(f!, y, jac, [prep,] backend, x, [contexts...]) -> (y, jac)

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

To improve performance via operator preparation, refer to prepare_jacobian.

SecondOrder

Combination of two backends for second-order differentiation.

Constructor

SecondOrder(outer_backend, inner_backend)

Fields

• outer::AbstractADType: backend for the outer differentiation
• inner::AbstractADType: backend for the inner differentiation

prepare_second_derivative(f, backend, x, [contexts...]; strict=Val(false)) -> prep

Create a prep object that can be given to second_derivative and its variants to speed them up.

Depending on the backend, this can have several effects (preallocating memory, recording an execution trace) which are transparent to the user.

When strict=Val(true), type checking is enforced between preparation and execution (but size checking is left to the user).

second_derivative(f, [prep,] backend, x, [contexts...]) -> der2

Compute the second derivative of the function f at point x.

To improve performance via operator preparation, refer to prepare_second_derivative.

second_derivative!(f, der2, [prep,] backend, x, [contexts...]) -> der2

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

To improve performance via operator preparation, refer to prepare_second_derivative.

value_derivative_and_second_derivative(f, [prep,] backend, x, [contexts...]) -> (y, der, der2)

Compute the value, first derivative and second derivative of the function f at point x.

To improve performance via operator preparation, refer to prepare_second_derivative.

value_derivative_and_second_derivative!(f, der, der2, [prep,] backend, x, [contexts...]) -> (y, der, der2)

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

To improve performance via operator preparation, refer to prepare_second_derivative.

prepare_hvp(f, backend, x, tx, [contexts...]; strict=Val(false)) -> prep

Create a prep object that can be given to hvp and its variants to speed them up.

Depending on the backend, this can have several effects (preallocating memory, recording an execution trace) which are transparent to the user.

When strict=Val(true), type checking is enforced between preparation and execution (but size checking is left to the user).

prepare_hvp_same_point(f, backend, x, tx, [contexts...]; strict=Val(false)) -> prep_same

Create a prep object that can be given to hvp and its variants to speed them up, if they are applied at the same point x and with the same contexts.

Depending on the backend, this can have several effects (preallocating memory, recording an execution trace) which are transparent to the user.

When strict=Val(true), type checking is enforced between preparation and execution (but size checking is left to the user).

hvp(f, [prep,] backend, x, tx, [contexts...]) -> tg

Compute the Hessian-vector product of f at point x with a tuple of tangents tx.

To improve performance via operator preparation, refer to prepare_hvp and prepare_hvp_same_point.

hvp!(f, tg, [prep,] backend, x, tx, [contexts...]) -> tg

Compute the Hessian-vector product of f at point x with a tuple of tangents tx, overwriting tg.

To improve performance via operator preparation, refer to prepare_hvp and prepare_hvp_same_point.

gradient_and_hvp(f, [prep,] backend, x, tx, [contexts...]) -> (grad, tg)

Compute the gradient and the Hessian-vector product of f at point x with a tuple of tangents tx.

To improve performance via operator preparation, refer to prepare_hvp and prepare_hvp_same_point.

gradient_and_hvp!(f, grad, tg, [prep,] backend, x, tx, [contexts...]) -> (grad, tg)

Compute the gradient and the Hessian-vector product of f at point x with a tuple of tangents tx, overwriting grad and tg.

To improve performance via operator preparation, refer to prepare_hvp and prepare_hvp_same_point.

prepare_hessian(f, backend, x, [contexts...]; strict=Val(false)) -> prep

Create a prep object that can be given to hessian and its variants to speed them up.

Depending on the backend, this can have several effects (preallocating memory, recording an execution trace) which are transparent to the user.

When strict=Val(true), type checking is enforced between preparation and execution (but size checking is left to the user).

hessian(f, [prep,] backend, x, [contexts...]) -> hess

Compute the Hessian matrix of the function f at point x.

To improve performance via operator preparation, refer to prepare_hessian.

hessian!(f, hess, [prep,] backend, x, [contexts...]) -> hess

Compute the Hessian matrix of the function f at point x, overwriting hess.

To improve performance via operator preparation, refer to prepare_hessian.

value_gradient_and_hessian(f, [prep,] backend, x, [contexts...]) -> (y, grad, hess)

Compute the value, gradient vector and Hessian matrix of the function f at point x.

To improve performance via operator preparation, refer to prepare_hessian.

value_gradient_and_hessian!(f, grad, hess, [prep,] backend, x, [contexts...]) -> (y, grad, hess)

Compute the value, gradient vector and Hessian matrix of the function f at point x, overwriting grad and hess.

To improve performance via operator preparation, refer to prepare_hessian.

check_available(backend)

Check whether backend is available (i.e. whether the extension is loaded) and return a Bool.

check_inplace(backend)

Check whether backend supports differentiation of in-place functions and return a Bool.

outer(backend::SecondOrder)
outer(backend::AbstractADType)

Return the outer backend of a SecondOrder object, tasked with differentiation at the second order.

For any other backend type, this function acts like the identity.

inner(backend::SecondOrder)
inner(backend::AbstractADType)

Return the inner backend of a SecondOrder object, tasked with differentiation at the first order.

For any other backend type, this function acts like the identity.

DifferentiateWith

Function wrapper that enforces differentiation with a "substitute" AD backend, possible different from the "true" AD backend that is called.

For instance, suppose a function f is not differentiable with Zygote because it involves mutation, but you know that it is differentiable with Enzyme. Then f2 = DifferentiateWith(f, AutoEnzyme()) is a new function that behaves like f, except that f2 is differentiable with Zygote (thanks to a chain rule which calls Enzyme under the hood). Moreover, any larger algorithm alg that calls f2 instead of f will also be differentiable with Zygote (as long as f was the only Zygote blocker).

Fields

• f: the function in question, with signature f(x)
• backend::AbstractADType: the substitute backend to use for differentiation

Constructor

DifferentiateWith(f, backend)

Example

julia> using DifferentiationInterface

julia> import FiniteDiff, ForwardDiff, Zygote

julia> function f(x::Vector{Float64})
           a = Vector{Float64}(undef, 1)  # type constraint breaks ForwardDiff
           a[1] = sum(abs2, x)  # mutation breaks Zygote
           return a[1]
       end;

julia> f2 = DifferentiateWith(f, AutoFiniteDiff());

julia> f([3.0, 5.0]) == f2([3.0, 5.0])
true

julia> alg(x) = 7 * f2(x);

julia> ForwardDiff.gradient(alg, [3.0, 5.0])
2-element Vector{Float64}:
 42.0
 70.0

julia> Zygote.gradient(alg, [3.0, 5.0])[1]
2-element Vector{Float64}:
 42.0
 70.0

MixedMode

Combination of a forward and a reverse mode backend for mixed-mode sparse Jacobian computation.

Constructor

MixedMode(forward_backend, reverse_backend)

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). This process can be very slow, and should only be used if its output can be exploited multiple times to compute many sparse matrices.

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  ⋅

AutoForwardFromPrimitive

Wrapper which forces a given backend to act as a reverse-mode backend.

Used in internal testing.

AutoReverseFromPrimitive

Wrapper which forces a given backend to act as a reverse-mode backend.

Used in internal testing.

AutoSimpleFiniteDiff <: ADTypes.AbstractADType

Forward mode backend based on the finite difference (f(x + ε) - f(x)) / ε, with artificial chunk size to mimick ForwardDiff.

Constructor

AutoSimpleFiniteDiff(ε=1e-5; chunksize=nothing)

AutoZeroForward <: ADTypes.AbstractADType

Trivial backend that sets all derivatives to zero. Used in testing and benchmarking.

AutoZeroReverse <: ADTypes.AbstractADType

Trivial backend that sets all derivatives to zero. Used in testing and benchmarking.

BatchSizeSettings{B,singlebatch,aligned}

Configuration for the batch size deduced from a backend and a sample array of length N.

Type parameters

• B::Int: batch size
• singlebatch::Bool: whether B == N (B > N is not allowed)
• aligned::Bool: whether N % B == 0

Fields

• N::Int: array length
• A::Int: number of batches A = div(N, B, RoundUp)
• B_last::Int: size of the last batch (if aligned is false)

DerivativePrep

Abstract type for additional information needed by derivative and its variants.

DontPrepareInner

Trait identifying outer backends for which the inner backend in second-order autodiff should not be prepared at all.

FixTail

Closure around a function f and a set of tail argument tail_args such that

(ft::FixTail)(args...) = ft.f(args..., ft.tail_args...)

ForwardAndReverseMode <: ADTypes.AbstractMode

Appropriate mode type for MixedMode backends.

ForwardOverForward

Traits identifying second-order backends that compute HVPs in forward over forward mode (inefficient).

ForwardOverReverse

Traits identifying second-order backends that compute HVPs in forward over reverse mode.

FunctionContext

Private type of Context argument used for passing functions inside second-order differentiation.

Behaves differently for Enzyme only, where the function can be annotated.

GradientPrep

Abstract type for additional information needed by gradient and its variants.

HVPPrep

Abstract type for additional information needed by hvp and its variants.

HessianPrep

Abstract type for additional information needed by hessian and its variants.

InPlaceNotSupported

Trait identifying backends that do not support in-place functions f!(y, x).

InPlaceSupported

Trait identifying backends that support in-place functions f!(y, x).

JacobianPrep

Abstract type for additional information needed by jacobian and its variants.

OneElement

Efficient storage for a one-hot array, aka an array in the standard Euclidean basis.

PrepareInnerOverload

Trait identifying outer backends for which the inner backend in second-order autodiff should be prepared with an overloaded input type.

PrepareInnerSimple

Trait identifying outer backends for which the inner backend in second-order autodiff should be prepared with the same input type.

PullbackFast

Trait identifying backends that support efficient pullbacks.

PullbackPrep

Abstract type for additional information needed by pullback and its variants.

PullbackSlow

Trait identifying backends that do not support efficient pullbacks.

PushforwardFast

Trait identifying backends that support efficient pushforwards.

PushforwardPrep

Abstract type for additional information needed by pushforward and its variants.

PushforwardSlow

Trait identifying backends that do not support efficient pushforwards.

ReverseOverForward

Traits identifying second-order backends that compute HVPs in reverse over forward mode.

ReverseOverReverse

Traits identifying second-order backends that compute HVPs in reverse over reverse mode.

Rewrap

Utility for recording context types of additional arguments (e.g. Constant or Cache) and re-wrapping them into their types after they have been unwrapped.

Useful for second-order differentiation.

SecondDerivativePrep

Abstract type for additional information needed by second_derivative and its variants.

mode(backend::SecondOrder)

Return the outer mode of the second-order backend.

basis(a::AbstractArray, i)

Construct the i-th standard basis array in the vector space of a.

fix_tail(f, tail_args...)

Convenience for constructing a FixTail, with a shortcut when there are no tail arguments.

forward_backend(m::MixedMode)

Return the forward-mode part of a MixedMode backend.

hessian_sparsity_with_contexts(f, detector, x, contexts...)

Wrapper around ADTypes.hessian_sparsity enabling the allocation of caches with proper element types.

hvp_mode(backend)

Return the best combination of modes for hvp and its variants, among the following options:

• ForwardOverForward
• ForwardOverReverse
• ReverseOverForward
• ReverseOverReverse

inner_preparation_behavior(backend::AbstractADType)

Return PrepareInnerSimple, PrepareInnerOverload or DontPrepareInner in a statically predictable way.

inplace_support(backend)

Return InPlaceSupported or InPlaceNotSupported in a statically predictable way.

ismutable_array(x)

Check whether x is a mutable array and return a Bool.

At the moment, this only returns false for StaticArrays.SArray.

jacobian_sparsity_with_contexts(f, detector, x, contexts...)
jacobian_sparsity_with_contexts(f!, y, detector, x, contexts...)

Wrapper around ADTypes.jacobian_sparsity enabling the allocation of caches with proper element types.

multibasis(a::AbstractArray, inds)

Construct the sum of the i-th standard basis arrays in the vector space of a for all i ∈ inds.

overloaded_input_type(prep)

If it exists, return the overloaded input type which will be passed to the differentiated function when preparation result prep is reused.

pick_batchsize(backend, x_or_y::AbstractArray)

Return a BatchSizeSettings appropriate for arrays of the same length as x_or_y with a given backend.

Note that the array in question can be either the input or the output of the function, depending on whether the backend performs forward- or reverse-mode AD.

pick_batchsize(backend, N::Integer)

Return a BatchSizeSettings appropriate for arrays of length N with a given backend.

prepare!_derivative(f,     prep, backend, x, [contexts...]) -> new_prep
prepare!_derivative(f!, y, prep, backend, x, [contexts...]) -> new_prep

Same behavior as prepare_derivative but can resize the contents of an existing prep object to avoid some allocations.

There is no guarantee that prep will be mutated, or that performance will be improved compared to preparation from scratch.

prepare!_gradient(f, prep, backend, x, [contexts...]) -> new_prep

Same behavior as prepare_gradient but can resize the contents of an existing prep object to avoid some allocations.

There is no guarantee that prep will be mutated, or that performance will be improved compared to preparation from scratch.

prepare!_hessian(f, backend, x, [contexts...]) -> new_prep

Same behavior as prepare_hessian but can resize the contents of an existing prep object to avoid some allocations.

There is no guarantee that prep will be mutated, or that performance will be improved compared to preparation from scratch.

prepare!_hvp(f, backend, x, tx, [contexts...]) -> new_prep

Create a prep object that can be given to hvp and its variants to speed them up.

Depending on the backend, this can have several effects (preallocating memory, recording an execution trace) which are transparent to the user.

When strict=Val(true), type checking is enforced between preparation and execution (but size checking is left to the user).

prepare!_jacobian(f,     prep, backend, x, [contexts...]) -> new_prep
prepare!_jacobian(f!, y, prep, backend, x, [contexts...]) -> new_prep

Same behavior as prepare_jacobian but can resize the contents of an existing prep object to avoid some allocations.

There is no guarantee that prep will be mutated, or that performance will be improved compared to preparation from scratch.

prepare!_pullback(f,     prep, backend, x, ty, [contexts...]) -> new_prep
prepare!_pullback(f!, y, prep, backend, x, ty, [contexts...]) -> new_prep

Same behavior as prepare_pullback but can resize the contents of an existing prep object to avoid some allocations.

There is no guarantee that prep will be mutated, or that performance will be improved compared to preparation from scratch.

prepare!_pushforward(f,     prep, backend, x, tx, [contexts...]) -> new_prep
prepare!_pushforward(f!, y, prep, backend, x, tx, [contexts...]) -> new_prep

Same behavior as prepare_pushforward but can resize the contents of an existing prep object to avoid some allocations.

There is no guarantee that prep will be mutated, or that performance will be improved compared to preparation from scratch.

prepare!_second_derivative(f, prep, backend, x, [contexts...]) -> new_prep

Same behavior as prepare_second_derivative but can resize the contents of an existing prep object to avoid some allocations.

There is no guarantee that prep will be mutated, or that performance will be improved compared to preparation from scratch.

pullback_performance(backend)

Return PullbackFast or PullbackSlow in a statically predictable way.

pushforward_performance(backend)

Return PushforwardFast or PushforwardSlow in a statically predictable way.

reasonable_batchsize(N::Integer, Bmax::Integer)

Reproduces the heuristic from ForwardDiff to minimize

1. the number of batches necessary to cover an array of length N
2. the number of leftover indices in the last partial batch

Source: https://github.com/JuliaDiff/ForwardDiff.jl/blob/ec74fbc32b10bbf60b3c527d8961666310733728/src/prelude.jl#L19-L29

recursive_similar(x, T)

Apply similar(_, T) recursively to x or its components.

Works if x is an AbstractArray or a (nested) NTuple / NamedTuple of AbstractArrays.

reverse_backend(m::MixedMode)

Return the reverse-mode part of a MixedMode backend.

threshold_batchsize(backend::AbstractADType, B::Integer)

If the backend object has a fixed batch size B0, return a new backend where the fixed batch size is min(B0, B). Otherwise, act as the identity.