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_NAMEThe 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
sourceReactant.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.
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 ofSymbols representing the names of the flattened arguments.flatten_code: A list ofExprs 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.
sourceReactant.Compiler.codegen_xla_call Function
codegen_xla_callGenerate Julia code to call the XLA executable.
Arguments
flatten_names: A list ofSymbols 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)
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 ForwardDefault instance of ForwardMode that doesn't return the primal
EnzymeCore.ForwardWithPrimal Constant
const ForwardWithPrimalDefault instance of ForwardMode that also returns the primal
EnzymeCore.Reverse Constant
const ReverseDefault instance of ReverseMode that doesn't return the primal
EnzymeCore.ReverseHolomorphic Constant
const ReverseHolomorphicHolomorphic instance of ReverseMode that doesn't return the primal
EnzymeCore.ReverseHolomorphicWithPrimal Constant
const ReverseHolomorphicWithPrimalHolomorphic instance of ReverseMode that also returns the primal
EnzymeCore.ReverseSplitNoPrimal Constant
const ReverseSplitNoPrimalDefault instance of ReverseModeSplit that doesn't return the primal
EnzymeCore.ReverseSplitWithPrimal Constant
const ReverseSplitWithPrimalDefault instance of ReverseModeSplit that also returns the primal
EnzymeCore.ReverseWithPrimal Constant
const ReverseWithPrimalDefault instance of ReverseMode that also returns the primal.
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.
EnzymeCore.Annotation Type
abstract type Annotation{T}Abstract type for autodiff function argument wrappers like Const, Active and Duplicated.
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.
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.
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.
EnzymeCore.Const Type
Const(x)Mark a function argument x of autodiff as constant, Enzyme will not auto-differentiate in respect Const arguments.
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 ∂f_∂x.
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.0For 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.
EnzymeCore.FFIABI Type
struct FFIABI <: ABIForeign function call ABI. JIT the differentiated function, then inttoptr call the address.
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:
You can modify them with the following helper functions:
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.
EnzymeCore.Mode Type
abstract type Mode{ABI,ErrIfFuncWritten,RuntimeActivity,StrongZero}Abstract type for which differentiation mode will be used.
Subtypes
Type parameters
ABI: what runtimeABIto useErrIfFuncWritten: 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 FAQStrongZero: 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:
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 computed/dzother 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:
You can modify them with the following helper functions:
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 (pick0to derive it automatically)ModifiedBetween:Tupleof each argument's "modified between" state (picktrueto 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:
You can modify them with the following helper functions:
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.
EnzymeCore.NoPrimal Method
NoPrimal(::Mode)Return a new mode which excludes the primal value.
sourceEnzymeCore.ReverseSplitModified Method
ReverseSplitModified(::ReverseModeSplit, ::Val{MB})Return a new instance of ReverseModeSplit mode where ModifiedBetween is set to MB.
EnzymeCore.ReverseSplitWidth Method
ReverseSplitWidth(::ReverseModeSplit, ::Val{W})Return a new instance of ReverseModeSplit mode where Width is set to W.
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.
EnzymeCore.WithPrimal Method
WithPrimal(::Mode)Return a new mode which includes the primal value.
sourceEnzymeCore.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.
sourceEnzymeCore.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.
sourceEnzymeCore.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.
sourceEnzymeCore.compiler_job_from_backend Function
compiler_job_from_backend(::KernelAbstractions.Backend, F::Type, TT:Type)::GPUCompiler.CompilerJobReturns 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()))
endEnzymeCore.ignore_derivatives Method
ignore_derivatives(x::T)::TBehaves 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.
EnzymeCore.make_zero Function
make_zero(::Type{T}, seen::IdDict, prev::T, ::Val{copy_if_inactive}=Val(false))::TRecursively 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.
EnzymeCore.make_zero! Function
make_zero!(val::T, seen::IdSet{Any}=IdSet())::NothingRecursively set a variables differentiable fields to zero.
Warn
Only applicable for mutable types T.
EnzymeCore.needs_primal Method
needs_primal(::Mode)
needs_primal(::Type{Mode})Returns true if the mode needs the primal value, otherwise false.
EnzymeCore.remake_zero! Function
remake_zero!(val::T, seen::IdSet{Any}=IdSet())::NothingRecursively 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.
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
sourceEnzymeCore.set_abi Function
set_abi(::Mode, ::Type{ABI})Return a new mode with its ABI set to the chosen type.
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.
sourceEnzymeCore.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.
sourceEnzymeCore.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.
sourceEnzymeCore.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
sourceEnzymeCore.within_autodiff Method
within_autodiff()Returns true if within autodiff, otherwise false.
sourceEnzymeCore.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.
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.
NeedsPrimalandNeedsShadow: 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.
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).
NeedsPrimalandNeedsShadow: 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.
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.
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.
sourceEnzymeCore.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.
sourceEnzymeCore.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.
sourceEnzymeCore.EnzymeRules.inactive_type Method
inactive_type(::Type{Ty})Mark a particular type Ty as always being inactive.
EnzymeCore.EnzymeRules.needs_shadow Method
needs_shadow(::FwdConfig)
needs_shadow(::RevConfig)Whether a custom rule should return the shadow (derivative) of the function result.
sourceEnzymeCore.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.
sourceEnzymeCore.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).
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
sourceEnzymeCore.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.
sourceEnzymeCore.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
sourceEnzymeCore.needs_primal Method
needs_primal(::FwdConfig)
needs_primal(::RevConfig)Whether a custom rule should return the original result of the function.
source