Enzyme.gradient!Method
gradient!(::ReverseMode, dx, f, x)

Compute the gradient of an array-input function f using reverse mode, storing the derivative result in an existing array dx. Both x and dx must be Arrays of the same type.

Example:

f(x) = x[1]*x[2]

dx = [0.0, 0.0]
gradient!(Reverse, dx, f, [2.0, 3.0])

# output

2-element Vector{Float64}:
 3.0
 2.0
Enzyme.gradientMethod
gradient(::ForwardMode, f, x; shadow=onehot(x))

Compute the gradient of an array-input function f using forward mode. The optional keyword argument shadow is a vector of one-hot vectors of type x which are used to forward-propagate into the return. For performance reasons, this should be computed once, outside the call to gradient, rather than within this call.

Example:

f(x) = x[1]*x[2]

grad = gradient(Forward, f, [2.0, 3.0])

# output

(3.0, 2.0)
Enzyme.gradientMethod
gradient(::ReverseMode, f, x)

Compute the gradient of a real-valued function f using reverse mode. This will allocate and return new array make_zero(x) with the gradient result.

Besides arrays, for struct x it returns another instance of the same type, whose fields contain the components of the gradient. In the result, grad.a contains ∂f/∂x.a for any differential x.a, while grad.c == x.c for other types.

Examples:

f(x) = x[1]*x[2]

grad = gradient(Reverse, f, [2.0, 3.0])

# output

2-element Vector{Float64}:
 3.0
 2.0
grad = gradient(Reverse, only ∘ f, (a = 2.0, b = [3.0], c = "str"))

# output

(a = 3.0, b = [2.0], c = "str")
Enzyme.gradientMethod
gradient(::ForwardMode, f, x::Union{Array,NTuple}, ::Val{chunk}; shadow=onehot(x))

Compute the gradient of an array-input function f using vector forward mode. Like gradient, except it uses a chunk size of chunk to compute chunk derivatives in a single call.

Example:

f(x) = x[1]*x[2]

grad = gradient(Forward, f, [2.0, 3.0], Val(2))

# output

(3.0, 2.0)
Enzyme.jacobianMethod
jacobian(::ForwardMode, f, x; shadow=onehot(x))
jacobian(::ForwardMode, f, x, ::Val{chunk}; shadow=onehot(x))

Compute the jacobian of an array-input function f using (potentially vector) forward mode. This is a simple rename of the gradient function, and all relevant arguments apply here.

Example:

f(x) = [x[1]*x[2], x[2]]

grad = jacobian(Forward, f, [2.0, 3.0])

# output

2×2 Matrix{Float64}:
 3.0  2.0
 0.0  1.0
Enzyme.jacobianMethod
jacobian(::ReverseMode, f, x, ::Val{num_outs}, ::Val{chunk})

Compute the jacobian of an array-input function f using (potentially vector) reverse mode. The chunk argument denotes the chunk size to use and num_outs denotes the number of outputs f will return in an array.

Example:

f(x) = [x[1]*x[2], x[2]]

grad = jacobian(Reverse, f, [2.0, 3.0], Val(2))

# output

2×2 Matrix{Float64}:
 3.0  2.0
 0.0  1.0
Enzyme.typetreeFunction
function typetree(T, ctx, dl, seen=TypeTreeTable())

Construct a Enzyme typetree from a Julia type.

Warning

When using a memoized lookup by providing seen across multiple calls to typtree the user must call copy on the returned value before mutating it.

Enzyme.unsafe_to_pointerMethod
unsafe_to_pointer
Warning

Assumes that val is globally rooted and pointer to it can be leaked. Prefer pointer_from_objref. Only use inside Enzyme.jl should be for Types.

EnzymeCore.autodiffMethod
autodiff(::ForwardMode, f, Activity, args::Vararg{<:Annotation, Nargs})

Auto-differentiate function f at arguments args using forward mode.

args may be numbers, arrays, structs of numbers, structs of arrays and so on. Enzyme will only differentiate in respect to arguments that are wrapped in a Duplicated or similar argument. Unlike reverse mode in autodiff, Active arguments are not allowed here, since all derivative results of immutable objects will be returned and should instead use Duplicated or variants like DuplicatedNoNeed.

Activity is the Activity of the return value, it may be:

  • Const if the return is not to be differentiated with respect to
  • Duplicated, if the return is being differentiated with respect to and both the original value and the derivative return are desired
  • DuplicatedNoNeed, if the return is being differentiated with respect to and only the derivative return is desired.
  • BatchDuplicated, like Duplicated, but computing multiple derivatives at once. All batch sizes must be the same for all arguments.
  • BatchDuplicatedNoNeed, like DuplicatedNoNeed, but computing multiple derivatives at one. All batch sizes must be the same for all arguments.

Example returning both original return and derivative:

f(x) = x*x
res, ∂f_∂x = autodiff(Forward, f, Duplicated, Duplicated(3.14, 1.0))

# output

(9.8596, 6.28)

Example returning just the derivative:

f(x) = x*x
∂f_∂x = autodiff(Forward, f, DuplicatedNoNeed, Duplicated(3.14, 1.0))

# output

(6.28,)
EnzymeCore.autodiffMethod
autodiff(mode::Mode, f, ::Type{A}, args::Vararg{Annotation, Nargs})

Like autodiff but will try to extend f to an annotation, if needed.

EnzymeCore.autodiffMethod
autodiff(::ReverseMode, f, Activity, args::Vararg{<:Annotation, Nargs})

Auto-differentiate function f at arguments args using reverse mode.

Limitations:

  • f may only return a Real (of a built-in/primitive type) or nothing, not an array, struct, BigFloat, etc. To handle vector-valued return types, use a mutating f! that returns nothing and stores it's return value in one of the arguments, which must be wrapped in a Duplicated.

args may be numbers, arrays, structs of numbers, structs of arrays and so on. Enzyme will only differentiate in respect to arguments that are wrapped in an Active (for arguments whose derivative result must be returned rather than mutated in place, such as primitive types and structs thereof) or Duplicated (for mutable arguments like arrays, Refs and structs thereof).

Activity is the Activity of the return value, it may be Const or Active.

Example:

a = 4.2
b = [2.2, 3.3]; ∂f_∂b = zero(b)
c = 55; d = 9

f(a, b, c, d) = a * √(b[1]^2 + b[2]^2) + c^2 * d^2
∂f_∂a, _, _, ∂f_∂d = autodiff(Reverse, f, Active, Active(a), Duplicated(b, ∂f_∂b), Const(c), Active(d))[1]

# output

(3.966106403010388, nothing, nothing, 54450.0)

here, autodiff returns a tuple $(\partial f/\partial a, \partial f/\partial d)$, while $\partial f/\partial b$ will be added to ∂f_∂b (but not returned). c will be treated as Const(c).

One can also request the original returned value of the computation.

Example:

Enzyme.autodiff(ReverseWithPrimal, x->x*x, Active(3.0))

# output

((6.0,), 9.0)
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.autodiff_deferredMethod
autodiff_deferred(::ForwardMode, f, Activity, args::Vararg{<:Annotation, Nargs})

Same as autodiff(::ForwardMode, f, Activity, args) but uses deferred compilation to support usage in GPU code, as well as high-order differentiation.

EnzymeCore.autodiff_deferredMethod
autodiff_deferred(::ReverseMode, f, Activity, args::Vararg{<:Annotation, Nargs})

Same as autodiff but uses deferred compilation to support usage in GPU code, as well as high-order differentiation.

EnzymeCore.autodiff_deferred_thunkMethod
autodiff_deferred_thunk(::ReverseModeSplit, ftype, Activity, argtypes::Vararg{Type{<:Annotation}, Nargs})

Provide the split forward and reverse pass functions for annotated function type ftype when called with args of type argtypes when using reverse mode.

Activity is the Activity of the return value, it may be Const, Active, or Duplicated (or its variants DuplicatedNoNeed, BatchDuplicated, and BatchDuplicatedNoNeed).

The forward function will return a tape, the primal (or nothing if not requested), and the shadow (or nothing if not a Duplicated variant), and tapes the corresponding type arguements provided.

The reverse function will return the derivative of Active arguments, updating the Duplicated arguments in place. The same arguments to the forward pass should be provided, followed by the adjoint of the return (if the return is active), and finally the tape from the forward pass.

Example:


A = [2.2]; ∂A = zero(A)
v = 3.3

function f(A, v)
    res = A[1] * v
    A[1] = 0
    res
end

TapeType = tape_type(ReverseSplitWithPrimal, Const{typeof(f)}, Active, Duplicated{typeof(A)}, Active{typeof(v)})
forward, reverse = autodiff_deferred_thunk(ReverseSplitWithPrimal, TapeType, Const{typeof(f)}, Active{Float64}, Duplicated{typeof(A)}, Active{typeof(v)})

tape, result, shadow_result  = forward(Const(f), Duplicated(A, ∂A), Active(v))
_, ∂v = reverse(Const(f), Duplicated(A, ∂A), Active(v), 1.0, tape)[1]

result, ∂v, ∂A 

# output

(7.26, 2.2, [3.3])
EnzymeCore.autodiff_thunkMethod
autodiff_thunk(::ForwardMode, ftype, Activity, argtypes::Vararg{Type{<:Annotation}, Nargs})

Provide the thunk forward mode function for annotated function type ftype when called with args of type argtypes.

Activity is the Activity of the return value, it may be Const or Duplicated (or its variants DuplicatedNoNeed, BatchDuplicated, andBatchDuplicatedNoNeed).

The forward function will return the primal (if requested) and the shadow (or nothing if not a Duplicated variant).

Example returning both original return and derivative:

a = 4.2
b = [2.2, 3.3]; ∂f_∂b = zero(b)
c = 55; d = 9

f(x) = x*x
forward = autodiff_thunk(Forward, Const{typeof(f)}, Duplicated, Duplicated{Float64})
res, ∂f_∂x = forward(Const(f), Duplicated(3.14, 1.0))

# output

(9.8596, 6.28)

Example returning just the derivative:

a = 4.2
b = [2.2, 3.3]; ∂f_∂b = zero(b)
c = 55; d = 9

f(x) = x*x
forward = autodiff_thunk(Forward, Const{typeof(f)}, DuplicatedNoNeed, Duplicated{Float64})
∂f_∂x = forward(Const(f), Duplicated(3.14, 1.0))

# output

(6.28,)
EnzymeCore.autodiff_thunkMethod
autodiff_thunk(::ReverseModeSplit, ftype, Activity, argtypes::Vararg{Type{<:Annotation, Nargs})

Provide the split forward and reverse pass functions for annotated function type ftype when called with args of type argtypes when using reverse mode.

Activity is the Activity of the return value, it may be Const, Active, or Duplicated (or its variants DuplicatedNoNeed, BatchDuplicated, and BatchDuplicatedNoNeed).

The forward function will return a tape, the primal (or nothing if not requested), and the shadow (or nothing if not a Duplicated variant), and tapes the corresponding type arguements provided.

The reverse function will return the derivative of Active arguments, updating the Duplicated arguments in place. The same arguments to the forward pass should be provided, followed by the adjoint of the return (if the return is active), and finally the tape from the forward pass.

Example:


A = [2.2]; ∂A = zero(A)
v = 3.3

function f(A, v)
    res = A[1] * v
    A[1] = 0
    res
end

forward, reverse = autodiff_thunk(ReverseSplitWithPrimal, Const{typeof(f)}, Active, Duplicated{typeof(A)}, Active{typeof(v)})

tape, result, shadow_result  = forward(Const(f), Duplicated(A, ∂A), Active(v))
_, ∂v = reverse(Const(f), Duplicated(A, ∂A), Active(v), 1.0, tape)[1]

result, ∂v, ∂A 

# output

(7.26, 2.2, [3.3])
Enzyme.@import_fruleMacro
import_frule(::fn, tys...)

Automatically import a ChainRulesCore.fruleas a custom forward modeEnzymeRule. When called in batch mode, this will end up calling the primal multiple times, which may result in incorrect behavior if the function mutates, and slow code, always. Importing the rule fromChainRules` is also likely to be slower than writing your own rule, and may also be slower than not having a rule at all.

Use with caution.

Enzyme.@import_frule(typeof(Base.sort), Any);

x=[1.0, 2.0, 0.0]; dx=[0.1, 0.2, 0.3]; ddx = [0.01, 0.02, 0.03];

Enzyme.autodiff(Forward, sort, Duplicated, BatchDuplicated(x, (dx,ddx)))
Enzyme.autodiff(Forward, sort, DuplicatedNoNeed, BatchDuplicated(x, (dx,ddx)))
Enzyme.autodiff(Forward, sort, DuplicatedNoNeed, BatchDuplicated(x, (dx,)))
Enzyme.autodiff(Forward, sort, Duplicated, BatchDuplicated(x, (dx,)))

# output

(var"1" = [0.0, 1.0, 2.0], var"2" = (var"1" = [0.3, 0.1, 0.2], var"2" = [0.03, 0.01, 0.02]))
(var"1" = (var"1" = [0.3, 0.1, 0.2], var"2" = [0.03, 0.01, 0.02]),)
(var"1" = [0.3, 0.1, 0.2],)
(var"1" = [0.0, 1.0, 2.0], var"2" = [0.3, 0.1, 0.2])
Enzyme.@import_rruleMacro
import_rrule(::fn, tys...)

Automatically import a ChainRules.rrule as a custom reverse mode EnzymeRule. When called in batch mode, this will end up calling the primal multiple times which results in slower code. This macro assumes that the underlying function to be imported is read-only, and returns a Duplicated or Const object. This macro also assumes that the inputs permit a .+= operation and that the output has a valid Enzyme.make_zero function defined. It also assumes that overwritten(x) accurately describes if there is any non-preserved data from forward to reverse, not just the outermost data structure being overwritten as provided by the specification.

Finally, this macro falls back to almost always caching all of the inputs, even if it may not be needed for the derivative computation.

As a result, this auto importer is also likely to be slower than writing your own rule, and may also be slower than not having a rule at all.

Use with caution.

Enzyme.@import_rrule(typeof(Base.sort), Any);
Enzyme.API.fast_math!Method
fast_math!(val::Bool)

Whether generated derivatives have fast math on or off, default on.

Enzyme.API.inlineall!Method
inlineall!(val::Bool)

Whether to inline all (non-recursive) functions generated by Julia within a single compilation unit. This may improve Enzyme's ability to successfully differentiate code and improve performance of the original and generated derivative program. It often, however, comes with an increase in compile time. This is off by default.

Enzyme.API.instname!Method
instname!(val::Bool)

Whether to add a name to all LLVM values. This may be helpful for debugging generated programs, both primal and derivative. Off by default.

Enzyme.API.looseTypeAnalysis!Method
looseTypeAnalysis!(val::Bool)

Enzyme runs a type analysis to deduce the corresponding types of all values being differentiated. This is necessary to compute correct derivatives of various values. For example, a copy of Float32's requires a different derivative than a memcpy of Float64's, Ptr's, etc. In some cases Enzyme may not be able to deduce all the types necessary and throw an unknown type error. If this is the case, open an issue. One can silence these issues by setting looseTypeAnalysis!(true) which tells Enzyme to make its best guess. This will remove the error and allow differentiation to continue, however, it may produce incorrect results. Alternatively one can consider increasing the space of the evaluated type lattice which gives Enzyme more time to run a more thorough analysis through the use of maxtypeoffset!

Enzyme.API.maxtypedepth!Method
maxtypedepth!(val::Bool)

Enzyme runs a type analysis to deduce the corresponding types of all values being differentiated. This is necessary to compute correct derivatives of various values. To ensure this analysis temrinates, it operates on a finite lattice of possible states. This function sets the maximum depth into a type that Enzyme will consider. A smaller value will cause type analysis to run faster, but may result in some necessary types not being found and result in unknown type errors. A larger value may result in unknown type errors being resolved by searching a larger space, but may run longer. The default setting is 6.

Enzyme.API.maxtypeoffset!Method
maxtypeoffset!(val::Bool)

Enzyme runs a type analysis to deduce the corresponding types of all values being differentiated. This is necessary to compute correct derivatives of various values. To ensure this analysis temrinates, it operates on a finite lattice of possible states. This function sets the maximum offset into a type that Enzyme will consider. A smaller value will cause type analysis to run faster, but may result in some necessary types not being found and result in unknown type errors. A larger value may result in unknown type errors being resolved by searching a larger space, but may run longer. The default setting is 512.

Enzyme.API.memmove_warning!Method
memmove_warning!(val::Bool)

Whether to issue a warning when differentiating memmove. Off by default.

Enzyme.API.printactivity!Method
printactivity!(val::Bool)

An debugging option for developers of Enzyme. If one sets this flag prior to the first differentiation of a function, Enzyme will print (to stderr) a log of all decisions made during Activity Analysis (the analysis which determines what values/instructions are differentiated). This may be useful for debugging MixedActivity errors, correctness, and performance errors. Off by default

Enzyme.API.printall!Method
printall!(val::Bool)

An debugging option for developers of Enzyme. If one sets this flag prior to the first differentiation of a function, Enzyme will print (to stderr) the LLVM function being differentiated, as well as all generated derivatives immediately after running Enzyme (but prior to any other optimizations). Off by default

Enzyme.API.printdiffuse!Method
printdiffuse!(val::Bool)

An debugging option for developers of Enzyme. If one sets this flag prior to the first differentiation of a function, Enzyme will print (to stderr) information about each LLVM value – specifically whether it and its shadow is required for computing the derivative. In contrast to printunnecessary!, this flag prints debug log for the analysis which determines for each value and shadow value, whether it can find a user which would require it to be kept around (rather than being deleted). This is prior to any cache optimizations and a debug log of Differential Use Analysis. This may be helpful for debugging caching, phi node deletion, performance, and other errors. Off by default

Enzyme.API.printperf!Method
printperf!(val::Bool)

An debugging option for developers of Enzyme. If one sets this flag prior to the first differentiation of a function, Enzyme will print (to stderr) performance information about generated derivative programs. It will provide debug information that warns why particular values are cached for the reverse pass, and thus require additional computation/storage. This is particularly helpful for debugging derivatives which OOM or otherwise run slow. ff by default

Enzyme.API.printtype!Method
printtype!(val::Bool)

An debugging option for developers of Enzyme. If one sets this flag prior to the first differentiation of a function, Enzyme will print (to stderr) a log of all decisions made during Type Analysis (the analysis which Enzyme determines the type of all values in the program). This may be useful for debugging correctness errors, illegal type analysis errors, insufficient type information errors, correctness, and performance errors. Off by default

Enzyme.API.printunnecessary!Method
printunnecessary!(val::Bool)

An debugging option for developers of Enzyme. If one sets this flag prior to the first differentiation of a function, Enzyme will print (to stderr) information about each LLVM value – specifically whether it and its shadow is required for computing the derivative. In contrast to printdiffuse!, this flag prints the final results after running cache optimizations such as minCut (see Recompute vs Cache Heuristics from this paper and slides 31-33 from this presentation) for a description of the caching algorithm. This may be helpful for debugging caching, phi node deletion, performance, and other errors. Off by default

Enzyme.API.runtimeActivity!Method
runtimeActivity!(val::Bool)

Enzyme runs an activity analysis which deduces which values, instructions, etc are necessary to be differentiated and therefore involved in the differentiation procedure. This runs at compile time. However, there may be implementation flaws in this analysis that means that Enzyme cannot deduce that an inactive (const) value is actually const. Alternatively, there may be some data which is conditionally active, depending on which runtime branch is taken. In these cases Enzyme conservatively presumes the value is active.

However, in certain cases, an insufficiently aggressive activity analysis may result in derivative errors – for example by mistakenly using the primal (const) argument and mistaking it for the duplicated shadow. As a result this may result in incorrect results, or accidental updates to the primal.

This flag enables runntime activity which tells all load/stores to check at runtime whether the value they are updating is indeed active (in addition to the compile-time activity analysis). This will remedy these such errors, but at a performance penalty of performing such checks.

It is on the Enzyme roadmap to add a PotentiallyDuplicated style activity, in addition to the current Const and Duplicated styles that will disable the need for this, which does not require the check when a value is guaranteed active, but still supports runtime-based activity information.

This function takes an argument to set the runtime activity value, true means it is on, and false means off. By default it is off.

Enzyme.API.strictAliasing!Method
strictAliasing!(val::Bool)

Whether Enzyme's type analysis will assume strict aliasing semantics. When strict aliasing semantics are on (the default), Enzyme can propagate type information up through conditional branches. This may lead to illegal type errors when analyzing code with unions. Disabling strict aliasing will enable these union types to be correctly analyzed. However, it may lead to some errors that sufficient type information cannot be deduced. One can turn these insufficient type information errors into to warnings by calling looseTypeAnalysis!(true) which tells Enzyme to use its best guess in such scenarios.

Enzyme.API.strong_zero!Method
strong_zero!(val::Bool)

Whether to enforce multiplication by zero as enforcing a zero result even if multiplying against a NaN or infinity. Necessary for some programs in which a value has a zero derivative since it is unused, even if it has an otherwise infinite or nan derivative.

Enzyme.API.typeWarning!Method
typeWarning!(val::Bool)

Whether to print a warning when Type Analysis learns informatoin about a value's type which cannot be represented in the current size of the lattice. See maxtypeoffset! for more information. Off by default.