Merge pull request #84 from devmotion/dw/enzyme

Add support for `Enzyme`

Author: tpapp

Project.toml

@@ -1,7 +1,7 @@
 name = "LogDensityProblems"
 uuid = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 authors = ["Tamas K. Papp <tkpapp@gmail.com>"]
-version = "0.11.3"
+version = "0.11.4"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
@@ -27,6 +27,7 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
@@ -35,4 +36,4 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["BenchmarkTools", "Distributions", "Documenter", "ForwardDiff", "Test", "StatsBase", "StatsFuns", "Tracker", "Zygote", "ReverseDiff"]
+test = ["BenchmarkTools", "Distributions", "Documenter", "ForwardDiff", "Pkg", "ReverseDiff", "StatsBase", "StatsFuns", "Test", "Tracker", "Zygote"]

src/AD_Enzyme.jl

@@ -0,0 +1,61 @@
+import .Enzyme
+
+struct EnzymeGradientLogDensity{L,M<:Union{Enzyme.ForwardMode,Enzyme.ReverseMode},S} <: ADGradientWrapper
+    ℓ::L
+    mode::M
+    shadow::S # only used in forward mode
+end
+
+"""
+    ADgradient(:Enzyme, ℓ; kwargs...)
+    ADgradient(Val(:Enzyme), ℓ; kwargs...)
+
+Gradient using algorithmic/automatic differentiation via Enzyme.
+
+# Keyword arguments
+
+- `mode::Enzyme.Mode`: Differentiation mode (default: `Enzyme.Reverse`).
+  Currently only `Enzyme.Reverse` and `Enzyme.Forward` are supported.
+
+- `shadow`: Collection of one-hot vectors for each entry of the inputs `x` to the log density
+  `ℓ`, or `nothing` (default: `nothing`). This keyword argument is only used in forward
+  mode. By default, it will be recomputed in every call of `logdensity_and_gradient(ℓ, x)`.
+  For performance reasons it is recommended to compute it only once when calling `ADgradient`.
+  The one-hot vectors can be constructed, e.g., with `Enzyme.onehot(x)`.
+"""
+function ADgradient(::Val{:Enzyme}, ℓ; mode::Enzyme.Mode = Enzyme.Reverse, shadow = nothing)
+    mode isa Union{Enzyme.ForwardMode,Enzyme.ReverseMode} ||
+        throw(ArgumentError("currently automatic differentiation via Enzyme only supports " *
+                            "`Enzyme.Forward` and `Enzyme.Reverse` modes"))
+    if mode isa Enzyme.ReverseMode && shadow !== nothing
+        @info "keyword argument `shadow` is ignored in reverse mode"
+        shadow = nothing
+    end
+    return EnzymeGradientLogDensity(ℓ, mode, shadow)
+end
+
+function Base.show(io::IO, ∇ℓ::EnzymeGradientLogDensity)
+    print(io, "Enzyme AD wrapper for ", ∇ℓ.ℓ, " with ",
+          ∇ℓ.mode isa Enzyme.ForwardMode ? "forward" : "reverse", " mode")
+end
+
+function logdensity_and_gradient(∇ℓ::EnzymeGradientLogDensity{<:Any,<:Enzyme.ForwardMode},
+                                 x::AbstractVector)
+    @unpack ℓ, mode, shadow = ∇ℓ
+    _shadow = shadow === nothing ? Enzyme.onehot(x) : shadow
+    y, ∂ℓ_∂x = Enzyme.autodiff(mode, Base.Fix1(logdensity, ℓ), Enzyme.BatchDuplicated,
+                               Enzyme.BatchDuplicated(x, _shadow))
+    return y, collect(∂ℓ_∂x)
+end
+
+function logdensity_and_gradient(∇ℓ::EnzymeGradientLogDensity{<:Any,<:Enzyme.ReverseMode},
+                                 x::AbstractVector)
+    @unpack ℓ, mode = ∇ℓ
+    # Currently it is not possible to retrieve the primal together with the derivatives.
+    # Ref: https://github.com/EnzymeAD/Enzyme.jl/issues/107
+    y = logdensity(ℓ, x)
+    ∂ℓ_∂x = zero(x)
+    Enzyme.autodiff(mode, Base.Fix1(logdensity, ℓ), Enzyme.Active,
+                    Enzyme.Duplicated(x, ∂ℓ_∂x))
+    y, ∂ℓ_∂x
+end

src/LogDensityProblems.jl

@@ -241,6 +241,7 @@ function __init__()
     @require Tracker="9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" include("AD_Tracker.jl")
     @require Zygote="e88e6eb3-aa80-5325-afca-941959d7151f" include("AD_Zygote.jl")
     @require ReverseDiff="37e2e3b7-166d-5795-8a7a-e32c996b4267" include("AD_ReverseDiff.jl")
+    @require Enzyme="7da242da-08ed-463a-9acd-ee780be4f1d9" include("AD_Enzyme.jl")
 end
 
 ####

test/runtests.jl

@@ -1,3 +1,14 @@
+@static if VERSION >= v"1.6"
+    # Enzyme only supports Julia >= 1.6
+    # We add it without messing with the existing, possibly precompiled, dependencies
+    using Pkg
+    Pkg.add(Pkg.PackageSpec(; name="Enzyme", uuid="7da242da-08ed-463a-9acd-ee780be4f1d9");
+            preserve=Pkg.PRESERVE_ALL)
+
+    import Enzyme
+    struct EnzymeTestMode <: Enzyme.Mode end
+end
+
 using LogDensityProblems, Test, Distributions, TransformVariables, BenchmarkTools
 import LogDensityProblems: capabilities, dimension, logdensity
 using LogDensityProblems: logdensity_and_gradient, LogDensityOrder
@@ -149,6 +160,39 @@ end
     end
 end
 
+@static if VERSION >= v"1.6"
+    @testset "AD via Enzyme" begin
+        ℓ = TestLogDensity(test_logdensity1)
+
+        ∇ℓ_reverse = ADgradient(:Enzyme, ℓ)
+        @test ∇ℓ_reverse === ADgradient(:Enzyme, ℓ; mode=Enzyme.Reverse)
+        @test repr(∇ℓ_reverse) == "Enzyme AD wrapper for " * repr(ℓ) * " with reverse mode"
+
+        ∇ℓ_forward = ADgradient(:Enzyme, ℓ; mode=Enzyme.Forward)
+        ∇ℓ_forward_shadow = ADgradient(:Enzyme, ℓ;
+                                       mode=Enzyme.Forward,
+                                       shadow=Enzyme.onehot(Vector{Float64}(undef, dimension(ℓ))))
+        for ∇ℓ in (∇ℓ_forward, ∇ℓ_forward_shadow)
+            @test repr(∇ℓ) == "Enzyme AD wrapper for " * repr(ℓ) * " with forward mode"
+        end
+
+        for ∇ℓ in (∇ℓ_reverse, ∇ℓ_forward, ∇ℓ_forward_shadow)
+            @test dimension(∇ℓ) == 3
+            @test capabilities(∇ℓ) ≡ LogDensityOrder(1)
+            for _ in 1:100
+                x = randn(3)
+                @test @inferred(logdensity(∇ℓ, x)) ≅ test_logdensity1(x)
+                @test logdensity_and_gradient(∇ℓ, x) ≅ (test_logdensity1(x), test_gradient(x))
+            end
+        end
+
+        # Branches in `ADgradient`
+        @test_throws ArgumentError ADgradient(:Enzyme, ℓ; mode=EnzymeTestMode())
+        ∇ℓ = @test_logs (:info, "keyword argument `shadow` is ignored in reverse mode") ADgradient(:Enzyme, ℓ; shadow = (1,))
+        @test ∇ℓ.shadow === nothing
+    end
+end
+
 @testset "ADgradient missing method" begin
     msg = "Don't know how to AD with Foo, consider `import Foo` if there is such a package."
     P = TransformedLogDensity(as(Array, 1), x -> sum(abs2, x))