Explicit Enzyme rules

TuringLang · Nov 8, 2024 · a2ff88a · a2ff88a
1 parent e99db76
commit a2ff88a
Show file tree

Hide file tree

Showing 6 changed files with 75 additions and 20 deletions.
diff --git a/Project.toml b/Project.toml
@@ -25,7 +25,6 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [weakdeps]
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
-Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"
@@ -36,7 +35,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 BijectorsDistributionsADExt = "DistributionsAD"
-BijectorsEnzymeExt = ["Enzyme", "EnzymeCore"]
+BijectorsEnzymeCoreExt = ["EnzymeCore"]
 BijectorsForwardDiffExt = "ForwardDiff"
 BijectorsLazyArraysExt = "LazyArrays"
 BijectorsMooncakeExt = "Mooncake"

diff --git a/ext/BijectorsEnzymeCoreExt.jl b/ext/BijectorsEnzymeCoreExt.jl
@@ -0,0 +1,55 @@
+module BijectorsEnzymeCoreExt
+
+if isdefined(Base, :get_extension)
+    using EnzymeCore: Const, Duplicated, DuplicatedNoNeed, EnzymeRules
+    using Bijectors: find_alpha
+else
+    using ..EnzymeCore: Const, Duplicated, DuplicatedNoNeed, EnzymeRules
+    using ..Bijectors: find_alpha
+end
+
+_add(x::Real, ::Nothing) = x
+_add(x::Real, y::Real) = x + y
+
+_muladd(x::Real, y::Real, ::Nothing) = x * y
+_muladd(x::Real, y::Real, z::Real) = muladd(x, y, z)
+
+function EnzymeRules.forward(
+    ::Const{typeof(find_alpha)},
+    ::Type{RT},
+    wt_y::Union{Const,Duplicated},
+    wt_u_hat::Union{Const,Duplicated},
+    b::Union{Const,Duplicated},
+) where {RT<:Union{Const,DuplicatedNoNeed,Duplicated}}
+    # Compute primal value
+    Ω = find_alpha(wt_y.val, wt_u_hat.val, b.val)
+
+    # Early exit if no derivatives are requested
+    if RT <: Const
+        return Ω
+    end
+
+    if wt_y isa Const && wt_u_hat isa Const && b isa Const
+        # Trivial case: All partial derivatives are 0
+        Ω̇ = zero(Ω)
+    else
+        # In all other cases we have to compute the partial derivatives
+        # As in the rule for ChainRules, we reuse the following term in the computation of all derivatives
+        c = wt_u_hat.val * sech(Ω + b.val)^2
+
+        x = b isa Duplicated ? (-c) * b.dval : nothing
+        y = wt_u_hat isa Duplicated ? _muladd(-tanh(Ω + b.val), wt_u_hat.dval, x) : x
+        z = wt_y isa Duplicated ? _add(wt_y.dval, y) : y
+
+        Ω̇ = z / (1 + c)
+    end
+
+    if RT <: Duplicated
+        return Duplicated(Ω, Ω̇)
+    else
+        @assert RT <: DuplicatedNoNeed
+        return Ω̇
+    end
+end
+
+end  # module
diff --git a/ext/BijectorsEnzymeExt.jl b/ext/BijectorsEnzymeExt.jl
diff --git a/test/Project.toml b/test/Project.toml
@@ -7,6 +7,7 @@ Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+EnzymeTestUtils = "12d8515a-0907-448a-8884-5fe00fdf1c5a"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -33,6 +34,7 @@ Combinatorics = "1.0.2"
 Compat = "3.46, 4.2"
 DistributionsAD = "0.6.3"
 Enzyme = "0.12.22"
+EnzymeTestUtils = "0.1.8"
 FillArrays = "1"
 FiniteDifferences = "0.11, 0.12"
 ForwardDiff = "0.10.12"

diff --git a/test/ad/enzyme.jl b/test/ad/enzyme.jl
@@ -0,0 +1,15 @@
+@testset "Enzyme: Bijectors.find_alpha" begin
+    x = randn()
+    y = expm1(randn())
+    z = randn()
+
+    @testset "forward" begin
+        @testset for RT in (Const, DuplicatedNoNeed, Duplicated),
+            Tx in (Const, Duplicated),
+            Ty in (Const, Duplicated),
+            Tz in (Const, Duplicated)
+
+            test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz))
+        end
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -4,6 +4,7 @@ using ChainRulesTestUtils
 using Combinatorics
 using DistributionsAD
 using Enzyme
+using EnzymeTestUtils
 using FiniteDifferences
 using ForwardDiff
 using Functors
@@ -68,6 +69,7 @@ end
 
 if GROUP == "All" || GROUP == "AD"
     include("ad/chainrules.jl")
+    include("ad/enzyme.jl")
     include("ad/flows.jl")
     include("ad/pd.jl")
     include("ad/corr.jl")