Explicit Enzyme rules

.github/workflows/AD.yml

@@ -14,6 +14,7 @@ jobs:
       matrix:
         version:
           - '1.6'
+          - '1.10'
           - '1'
         os:
           - ubuntu-latest

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"

ext/BijectorsEnzymeCoreExt.jl

@@ -0,0 +1,222 @@
+module BijectorsEnzymeCoreExt
+
+if isdefined(Base, :get_extension)
+    using EnzymeCore:
+        Active,
+        Const,
+        Duplicated,
+        DuplicatedNoNeed,
+        BatchDuplicated,
+        BatchDuplicatedNoNeed,
+        EnzymeRules
+    using Bijectors: find_alpha
+else
+    using ..EnzymeCore:
+        Active,
+        Const,
+        Duplicated,
+        DuplicatedNoNeed,
+        BatchDuplicated,
+        BatchDuplicatedNoNeed,
+        EnzymeRules
+    using ..Bijectors: find_alpha
+end
+
+# Compute a tuple of partial derivatives wrt non-`Const` arguments
+# and `nothing`s for `Const` arguments
+function ∂find_alpha(
+    Ω::Real,
+    wt_y::Union{Const,Active,Duplicated,BatchDuplicated},
+    wt_u_hat::Union{Const,Active,Duplicated,BatchDuplicated},
+    b::Union{Const,Active,Duplicated,BatchDuplicated},
+)
+    # We reuse the following term in the computation of the derivatives
+    Ωpb = Ω + b.val
+    c = wt_u_hat.val * sech(Ωpb)^2
+    cp1 = c + 1
+
+    ∂Ω_∂wt_y = wt_y isa Const ? nothing : oneunit(wt_y.val) / cp1
+    ∂Ω_∂wt_u_hat = wt_u_hat isa Const ? nothing : -tanh(Ωpb) / cp1
+    ∂Ω_∂b = b isa Const ? nothing : -c / cp1
+
+    return (∂Ω_∂wt_y, ∂Ω_∂wt_u_hat, ∂Ω_∂b)
+end
+
+# `muladd` for partial derivatives that can deal with `nothing` derivatives
+_muladd_partial(::Nothing, ::Const, x::Union{Real,Tuple{Vararg{Real}},Nothing}) = x
+_muladd_partial(x::Real, y::Duplicated, z::Real) = muladd(x, y.dval, z)
+_muladd_partial(x::Real, y::Duplicated, ::Nothing) = x * y.dval
+function _muladd_partial(x::Real, y::BatchDuplicated{<:Real,N}, z::NTuple{N,Real}) where {N}
+    let x = x
+        map((a, b) -> muladd(x, a, b), y.dval, z)
+    end
+end
+_muladd_partial(x::Real, y::BatchDuplicated, ::Nothing) = map(Base.Fix1(*, x), y.dval)
+
+function EnzymeRules.forward(
+    ::Const{typeof(find_alpha)},
+    ::Type{RT},
+    wt_y::Union{Const,Duplicated,BatchDuplicated},
+    wt_u_hat::Union{Const,Duplicated,BatchDuplicated},
+    b::Union{Const,Duplicated,BatchDuplicated},
+) where {RT<:Union{Const,Duplicated,DuplicatedNoNeed,BatchDuplicated,BatchDuplicatedNoNeed}}
+    # Check that the types of the activities are consistent
+    if !(
+        RT <: Union{Const,Duplicated,DuplicatedNoNeed} &&
+        wt_y isa Union{Const,Duplicated} &&
+        wt_u_hat isa Union{Const,Duplicated} &&
+        b isa Union{Const,Duplicated}
+    ) && !(
+        RT <: Union{Const,BatchDuplicated,BatchDuplicatedNoNeed} &&
+        wt_y isa Union{Const,BatchDuplicated} &&
+        wt_u_hat isa Union{Const,BatchDuplicated} &&
+        b isa Union{Const,BatchDuplicated}
+    )
+        throw(ArgumentError("inconsistent activities"))
+    end
+
+    # 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
+        ∂Ω_∂wt_y, ∂Ω_∂wt_u_hat, ∂Ω_∂b = ∂find_alpha(Ω, wt_y, wt_u_hat, b)
+        _muladd_partial(
+            ∂Ω_∂wt_y,
+            wt_y,
+            _muladd_partial(∂Ω_∂wt_u_hat, wt_u_hat, _muladd_partial(∂Ω_∂b, b, nothing)),
+        )
+    end
+
+    if RT <: Duplicated
+        @assert Ω̇ isa Real
+        return Duplicated(Ω, Ω̇)
+    elseif RT <: DuplicatedNoNeed
+        @assert Ω̇ isa Real
+        return Ω̇
+    elseif RT <: BatchDuplicated
+        @assert Ω̇ isa Tuple{Vararg{Real}}
+        return BatchDuplicated(Ω, Ω̇)
+    else
+        @assert RT <: BatchDuplicatedNoNeed
+        @assert Ω̇ isa Tuple{Vararg{Real}}
+        return Ω̇
+    end
+end
+
+struct Zero{T}
+    x::T
+end
+(f::Zero)(_) = zero(f.x)
+
+function EnzymeRules.augmented_primal(
+    config::EnzymeRules.Config,
+    ::Const{typeof(find_alpha)},
+    ::Type{RT},
+    wt_y::Union{Const,Active},
+    wt_u_hat::Union{Const,Active},
+    b::Union{Const,Active},
+) where {RT<:Union{Const,Active}}
+    # Only compute the the original return value if it is actually needed
+    Ω =
+        if EnzymeRules.needs_primal(config) ||
+            EnzymeRules.needs_shadow(config) ||
+            !(RT <: Const || (wt_y isa Const && wt_u_hat isa Const && b isa Const))
+            find_alpha(wt_y.val, wt_u_hat.val, b.val)
+        else
+            nothing
+        end
+
+    tape = if RT <: Const || (wt_y isa Const && wt_u_hat isa Const && b isa Const)
+        # Trivial case: No differentiation or all derivatives are 0
+        # Thus no tape is needed
+        nothing
+    else
+        # Derivatives with respect to at least one argument needed
+        # They are computed in the reverse pass, and therefore the original return is cached
+        # In principle, the partial derivatives could be computed here and be cached
+        # But Enzyme only executes the reverse pass once,
+        # thus this would not increase efficiency but instead more values would have to be cached
+        Ω
+    end
+
+    # Ensure that we follow the interface requirements of `augmented_primal`
+    primal = EnzymeRules.needs_primal(config) ? Ω : nothing
+    shadow = if EnzymeRules.needs_shadow(config)
+        if EnzymeRules.width(config) === 1
+            zero(Ω)
+        else
+            ntuple(Zero(Ω), Val(EnzymeRules.width(config)))
+        end
+    else
+        nothing
+    end
+
+    return EnzymeRules.AugmentedReturn(primal, shadow, tape)
+end
+
+struct ZeroOrNothing{N} end
+(::ZeroOrNothing)(::Const) = nothing
+(::ZeroOrNothing{1})(x::Active) = zero(x.val)
+(::ZeroOrNothing{N})(x::Active) where {N} = ntuple(Zero(x.val), Val{N}())
+
+function EnzymeRules.reverse(
+    config::EnzymeRules.Config,
+    ::Const{typeof(find_alpha)},
+    ::Type{<:Const},
+    ::Nothing,
+    wt_y::Union{Const,Active},
+    wt_u_hat::Union{Const,Active},
+    b::Union{Const,Active},
+)
+    # Trivial case: Nothing to be differentiated (return activity is `Const`)
+    return map(ZeroOrNothing{EnzymeRules.width(config)}(), (wt_y, wt_u_hat, b))
+end
+function EnzymeRules.reverse(
+    ::EnzymeRules.Config,
+    ::Const{typeof(find_alpha)},
+    ::Active,
+    ::Nothing,
+    ::Const,
+    ::Const,
+    ::Const,
+)
+    # Trivial case: Tape does not exist sice all partial derivatives are 0
+    return (nothing, nothing, nothing)
+end
+
+struct MulPartialOrNothing{T<:Union{Real,Tuple{Vararg{Real}}}}
+    x::T
+end
+(::MulPartialOrNothing)(::Nothing) = nothing
+(f::MulPartialOrNothing{<:Real})(∂f_∂x::Real) = ∂f_∂x * f.x
+function (f::MulPartialOrNothing{<:NTuple{N,Real}})(∂f_∂x::Real) where {N}
+    return map(Base.Fix1(*, ∂f_∂x), f.x)
+end
+
+function EnzymeRules.reverse(
+    ::EnzymeRules.Config,
+    ::Const{typeof(find_alpha)},
+    ΔΩ::Active,
+    Ω::Real,
+    wt_y::Union{Const,Active},
+    wt_u_hat::Union{Const,Active},
+    b::Union{Const,Active},
+)
+    # Tape must be `nothing` if all arguments are `Const`
+    @assert !(wt_y isa Const && wt_u_hat isa Const && b isa Const)
+
+    # Compute partial derivatives
+    ∂Ω_∂xs = ∂find_alpha(Ω, wt_y, wt_u_hat, b)
+    return map(MulPartialOrNothing(ΔΩ.val), ∂Ω_∂xs)
+end
+
+end  # module

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"

test/ad/enzyme.jl

@@ -0,0 +1,38 @@
+# Segfaults on older Julia versions, probably never supported
+# TODO: Enable tests on Julia >= 1.11 when updating to Enzyme 0.13
+if v"1.10" <= VERSION < v"1.11"
+    @testset "Enzyme: Bijectors.find_alpha" begin
+        x = randn()
+        y = expm1(randn())
+        z = randn()
+
+        @testset "forward" begin
+            # No batches
+            @testset for RT in (Const, Duplicated, DuplicatedNoNeed),
+                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
+
+            # Batches
+            @testset for RT in (Const, BatchDuplicated, BatchDuplicatedNoNeed),
+                Tx in (Const, BatchDuplicated),
+                Ty in (Const, BatchDuplicated),
+                Tz in (Const, BatchDuplicated)
+
+                test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz))
+            end
+        end
+        @testset "reverse" begin
+            @testset for RT in (Const, Active),
+                Tx in (Const, Active),
+                Ty in (Const, Active),
+                Tz in (Const, Active)
+
+                test_reverse(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz))
+            end
+        end
+    end
+end

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")