Enzyme

Documentation for Enzyme.jl, the Julia bindings for Enzyme.

Enzyme performs automatic differentiation (AD) of statically analyzable LLVM. It is highly-efficient and its ability perform AD on optimized code allows Enzyme to meet or exceed the performance of state-of-the-art AD tools.

Enzyme.jl can be installed in the usual way Julia packages are installed:

] add Enzyme

The Enzyme binary dependencies will be installed automatically via Julia's binary actifact system.

The Enzyme.jl API revolves around the function autodiff, see it's documentation for details and a usage example. Also see Implementing pullbacks on how to use Enzyme.jl to implement back-propagation for functions with non-scalar results.

Caveats / Known-issues

Activity of temporary storage

If you have pass any temporary storage which may be involved in an active computation to a function you want to differentiate, you must also pass in a duplicated temporary storage for use in computing the derivatives.

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

# Incorrect [ returns (0.0,) ]
Enzyme.autodiff(f, Active(1.2), Const(Vector{Float64}(undef, 1)), Const(5))

# Correct [ returns (10.367999999999999,) == 1.2^4 * 5 ]
Enzyme.autodiff(f, Active(1.2), Duplicated(Vector{Float64}(undef, 1), Vector{Float64}(undef, 1)), Const(5))

CUDA.jl support

CUDA.jl is only support on Julia v1.7.0 and onwards. On 1.6 attempting to differentiate CUDA kernel functions, will not use device overloads correctly and thus return fundamentally wrong results.

Sparse Arrays

At the momment there is limited support for sparse linear algebra operations. Sparse arrays may be used, but care must be taken because backing arrays drop zeros in Julia (unless told not to).

using SparseArrays

a=sparse([2.0])
f(a)=sum(a)

# Incorrect: SparseMatrixCSC drops explicit zeros
# returns 1-element SparseVector{Float64, Int64} with 0 stored entries
da=sparse([0.0])

# Correct: Prevent SparseMatrixCSC from dropping zeros
# returns 1-element SparseVector{Float64, Int64} with 1 stored entry:
#  [1]  =  1.0
da=sparsevec([1],[0.0])

Enzyme.autodiff(Reverse,f,Active,Duplicated(a,da))
@show da