abstract type ABI

Abstract type for what ABI will be used.

Active(x)

Mark a function argument x of autodiff as active, Enzyme will auto-differentiate in respect Active arguments.

abstract type Annotation{T}

Abstract type for autodiff function argument wrappers like Const, Active and Duplicated.

BatchDuplicated(x, ∂f_∂xs)

Like Duplicated, except contains several shadows to compute derivatives for all at once. Argument ∂f_∂xs should be a tuple of the several values of type x.

BatchDuplicatedNoNeed(x, ∂f_∂xs)

Like DuplicatedNoNeed, except contains several shadows to compute derivatives for all at once. Argument ∂f_∂xs should be a tuple of the several values of type x.

BatchMixedDuplicated(x, ∂f_∂xs)

Like MixedDuplicated, except contains several shadows to compute derivatives for all at once. Only used within custom rules.

Const(x)

Mark a function argument x of autodiff as constant, Enzyme will not auto-differentiate in respect Const arguments.

Duplicated(x, ∂f_∂x)

Mark a function argument x of autodiff as duplicated, Enzyme will auto-differentiate in respect to such arguments, with dx acting as an accumulator for gradients (so $\partial f / \partial x$ will be added to) ∂f_∂x.

DuplicatedNoNeed(x, ∂f_∂x)

Like Duplicated, except also specifies that Enzyme may avoid computing the original result and only compute the derivative values.

struct FFIABI <: ABI

Foreign function call ABI. JIT the differentiated function, then inttoptr call the address.

struct Forward <: Mode

Forward mode differentiation

struct InlineABI <: ABI

Inlining function call ABI.

MixedDuplicated(x, ∂f_∂x)

Like Duplicated, except x may contain both active [immutable] and duplicated [mutable] data which is differentiable. Only used within custom rules.

abstract type Mode

Abstract type for what differentiation mode will be used.

struct NonGenABI <: ABI

Non-generated function ABI.

struct ReverseMode{ReturnPrimal,ABI,Holomorphic} <: Mode{ABI}

Reverse mode differentiation.

• ReturnPrimal: Should Enzyme return the primal return value from the augmented-forward.
• ABI: What runtime ABI to use
• Holomorphic: Whether the complex result function is holomorphic and we should compute d/dz

struct ReverseModeSplit{ReturnPrimal,ReturnShadow,Width,ModifiedBetween,ABI} <: Mode{ABI}

Reverse mode differentiation.

• ReturnPrimal: Should Enzyme return the primal return value from the augmented-forward.
• ReturnShadow: Should Enzyme return the shadow return value from the augmented-forward.
• Width: Batch Size (0 if to be automatically derived)
• ModifiedBetween: Tuple of each argument's modified between state (true if to be automatically derived).

compiler_job_from_backend(::KernelAbstractions.Backend, F::Type, TT:Type)::GPUCompiler.CompilerJob

Returns a GPUCompiler CompilerJob from a backend as specified by the first argument to the function.

For example, in CUDA one would do:

function EnzymeCore.compiler_job_from_backend(::CUDABackend, @nospecialize(F::Type), @nospecialize(TT::Type))
    mi = GPUCompiler.methodinstance(F, TT)
    return GPUCompiler.CompilerJob(mi, CUDA.compiler_config(CUDA.device()))
end

make_zero(::Type{T}, seen::IdDict, prev::T, ::Val{copy_if_inactive}=Val(false))::T

Recursively make a zero'd copy of the value `prev` of type `T`. The argument `copy_if_inactive` specifies
what to do if the type `T` is guaranteed to be inactive, use the primal (the default) or still copy the value.

make_zero!(val::T, seen::IdSet{Any}=IdSet())::Nothing

Recursively set a variables differentiable fields to zero. Only applicable for mutable types `T`.

make_zero(prev::T)

Helper function to recursively make zero.

AugmentedReturn(primal, shadow, tape)

Augment the primal return value of a function with its shadow, as well as any additional information needed to correctly compute the reverse pass, stored in tape.

Unless specified by the config that a variable is not overwritten, rules must assume any arrays/data structures/etc are overwritten between the forward and the reverse pass. Any floats or variables passed by value are always preserved as is (as are the arrays themselves, just not necessarily the values in the array).

See also augmented_primal.

Config{NeedsPrimal, NeedsShadow, Width, Overwritten}
ConfigWidth{Width} = Config{<:Any,<:Any, Width}

Configuration type to dispatch on in custom reverse rules (see augmented_primal and reverse).

• NeedsPrimal and NeedsShadow: boolean values specifying whether the primal and shadow (resp.) should be returned.
• Width: an integer that specifies the number of adjoints/shadows simultaneously being propagated.
• Overwritten: a tuple of booleans of whether each argument (including the function itself) is modified between the forward and reverse pass (true if potentially modified between).

Getters for the four type parameters are provided by needs_primal, needs_shadow, width, and overwritten.

augmented_primal(::Config, func::Annotation{typeof(f)}, RT::Type{<:Annotation}, args::Annotation...)

Must return an AugmentedReturn type.

• The primal must be the same type of the original return if needs_primal(config), otherwise nothing.
• The shadow must be nothing if needs_shadow(config) is false. If width is 1, the shadow should be the same type of the original return. If the width is greater than 1, the shadow should be NTuple{original return, width}.
• The tape can be any type (including Nothing) and is preserved for the reverse call.

forward(func::Annotation{typeof(f)}, RT::Type{<:Annotation}, args::Annotation...)

Calculate the forward derivative. The first argument func is the callable for which the rule applies to. Either wrapped in a Const), or a Duplicated if it is a closure. The second argument is the return type annotation, and all other arguments are the annotated function arguments.

inactive(func::typeof(f), args...)

Mark a particular function as always being inactive in both its return result and the function call itself.

inactive_noinl(func::typeof(f), args...)

Mark a particular function as always being inactive in both its return result and the function call itself, but do not prevent inlining of the function.

inactive_type(::Type{Ty})

Mark a particular type Ty as always being inactive.

noalias(func::typeof(f), args...)

Mark a particular function as always being a fresh allocation which does not alias any other accessible memory.

reverse(::Config, func::Annotation{typeof(f)}, dret::Active, tape, args::Annotation...)
reverse(::Config, func::Annotation{typeof(f)}, ::Type{<:Annotation), tape, args::Annotation...)

Takes gradient of derivative, activity annotation, and tape. If there is an active return dret is passed as Active{T} with the derivative of the active return val. Otherwise dret is passed as Type{Duplicated{T}}, etc.