Internal API

Private

These functions are not part of the public API and are subject to change at any time.

Reactant.REDUB_ARGUMENTS_NAME Constant

Reactant.REDUB_ARGUMENTS_NAME

The variable name bound to call_with_reactant's tuple of arguments in its @generated method definition.

This binding can be used to manually reference/destructure call_with_reactants arguments

This is required because user arguments could have a name which clashes with whatever name we choose for our argument. Thus we gensym to create it.

This originates from https://github.com/JuliaLabs/Cassette.jl/blob/c29b237c1ec0deda3a1037ec519eebe216952bfe/src/overdub.jl#L154

source

Reactant.Compiler.codegen_unflatten! Function

codegen_unflatten!

Generate Julia code to wrap the XLA buffers back into the output result datatypes. The name is due to its similarity to the unflatten function in jax.tree_util.register_pytree_node.

source

Reactant.Compiler.codegen_flatten! Function

codegen_flatten!

Generate Julia code to extract the XLA buffers from input arguments. The name is due to its similarity to the flatten function in jax.tree_util.register_pytree_node.

Arguments

• linear_args: A list of arguments to be flattened.

Returns

• flatten_names: A list of Symbols representing the names of the flattened arguments.

• flatten_code: A list of Exprs to extract the XLA buffers from the input arguments.

Note

The linearized arguments do not directly refer to the are the arguments that have been flattened into a single list.

source

Reactant.Compiler.codegen_xla_call Function

codegen_xla_call

Generate Julia code to call the XLA executable.

Arguments

• flatten_names: A list of Symbols representing the names of the flattened linear arguments.

• nresults: The number of results to expect.

• is_pure: Whether the function being compiled is pure (i.e., has no side effects)

source

Other Docstrings

Private

These docstrings are present here to prevent missing docstring warnings. For official Enzyme documentation checkout https://enzymead.github.io/Enzyme.jl/stable/.

EnzymeCore.Forward Constant

const Forward

Default instance of ForwardMode that doesn't return the primal

source

EnzymeCore.ForwardWithPrimal Constant

const ForwardWithPrimal

Default instance of ForwardMode that also returns the primal

source

EnzymeCore.Reverse Constant

const Reverse

Default instance of ReverseMode that doesn't return the primal

source

EnzymeCore.ReverseHolomorphic Constant

const ReverseHolomorphic

Holomorphic instance of ReverseMode that doesn't return the primal

source

EnzymeCore.ReverseHolomorphicWithPrimal Constant

const ReverseHolomorphicWithPrimal

Holomorphic instance of ReverseMode that also returns the primal

source

EnzymeCore.ReverseSplitNoPrimal Constant

const ReverseSplitNoPrimal

Default instance of ReverseModeSplit that doesn't return the primal

source

EnzymeCore.ReverseSplitWithPrimal Constant

const ReverseSplitWithPrimal

Default instance of ReverseModeSplit that also returns the primal

source

EnzymeCore.ReverseWithPrimal Constant

const ReverseWithPrimal

Default instance of ReverseMode that also returns the primal.

source

EnzymeCore.ABI Type

abstract type ABI

Abstract type for what ABI will be used.

Subtypes

• FFIABI (the default)

• InlineABI

• NonGenABI

source

EnzymeCore.Active Type

Active(x)

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

Note

Enzyme gradients with respect to integer values are zero. Active will automatically convert plain integers to floating point values, but cannot do so for integer values in tuples and structs.

source

EnzymeCore.Annotation Type

abstract type Annotation{T}

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

source

EnzymeCore.BatchDuplicated Type

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.

source

EnzymeCore.BatchDuplicatedNoNeed Type

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.

source

EnzymeCore.BatchMixedDuplicated Type

BatchMixedDuplicated(x, ∂f_∂xs)

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

source

EnzymeCore.Const Type

Const(x)

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

source

EnzymeCore.Duplicated Type

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.

source

EnzymeCore.DuplicatedNoNeed Type

DuplicatedNoNeed(x, ∂f_∂x)

Like Duplicated, except also specifies that Enzyme may avoid computing the original result and only compute the derivative values. This creates opportunities for improved performance.

function square_byref(out, v)
    out[] = v * v
    nothing
end

out = Ref(0.0)
dout = Ref(1.0)
Enzyme.autodiff(Reverse, square_byref, DuplicatedNoNeed(out, dout), Active(1.0))
dout[]

# output
0.0

For example, marking the out variable as DuplicatedNoNeed instead of Duplicated allows Enzyme to avoid computing v * v (while still computing its derivative).

This should only be used if x is a write-only variable. Otherwise, if the differentiated function stores values in x and reads them back in subsequent computations, using DuplicatedNoNeed may result in incorrect derivatives. In particular, DuplicatedNoNeed should not be used for preallocated workspace, even if the user might not care about its final value, as marking a variable as NoNeed means that reads from the variable are now undefined.

source

EnzymeCore.FFIABI Type

struct FFIABI <: ABI

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

source

EnzymeCore.ForwardMode Type

struct ForwardMode{
    ReturnPrimal,
    ABI,
    ErrIfFuncWritten,
    RuntimeActivity,
    StrongZero
} <: Mode{ABI,ErrIfFuncWritten,RuntimeActivity,StrongZero}

Subtype of Mode for forward mode differentiation.

Type parameters

• ReturnPrimal: whether to return the primal return value from the augmented-forward.

• other parameters: see Mode

Warning

The type parameters of ForwardMode are not part of the public API and can change without notice. Please use one of the following concrete instantiations instead:

• Forward

• ForwardWithPrimal

You can modify them with the following helper functions:

• WithPrimal / NoPrimal

• set_err_if_func_written / clear_err_if_func_written

• set_runtime_activity / clear_runtime_activity

• set_abi

source

EnzymeCore.InlineABI Type

struct InlineABI <: ABI

Inlining function call ABI.

source

EnzymeCore.MixedDuplicated Type

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.

source

EnzymeCore.Mode Type

abstract type Mode{ABI,ErrIfFuncWritten,RuntimeActivity,StrongZero}

Abstract type for which differentiation mode will be used.

Subtypes

• ForwardMode

• ReverseMode

• ReverseModeSplit

Type parameters

• ABI: what runtime ABI to use

• ErrIfFuncWritten: whether to error when the function differentiated is a closure and written to.

• RuntimeActivity: whether to enable runtime activity (default off). Runtime Activity is required is the differentiability of all mutable variables cannot be determined statically. For a deeper explanation see the FAQ

• StrongZero: whether to enforce that propagating a zero derivative input always ends up in zero derivative outputs. This is required to avoid nan's if one of the arguments may be infinite or nan. For a deeper explanation see the FAQ

Warning

The type parameters of Mode are not part of the public API and can change without notice. You can modify them with the following helper functions:

• WithPrimal / NoPrimal

• set_err_if_func_written / clear_err_if_func_written

• set_runtime_activity / clear_runtime_activity

• set_strong_zero / clear_strong_zero

• set_abi

source

EnzymeCore.NonGenABI Type

struct NonGenABI <: ABI

Non-generated function ABI.

source

EnzymeCore.ReverseMode Type

struct ReverseMode{
    ReturnPrimal,
    RuntimeActivity,
    StrongZero,
    ABI,
    Holomorphic,
    ErrIfFuncWritten
} <: Mode{ABI,ErrIfFuncWritten,RuntimeActivity,StrongZero}

Subtype of Mode for reverse mode differentiation.

Type parameters

• ReturnPrimal: whether to return the primal return value from the augmented-forward pass.

• Holomorphic: Whether the complex result function is holomorphic and we should compute d/dz

• other parameters: see Mode

Warning

The type parameters of ReverseMode are not part of the public API and can change without notice. Please use one of the following concrete instantiations instead:

• Reverse

• ReverseWithPrimal

• ReverseHolomorphic

• ReverseHolomorphicWithPrimal

You can modify them with the following helper functions:

• WithPrimal / NoPrimal

• set_err_if_func_written / clear_err_if_func_written

• set_runtime_activity / clear_runtime_activity

• set_strong_zero / clear_strong_zero

• set_abi

source

EnzymeCore.ReverseModeSplit Type

struct ReverseModeSplit{
    ReturnPrimal,
    ReturnShadow,
    Width,
    RuntimeActivity,
    StrongZero,
    ModifiedBetween,
    ABI,
    ErrFuncIfWritten
} <: Mode{ABI,ErrIfFuncWritten,RuntimeActivity,StrongZero}
    WithPrimal(::Enzyme.Mode)

Subtype of Mode for split reverse mode differentiation, to use in autodiff_thunk and variants.

Type parameters

• ReturnShadow: whether to return the shadow return value from the augmented-forward.

• Width: batch size (pick 0 to derive it automatically)

• ModifiedBetween: Tuple of each argument's "modified between" state (pick true to derive it automatically).

• other parameters: see ReverseMode

Warning

The type parameters of ReverseModeSplit are not part of the public API and can change without notice. Please use one of the following concrete instantiations instead:

• ReverseSplitNoPrimal

• ReverseSplitWithPrimal

You can modify them with the following helper functions:

• WithPrimal / NoPrimal

• set_err_if_func_written / clear_err_if_func_written

• set_runtime_activity / clear_runtime_activity

• set_strong_zero / clear_strong_zero

• set_abi

• ReverseSplitModified, ReverseSplitWidth

source

EnzymeCore.Combined Method

Combined(::ReverseMode)

Turn a ReverseModeSplit object into a ReverseMode object while preserving as many of the settings as possible.

This function acts as the identity on a ReverseMode.

See also Split.

source

EnzymeCore.NoPrimal Method

NoPrimal(::Mode)

Return a new mode which excludes the primal value.

source

EnzymeCore.ReverseSplitModified Method

ReverseSplitModified(::ReverseModeSplit, ::Val{MB})

Return a new instance of ReverseModeSplit mode where ModifiedBetween is set to MB.

source

EnzymeCore.ReverseSplitWidth Method

ReverseSplitWidth(::ReverseModeSplit, ::Val{W})

Return a new instance of ReverseModeSplit mode where Width is set to W.

source

EnzymeCore.Split Method

Split(
    ::ReverseMode, [::Val{ReturnShadow}, ::Val{Width}, ::Val{ModifiedBetween}, ::Val{ShadowInit}]
)

Turn a ReverseMode object into a ReverseModeSplit object while preserving as many of the settings as possible. The rest of the settings can be configured with optional positional arguments of Val type.

This function acts as the identity on a ReverseModeSplit.

See also Combined.

source

EnzymeCore.WithPrimal Method

WithPrimal(::Mode)

Return a new mode which includes the primal value.

source

EnzymeCore.clear_err_if_func_written Function

clear_err_if_func_written(::Mode)

Return a new mode which doesn't throw an error for attempts to write into an unannotated function object.

source

EnzymeCore.clear_runtime_activity Function

clear_runtime_activity(::Mode)

Return a new mode where runtime activity analysis is deactivated. See Enzyme.Mode for more information on runtime activity.

source

EnzymeCore.clear_strong_zero Function

clear_strong_zero(::Mode)

Return a new mode where strong_zero is deactivated. See Enzyme.Mode for more information on strong zero.

source

EnzymeCore.compiler_job_from_backend Function

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

source

EnzymeCore.ignore_derivatives Method

ignore_derivatives(x::T)::T

Behaves like the identity function, but disconnects the "shadow" associated with x. This has the effect of preventing any derivatives from being propagated through x.

Enzyme 0.13.74

Support for ignore_derivatives was added in Enzyme 0.13.74.

source

EnzymeCore.make_zero Function

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.

source

EnzymeCore.make_zero! Function

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

Recursively set a variables differentiable fields to zero.

Warn

Only applicable for mutable types T.

source

EnzymeCore.make_zero Method

make_zero(prev::T)

Helper function to recursively make zero.

source

EnzymeCore.needs_primal Method

needs_primal(::Mode)
needs_primal(::Type{Mode})

Returns true if the mode needs the primal value, otherwise false.

source

EnzymeCore.remake_zero! Function

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

Recursively set a variables differentiable fields to zero.

Warn

This assumes that the input value was previously generated by make_zero. Otherwise, this may not zero the immutable fields of a struct.

source

EnzymeCore.runtime_activity Method

runtime_activity(::Mode)
strong_zero(::Type{<:Mode})

Returns whether the given mode has runtime activity set. For a deeper explanation of what strong zero is see the FAQ

source

EnzymeCore.set_abi Function

set_abi(::Mode, ::Type{ABI})

Return a new mode with its ABI set to the chosen type.

source

EnzymeCore.set_err_if_func_written Function

set_err_if_func_written(::Mode)

Return a new mode which throws an error for any attempt to write into an unannotated function object.

source

EnzymeCore.set_runtime_activity Function

set_runtime_activity(::Mode)
set_runtime_activity(::Mode, activity::Bool)
set_runtime_activity(::Mode, config::Union{FwdConfig,RevConfig})
set_runtime_activity(::Mode, prev::Mode)

Return a new mode where runtime activity analysis is activated / set to the desired value. See the FAQ for more information.

source

EnzymeCore.set_strong_zero Function

set_strong_zero(::Mode)
set_strong_zero(::Mode, activity::Bool)
set_strong_zero(::Mode, config::Union{FwdConfig,RevConfig})
set_strong_zero(::Mode, prev::Mode)

Return a new mode where strong zero is activated / set to the desired value. See the FAQ for more information.

source

EnzymeCore.strong_zero Method

strong_zero(::Mode)
strong_zero(::Type{<:Mode})

Returns whether the given mode has strong zero set. For a deeper explanation of what strong zero is see the FAQ

source

EnzymeCore.within_autodiff Method

within_autodiff()

Returns true if within autodiff, otherwise false.

source

EnzymeCore.EnzymeRules.AugmentedReturn Type

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.

source

EnzymeCore.EnzymeRules.FwdConfig Type

FwdConfig{NeedsPrimal, NeedsShadow, Width, RuntimeActivity, StrongZero}
FwdConfigWidth{Width} = FwdConfig{<:Any, <:Any, Width}

Configuration type to dispatch on in custom forward rules (see forward.

• 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.

• RuntimeActivity: whether runtime activity is enabled. See the FAQ for more information.

• StrongZero: whether strong zero is enabled. See the FAQ for more information.

Getters for the type parameters are provided by needs_primal, needs_shadow, width runtime_activity, and strong_zero.

source

EnzymeCore.EnzymeRules.RevConfig Type

RevConfig{NeedsPrimal, NeedsShadow, Width, Overwritten, RuntimeActivity, StrongZero}
RevConfigWidth{Width} = RevConfig{<: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).

• RuntimeActivity: whether runtime activity is enabled. See the FAQ for more information.

• StrongZero: whether strong zero is enabled. See the FAQ for more information.

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

source

EnzymeCore.EnzymeRules.augmented_primal Function

augmented_primal(::RevConfig, 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.

source

EnzymeCore.EnzymeRules.forward Function

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

Calculate the forward derivative. The first argument is a `FwdConfig](/api/internal#EnzymeCore.EnzymeRules.FwdConfig) object describing parameters of the differentiation. The second argument `func` is the callable for which the rule applies to. Either wrapped in a [`Const`](/api/internal#EnzymeCore.Const)), or a [`Duplicated` if it is a closure. The third argument is the return type annotation, and all other arguments are the annotated function arguments.

source

EnzymeCore.EnzymeRules.inactive Function

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

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

source

EnzymeCore.EnzymeRules.inactive_noinl Function

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.

source

EnzymeCore.EnzymeRules.inactive_type Method

inactive_type(::Type{Ty})

Mark a particular type Ty as always being inactive.

source

EnzymeCore.EnzymeRules.needs_shadow Method

needs_shadow(::FwdConfig)
needs_shadow(::RevConfig)

Whether a custom rule should return the shadow (derivative) of the function result.

source

EnzymeCore.EnzymeRules.noalias Function

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

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

source

EnzymeCore.EnzymeRules.overwritten Method

overwritten(::RevConfig)

A tuple of booleans for each argument (including the function itself), indicating if it is modified between the forward and reverse pass (true if potentially modified between).

source

EnzymeCore.EnzymeRules.primal_type Method

primal_type(::FwdConfig, ::Type{<:Annotation{RT}})
primal_type(::RevConfig, ::Type{<:Annotation{RT}})

Compute the exepcted primal return type given a reverse mode config and return activity

source

EnzymeCore.EnzymeRules.reverse Function

reverse(::RevConfig, func::Annotation{typeof(f)}, dret::Active, tape, args::Annotation...)
reverse(::RevConfig, 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.

source

EnzymeCore.EnzymeRules.shadow_type Method

shadow_type(::FwdConfig, ::Type{<:Annotation{RT}})
shadow_type(::RevConfig, ::Type{<:Annotation{RT}})

Compute the exepcted shadow return type given a reverse mode config and return activity

source

EnzymeCore.needs_primal Method

needs_primal(::FwdConfig)
needs_primal(::RevConfig)

Whether a custom rule should return the original result of the function.

source