Enzyme support

Project.toml

@@ -1,31 +1,15 @@
 name = "AbstractDifferentiation"
 uuid = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"
 authors = ["Mohamed Tarek <[email protected]> and contributors"]
-version = "0.5.2"
+version = "0.6.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 
-[weakdeps]
-ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
-FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
-ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
-ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
-Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
-
-[extensions]
-AbstractDifferentiationChainRulesCoreExt = "ChainRulesCore"
-AbstractDifferentiationFiniteDifferencesExt = "FiniteDifferences"
-AbstractDifferentiationForwardDiffExt = ["DiffResults", "ForwardDiff"]
-AbstractDifferentiationReverseDiffExt = ["DiffResults", "ReverseDiff"]
-AbstractDifferentiationTrackerExt = "Tracker"
-AbstractDifferentiationZygoteExt = "Zygote"
-
 [compat]
 ChainRulesCore = "1"
 DiffResults = "1"
@@ -37,6 +21,16 @@ ReverseDiff = "1"
 Tracker = "0.2"
 Zygote = "0.6"
 julia = "1.6"
+Enzyme = "0.11"
+
+[extensions]
+AbstractDifferentiationChainRulesCoreExt = "ChainRulesCore"
+AbstractDifferentiationFiniteDifferencesExt = "FiniteDifferences"
+AbstractDifferentiationForwardDiffExt = ["DiffResults", "ForwardDiff"]
+AbstractDifferentiationReverseDiffExt = ["DiffResults", "ReverseDiff"]
+AbstractDifferentiationTrackerExt = "Tracker"
+AbstractDifferentiationZygoteExt = "Zygote"
+AbstractDifferentiationEnzymeExt = "Enzyme"
 
 [extras]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -48,6 +42,17 @@ ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 
 [targets]
-test = ["Test", "ChainRulesCore", "DiffResults", "FiniteDifferences", "ForwardDiff", "Random", "ReverseDiff", "Tracker", "Zygote"]
+test = ["Test", "ChainRulesCore", "DiffResults", "FiniteDifferences", "ForwardDiff", "Random", "ReverseDiff", "Tracker", "Zygote", "Enzyme"]
+
+[weakdeps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
+Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"

ext/AbstractDifferentiationEnzymeExt.jl

@@ -0,0 +1,60 @@
+module AbstractDifferentiationEnzymeExt
+
+if isdefined(Base, :get_extension)
+    import AbstractDifferentiation as AD
+    using Enzyme: Enzyme
+else
+    import ..AbstractDifferentiation as AD
+    using ..Enzyme: Enzyme
+end
+
+AD.@primitive function jacobian(b::AD.EnzymeForwardBackend, f, x)
+    val = f(x)
+    if val isa Real
+        return adjoint.(AD.gradient(b, f, x))
+    else
+        if length(x) == 1 && length(val) == 1
+            # Enzyme.jacobian returns a vector of length 1 in this case
+            return (Matrix(adjoint(Enzyme.jacobian(Enzyme.Forward, f, x))),)
+        else
+            return (Enzyme.jacobian(Enzyme.Forward, f, x),)
+        end
+    end
+end
+function AD.jacobian(b::AD.EnzymeForwardBackend, f, x::Real)
+    return AD.derivative(b, f, x)
+end
+function AD.gradient(::AD.EnzymeForwardBackend, f, x::AbstractArray)
+    # Enzyme.gradient with Forward returns a tuple of the same length as the input
+    return ([Enzyme.gradient(Enzyme.Forward, f, x)...],)
+end
+function AD.gradient(b::AD.EnzymeForwardBackend, f, x::Real)
+    return AD.derivative(b, f, x)
+end
+function AD.derivative(::AD.EnzymeForwardBackend, f, x::Number)
+    (Enzyme.autodiff(Enzyme.Forward, f, Enzyme.Duplicated(x, one(x)))[1],)
+end
+
+AD.@primitive function jacobian(::AD.EnzymeReverseBackend, f, x)
+    val = f(x)
+    if val isa Real
+        return (adjoint(Enzyme.gradient(Enzyme.Reverse, f, x)),)
+    else
+        if length(x) == 1 && length(val) == 1
+            # Enzyme.jacobian returns an adjoint vector of length 1 in this case
+            return (Matrix(Enzyme.jacobian(Enzyme.Reverse, f, x, Val(1))),)
+        else
+            return (Enzyme.jacobian(Enzyme.Reverse, f, x, Val(length(val))),)
+        end
+    end
+end
+function AD.gradient(::AD.EnzymeReverseBackend, f, x::AbstractArray)
+    dx = similar(x)
+    Enzyme.gradient!(Enzyme.Reverse, dx, f, x)
+    return (dx,)
+end
+function AD.derivative(::AD.EnzymeReverseBackend, f, x::Number)
+    (Enzyme.autodiff(Enzyme.Reverse, f, Enzyme.Active(x))[1][1],)
+end
+
+end # module

src/AbstractDifferentiation.jl

@@ -637,6 +637,7 @@ end
         @require FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" include("../ext/AbstractDifferentiationFiniteDifferencesExt.jl")
         @require Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" include("../ext/AbstractDifferentiationTrackerExt.jl")
         @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" include("../ext/AbstractDifferentiationZygoteExt.jl")
+        @require Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" include("../ext/AbstractDifferentiationEnzymeExt.jl")
     end
 end
 

src/backends.jl

@@ -77,3 +77,23 @@ It is a special case of [`ReverseRuleConfigBackend`](@ref).
     To be able to use this backend, you have to load Zygote.
 """
 function ZygoteBackend end
+
+"""
+    EnzymeReverseBackend
+
+AD backend that uses reverse mode of Enzyme.jl.
+
+!!! note
+    To be able to use this backend, you have to load Enzyme.
+"""
+struct EnzymeReverseBackend <: AbstractReverseMode end
+
+"""
+    EnzymeForwardBackend
+
+AD backend that uses forward mode of Enzyme.jl.
+
+!!! note
+    To be able to use this backend, you have to load Enzyme.
+"""
+struct EnzymeForwardBackend <: AbstractForwardMode end

test/enzyme.jl

@@ -0,0 +1,47 @@
+import AbstractDifferentiation as AD
+using Test
+using Enzyme
+
+backends = [
+    "EnzymeForwardBackend" => AD.EnzymeForwardBackend(),
+    "EnzymeReverseBackend" => AD.EnzymeReverseBackend(),
+]
+
+@testset "$name" for (name, backend) in backends
+    if name == "EnzymeForwardBackend"
+        @test backend isa AD.AbstractForwardMode
+    else
+        @test backend isa AD.AbstractReverseMode
+    end
+
+    @testset "Derivative" begin
+        test_derivatives(backend, multiple_inputs = false)
+    end
+    @testset "Gradient" begin
+        test_gradients(backend, multiple_inputs = false)
+    end
+    @testset "Jacobian" begin
+        test_jacobians(backend, multiple_inputs = false)
+    end
+    # @testset "Hessian" begin
+    #     test_hessians(backend, multiple_inputs = false)
+    # end
+    @testset "jvp" begin
+        test_jvp(backend; multiple_inputs = false, vaugmented=true)
+    end
+    @testset "j′vp" begin
+        test_j′vp(backend, multiple_inputs = false)
+    end
+    @testset "Lazy Derivative" begin
+        test_lazy_derivatives(backend, multiple_inputs = false)
+    end
+    @testset "Lazy Gradient" begin
+        test_lazy_gradients(backend, multiple_inputs = false)
+    end
+    @testset "Lazy Jacobian" begin
+        test_lazy_jacobians(backend; multiple_inputs = false, vaugmented=true)
+    end
+    # @testset "Lazy Hessian" begin
+    #     test_lazy_hessians(backend, multiple_inputs = false)
+    # end
+end

test/runtests.jl

@@ -9,4 +9,5 @@ using Test
     include("finitedifferences.jl")
     include("tracker.jl")
     include("ruleconfig.jl")
+    include("enzyme.jl")
 end