Enzyme.jl

When Enzyme performs automatic differentiation, it often needs to allocate 'shadow objects' (copies used to track derivatives). If the primal object has a Julia finalizer attached (used for resource management like manual memory allocation), Enzyme must handle it to prevent resource leaks.

Enzyme's approach is to attach the finalizer to the shadow object but define it to be 'inactive'—meaning the finalizer contains no instructions relevant to the AD process itself, but it still executes to ensure that resources associated with the shadow object are correctly released.

mutable struct Obj
    x::Float64
    function Obj(x)
        o = new(x)
        finalizer(o) do o
            # resource management logic
        end
        return o
    end
end

function f(x)
    o = Obj(x)
    return o.x
end

# Enzyme will manage the shadow object for 'o' and its finalizer
Enzyme.autodiff(Forward, f, Duplicated(1.0, 1.0))

For high-performance pointer-based code, Enzyme's memory semantics allow you to differentiate functions that only touch specific offsets of a pointer.

Safe Approach: Pass a shadow pointer (dptr) that is congruent with the primal pointer (ptr) in terms of length and layout.

Optimized Approach: If a function only accesses a specific offset (e.g., unsafe_load(ptr, 47)), you can pass a smaller shadow allocation by offsetting the pointer. This ensures the derivative code's access to the shadow is in-bounds while minimizing allocation.

Warning: This requires a deep understanding of how the generated derivative code accesses memory to avoid out-of-bounds access.

# Example: Optimizing a shadow for a single pointer access at offset 47
ptr = Base.reinterpret(Ptr{Float64}, Libc.malloc(100*sizeof(Float64)))
# ... setup ptr ...

# Allocate only enough for one Float64
dptr = Base.reinterpret(Ptr{Float64}, Libc.calloc(sizeof(Float64), 1))

# Offset dptr so that unsafe_load(dptr, 47) hits the 0th byte of the allocation
# We subtract 46 * sizeof(Float64) because Julia is 1-indexed
autodiff(Reverse, f, Duplicated(ptr, dptr - 46 * sizeof(Float64)))

Enzyme does not have special-cased logic for sparse arrays; instead, it uses a universal memory layout principle: the shadow (derivative) memory layout is identical to the primal layout.

For a sparse array represented by a backing array of values and an index array, Enzyme's shadow will consist of a second backing array of the same size (to store derivatives) and the same index array.

Key implications:

1. Composition: This representation allows Enzyme rules to compose correctly via the chain rule (e.g., in f(A(x)) where A(x) is sparse).
2. Memory Efficiency: The memory requirements for the derivative are the same as the primal, avoiding the 'memory blow-up' seen in other AD tools that convert sparse structures to dense ones.
3. Top-level behavior: At the top level, Enzyme only updates the derivatives of elements actually read/accessed by the function. This can lead to a semantic mismatch where the sparse array printer shows zeros for elements that Enzyme has technically 'ignored' because they didn't contribute to the computation.

Enzyme does not assume a specific convention for complex differentiation (which can be ambiguous) and instead requires users to return real numbers or specify the desired convention.

To differentiate a function $f(z)$ that returns a complex number, you can treat it as a function of two real variables $f(x, y) = u(x, y) + i v(x, y)$ and differentiate the real and imaginary parts separately using Reverse or Forward modes.

Reverse Mode

In Reverse mode, you can compute the gradient of the real part and the imaginary part separately by wrapping the function to return real(f(z)) and imag(f(z)) respectively.

Forward Mode

In Forward mode, you provide a differential input (a 'shadow') to compute the derivative along a specific direction. For example, seeding with 1.0 + 0.0im computes the derivative along the real axis, while 0.0 + 1.0im computes it along the imaginary axis.

f(z) = z * z
z = 3.1 + 2.7im

# Reverse mode for real and imaginary parts
grad_u = Enzyme.autodiff(Reverse, z->real(f(z)), Active, Active(z))[1][1]
grad_v = Enzyme.autodiff(Reverse, z->imag(f(z)), Active, Active(z))[1][1]

# Forward mode for directional derivatives
d_dx = Enzyme.autodiff(Forward, f, Duplicated(z, 1.0+0.0im))[1]
d_dy = Enzyme.autodiff(Forward, f, Duplicated(z, 0.0+1.0im))[1]

Enzyme tracks differentiable dataflow through base types, primarily floating-point numbers. It can differentiate through any complex data structure (structs, arrays, linked lists) as long as they contain these base types.

Supported Base Types

• Float32, Float64, Float16, BFloat16, etc.

Unsupported Types

Enzyme does not currently support differentiating data contained in:

• Int (Integers)
• String
• Val (Type constants)

If you attempt to differentiate through these, you may encounter errors such as Return type ... not marked Const or errors regarding ghost or constant type.

# Example: Differentiating a simple float
f(x) = x * x
Enzyme.autodiff(Forward, f, Duplicated(3.0, 1.0))

# Example: Differentiating a struct containing floats
struct Pair
    lhs::Float64
    rhs::Float64
end
f_pair(x) = x.lhs * x.rhs
Enzyme.autodiff(Forward, f_pair, Duplicated(Pair(3.0, 2.0), Pair(1.0, 0.0)))

# Example: Differentiating a complex structure (Linked List)
struct LList
    prev::Union{Nothing, LList}
    value::Float64
end
# ... (list construction and sum logic) ...
Enzyme.autodiff(Forward, list_sum, Duplicated(list, dlist))

In combined reverse mode, Enzyme.autodiff only handles functions with scalar outputs. To implement pullbacks for functions that return arrays (or modify arrays), use a mutating function that returns nothing and stores the result in one of the arguments. These arguments must be wrapped in Duplicated.

Key behaviors:

• The mutating function is generally more efficient than allocating a new output array.
• The result of backpropagation is added to the provided shadow arguments (e.g., ∂z_∂A); they act as accumulators for gradient information.

To ensure the shadow array has the correct data layout and matches the primal, use Enzyme.make_zero(x) to initialize it.

using Enzyme, Random

# 1. Define a mutating function that returns nothing
function mymul!(R, A, B)
    @inbounds for j in axes(B, 2), i in axes(A, 1)
        @inbounds for k in axes(A,2)
            R[i,j] += A[i,k] * B[k,j]
        end
    end
    return nothing
end

# 2. Setup data and shadows
A = rand(5, 3)
B = rand(3, 7)
R = zeros(size(A,1), size(B,2))
∂z_∂R = rand(size(R)...)

# Use make_zero to ensure correct layout/type
∂z_∂A = Enzyme.make_zero(A)
∂z_∂B = Enzyme.make_zero(B)

# 3. Call autodiff with Duplicated wrappers
Enzyme.autodiff(Reverse, mymul!, Const, Duplicated(R, ∂z_∂R), Duplicated(A, ∂z_∂A), Duplicated(B, ∂z_∂B))

# Note: ∂z_∂A and ∂z_∂B now contain the accumulated gradients

To avoid multiple forward passes when computing derivatives for both the real and imaginary parts of a complex function, you can use Enzyme's batching capabilities.

Batched Reverse Mode

Use autodiff_thunk with ReverseSplitWidth to compute multiple reverse passes (e.g., one for the real part and one for the imaginary part) in a single operation. This requires using ReverseSplitNoPrimal or similar split modes.

Batched Forward Mode

Use BatchDuplicated to provide multiple differential inputs (shadows) simultaneously in a single autodiff call.

# Batched Forward Mode
Enzyme.autodiff(Forward, f, BatchDuplicated(z, (1.0+0.0im, 1.0+0.0im)))[1]

# Batched Reverse Mode using thunks
fwd, rev = Enzyme.autodiff_thunk(ReverseSplitWidth(ReverseSplitNoPrimal, Val(2)), Const{typeof(f)}, Active, Active{ComplexF64})
rev(Const(f), Active(z), (1.0 + 0.0im, 0.0 + 1.0im), fwd(Const(f), Active(z))[1])[1][1]

Enzyme supports differentiating code using sparse arrays, but you must ensure the shadow (derivative) memory layout matches the primal. A common pitfall in Julia is that standard sparse array constructors (like sparse()) automatically drop explicit zeros. If the shadow array drops these zeros, Enzyme cannot correctly store or update the derivatives for those positions.

To avoid this, use sparsevec to prevent the dropping of zeros, or use the Enzyme.make_zero(x) helper function, which automatically generates a correctly structured shadow data structure for any input x.

using SparseArrays

# Incorrect: sparse([0.0]) drops the zero, leading to an empty shadow
a = sparse([2.0])
da1 = sparse([0.0]) 
Enzyme.autodiff(Reverse, sum, Active, Duplicated(a, da1))

# Correct: use sparsevec to preserve the zero entry
da2 = sparsevec([1], [0.0])
Enzyme.autodiff(Reverse, sum, Active, Duplicated(a, da2))

# Recommended: use Enzyme.make_zero to automate shadow creation
# Enzyme.gradient(Reverse, sum, a) internally calls make_zero(a)

An EnzymeRuntimeActivityError occurs when a function's return value's differentiability depends on runtime data (e.g., an if statement returning either a Duplicated variable or a Const variable). This is known as 'activity instability'.

While you should ideally make your code 'activity-stable' (ensuring variables are always either differentiable or always constant), you can use Runtime Activity mode to differentiate such code.

Warning: When using Runtime Activity, if the computed derivative of a function is mutable, you must check if the primal and shadow represent the same pointer. If dout === out (for pointer-like types), the true derivative is actually zero.

To enable this, use Enzyme.set_runtime_activity(ForwardWithPrimal) (or the appropriate mode for your differentiation type) as the first argument to autodiff.

# Enable runtime activity to handle conditional differentiability
dout, out = Enzyme.autodiff(Enzyme.set_runtime_activity(ForwardWithPrimal), g, Const(condition), Duplicated(x, dx), Const(y))

# CRITICAL: Check for derivative aliasing
if dout === out
    # The actual derivative is zero
end

If a function being differentiated uses a temporary buffer or storage (e.g., an array used for intermediate calculations), you cannot mark that storage as Const.

Marking a buffer as Const tells Enzyme that all values loaded from or stored into that buffer are non-differentiable. If the function modifies the buffer during computation, the derivative calculation will be incorrect. Instead, you must pass the temporary storage as Duplicated(buffer, shadow_buffer) so Enzyme can track the derivatives of the values processed through that buffer.

function f(x, tmp, k, n)
    tmp[1] = 1.0
    for i in 1:n
        tmp[k] *= x
    end
    tmp[1]
end

# Incorrect: tmp is marked Const, so Enzyme ignores its internal modifications
Enzyme.autodiff(Reverse, f, Active(1.2), Const(Vector{Float64}(undef, 1)), Const(1), Const(5))

# Correct: tmp is marked Duplicated so its contents are differentiable
Enzyme.autodiff(Reverse, f, Active(1.2), Duplicated(Vector{Float64}(undef, 1), zeros(1)), Const(1), Const(5))