Enable Enzyme (#67)

* fix enzyme to match new interface, enable enzyme tests
* fix type instability tighten Enzyme compat
* add indirection to enforce type stability of `restructure`
* fix tests enable Enzyme inference tests only on 1.10

Author: Red-Portal

.github/workflows/CI.yml

@@ -20,6 +20,7 @@ jobs:
       matrix:
         version:
           - '1.7'
+          - '1.10'
         os:
           - ubuntu-latest
           - macOS-latest

Project.toml

@@ -42,7 +42,7 @@ ChainRulesCore = "1.16"
 DiffResults = "1"
 Distributions = "0.25.87"
 DocStringExtensions = "0.8, 0.9"
-Enzyme = "0.12"
+Enzyme = "0.12.32"
 FillArrays = "1.3"
 ForwardDiff = "0.10.36"
 Functors = "0.4"

ext/AdvancedVIEnzymeExt.jl

@@ -11,13 +11,39 @@ else
     using ..AdvancedVI: ADTypes, DiffResults
 end
 
+function AdvancedVI.restructure_ad_forward(::ADTypes.AutoEnzyme, restructure, params)
+    return restructure(params)::typeof(restructure.model)
+end
+
+function AdvancedVI.value_and_gradient!(
+    ::ADTypes.AutoEnzyme, f, x::AbstractVector{<:Real}, out::DiffResults.MutableDiffResult
+)
+    Enzyme.API.runtimeActivity!(true)
+    ∇x = DiffResults.gradient(out)
+    fill!(∇x, zero(eltype(∇x)))
+    _, y = Enzyme.autodiff(
+        Enzyme.ReverseWithPrimal, Enzyme.Const(f), Enzyme.Active, Enzyme.Duplicated(x, ∇x)
+    )
+    DiffResults.value!(out, y)
+    return out
+end
+
 function AdvancedVI.value_and_gradient!(
-    ad::ADTypes.AutoEnzyme, f, θ::AbstractVector{T}, out::DiffResults.MutableDiffResult
-) where {T<:Real}
-    ∇θ = DiffResults.gradient(out)
-    fill!(∇θ, zero(T))
+    ::ADTypes.AutoEnzyme,
+    f,
+    x::AbstractVector{<:Real},
+    aux,
+    out::DiffResults.MutableDiffResult,
+)
+    Enzyme.API.runtimeActivity!(true)
+    ∇x = DiffResults.gradient(out)
+    fill!(∇x, zero(eltype(∇x)))
     _, y = Enzyme.autodiff(
-        Enzyme.ReverseWithPrimal, Enzyme.Const(f), Enzyme.Active, Enzyme.Duplicated(θ, ∇θ)
+        Enzyme.ReverseWithPrimal,
+        Enzyme.Const(f),
+        Enzyme.Active,
+        Enzyme.Duplicated(x, ∇x),
+        Enzyme.Const(aux),
     )
     DiffResults.value!(out, y)
     return out

src/AdvancedVI.jl

@@ -41,18 +41,17 @@ Evaluate the value and gradient of a function `f` at `x` using the automatic dif
 function value_and_gradient! end
 
 """
-    stop_gradient(x)
+    restructure_ad_forward(adtype, restructure, params)
 
-Stop the gradient from propagating to `x` if the selected ad backend supports it.
-Otherwise, it is equivalent to `identity`.
+Apply `restructure` to `params`.
+This is an indirection for handling the type stability of `restructure`, as some AD backends require strict type stability in the AD path.
 
 # Arguments
-- `x`: Input
-
-# Returns
-- `x`: Same value as the input.
+- `ad::ADTypes.AbstractADType`: Automatic differentiation backend. 
+- `restructure`: Callable for restructuring the varitional distribution from `params`.
+- `params`: Variational Parameters.
 """
-function stop_gradient end
+restructure_ad_forward(::ADTypes.AbstractADType, restructure, params) = restructure(params)
 
 # Update for gradient descent step
 """

src/families/location_scale.jl

@@ -38,14 +38,14 @@ Functors.@functor MvLocationScale (location, scale)
 # is very inefficient.
 # begin
 struct RestructureMeanField{S<:Diagonal,D,L}
-    q::MvLocationScale{S,D,L}
+    model::MvLocationScale{S,D,L}
 end
 
 function (re::RestructureMeanField)(flat::AbstractVector)
     n_dims = div(length(flat), 2)
     location = first(flat, n_dims)
     scale = Diagonal(last(flat, n_dims))
-    return MvLocationScale(location, scale, re.q.dist, re.q.scale_eps)
+    return MvLocationScale(location, scale, re.model.dist, re.model.scale_eps)
 end
 
 function Optimisers.destructure(q::MvLocationScale{<:Diagonal,D,L}) where {D,L}

src/objectives/elbo/repgradelbo.jl

@@ -93,8 +93,8 @@ function estimate_objective(obj::RepGradELBO, q, prob; n_samples::Int=obj.n_samp
 end
 
 function estimate_repgradelbo_ad_forward(params′, aux)
-    @unpack rng, obj, problem, restructure, q_stop = aux
-    q = restructure(params′)
+    @unpack rng, obj, problem, adtype, restructure, q_stop = aux
+    q = restructure_ad_forward(adtype, restructure, params′)
     samples, entropy = reparam_with_entropy(rng, q, q_stop, obj.n_samples, obj.entropy)
     energy = estimate_energy_with_samples(problem, samples)
     elbo = energy + entropy
@@ -112,7 +112,14 @@ function estimate_gradient!(
     state,
 )
     q_stop = restructure(params)
-    aux = (rng=rng, obj=obj, problem=prob, restructure=restructure, q_stop=q_stop)
+    aux = (
+        rng=rng,
+        adtype=adtype,
+        obj=obj,
+        problem=prob,
+        restructure=restructure,
+        q_stop=q_stop,
+    )
     value_and_gradient!(adtype, estimate_repgradelbo_ad_forward, params, aux, out)
     nelbo = DiffResults.value(out)
     stat = (elbo=-nelbo,)

test/Project.toml

@@ -26,9 +26,10 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 [compat]
 ADTypes = "0.2.1, 1"
 Bijectors = "0.13"
+DiffResults = "1.0"
 Distributions = "0.25.100"
 DistributionsAD = "0.6.45"
-Enzyme = "0.12"
+Enzyme = "0.12.32"
 FillArrays = "1.6.1"
 ForwardDiff = "0.10.36"
 Functors = "0.4.5"

test/inference/repgradelbo_distributionsad.jl

@@ -1,4 +1,19 @@
 
+AD_distributionsad = if VERSION >= v"1.10"
+    Dict(
+        :ForwarDiff => AutoForwardDiff(),
+        #:ReverseDiff => AutoReverseDiff(), # DistributionsAD doesn't support ReverseDiff at the moment
+        :Zygote => AutoZygote(),
+        :Enzyme => AutoEnzyme(),
+    )
+else
+    Dict(
+        :ForwarDiff => AutoForwardDiff(),
+        #:ReverseDiff => AutoReverseDiff(), # DistributionsAD doesn't support ReverseDiff at the moment
+        :Zygote => AutoZygote(),
+    )
+end
+
 @testset "inference RepGradELBO DistributionsAD" begin
     @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in
                                                                      [Float64, Float32],
@@ -9,12 +24,7 @@
             :RepGradELBOStickingTheLanding =>
                 RepGradELBO(n_montecarlo; entropy=StickingTheLandingEntropy()),
         ),
-        (adbackname, adtype) in Dict(
-            :ForwarDiff => AutoForwardDiff(),
-            #:ReverseDiff => AutoReverseDiff(),
-            :Zygote => AutoZygote(),
-            #:Enzyme      => AutoEnzyme(),
-        )
+        (adbackname, adtype) in AD_distributionsad
 
         seed = (0x38bef07cf9cc549d)
         rng = StableRNG(seed)
@@ -31,8 +41,8 @@
         # where ρ = 1 - ημ, μ is the strong convexity constant.
         contraction_rate = 1 - η * strong_convexity
 
-        μ0 = Zeros(realtype, n_dims)
-        L0 = Diagonal(Ones(realtype, n_dims))
+        μ0 = zeros(realtype, n_dims)
+        L0 = Diagonal(ones(realtype, n_dims))
         q0 = TuringDiagMvNormal(μ0, diag(L0))
 
         @testset "convergence" begin

test/inference/repgradelbo_locationscale.jl

@@ -1,4 +1,19 @@
 
+AD_locationscale = if VERSION >= v"1.10"
+    Dict(
+        :ForwarDiff => AutoForwardDiff(),
+        :ReverseDiff => AutoReverseDiff(),
+        :Zygote => AutoZygote(),
+        :Enzyme => AutoEnzyme(),
+    )
+else
+    Dict(
+        :ForwarDiff => AutoForwardDiff(),
+        :ReverseDiff => AutoReverseDiff(),
+        :Zygote => AutoZygote(),
+    )
+end
+
 @testset "inference RepGradELBO VILocationScale" begin
     @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in
                                                                      [Float64, Float32],
@@ -10,12 +25,7 @@
             :RepGradELBOStickingTheLanding =>
                 RepGradELBO(n_montecarlo; entropy=StickingTheLandingEntropy()),
         ),
-        (adbackname, adtype) in Dict(
-            :ForwarDiff => AutoForwardDiff(),
-            :ReverseDiff => AutoReverseDiff(),
-            :Zygote => AutoZygote(),
-            #:Enzyme      => AutoEnzyme(),
-        )
+        (adbackname, adtype) in AD_locationscale
 
         seed = (0x38bef07cf9cc549d)
         rng = StableRNG(seed)

test/inference/repgradelbo_locationscale_bijectors.jl

@@ -1,4 +1,19 @@
 
+AD_locationscale_bijectors = if VERSION >= v"1.10"
+    Dict(
+        :ForwarDiff => AutoForwardDiff(),
+        :ReverseDiff => AutoReverseDiff(),
+        :Zygote => AutoZygote(),
+        :Enzyme => AutoEnzyme(),
+    )
+else
+    Dict(
+        :ForwarDiff => AutoForwardDiff(),
+        :ReverseDiff => AutoReverseDiff(),
+        :Zygote => AutoZygote(),
+    )
+end
+
 @testset "inference RepGradELBO VILocationScale Bijectors" begin
     @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in
                                                                      [Float64, Float32],
@@ -10,12 +25,7 @@
             :RepGradELBOStickingTheLanding =>
                 RepGradELBO(n_montecarlo; entropy=StickingTheLandingEntropy()),
         ),
-        (adbackname, adtype) in Dict(
-            :ForwarDiff => AutoForwardDiff(),
-            :ReverseDiff => AutoReverseDiff(),
-            #:Zygote      => AutoZygote(), 
-            #:Enzyme      => AutoEnzyme(),
-        )
+        (adbackname, adtype) in AD_locationscale_bijectors
 
         seed = (0x38bef07cf9cc549d)
         rng = StableRNG(seed)

test/interface/ad.jl

@@ -19,4 +19,23 @@ using Test
         @test ∇ ≈ (A + A') * λ / 2
         @test f ≈ λ' * A * λ / 2
     end
+
+    @testset "$(adname) with auxiliary input" for (adname, adsymbol) in Dict(
+        :ForwardDiff => AutoForwardDiff(),
+        :ReverseDiff => AutoReverseDiff(),
+        :Zygote => AutoZygote(),
+        :Enzyme => AutoEnzyme(),
+    )
+        D = 10
+        A = randn(D, D)
+        λ = randn(D)
+        b = randn(D)
+        grad_buf = DiffResults.GradientResult(λ)
+        f(λ′, aux) = λ′' * A * λ′ / 2 + dot(aux.b, λ′)
+        AdvancedVI.value_and_gradient!(adsymbol, f, λ, (b=b,), grad_buf)
+        ∇ = DiffResults.gradient(grad_buf)
+        f = DiffResults.value(grad_buf)
+        @test ∇ ≈ (A + A') * λ / 2 + b
+        @test f ≈ λ' * A * λ / 2 + dot(b, λ)
+    end
 end

test/interface/repgradelbo.jl

@@ -35,7 +35,10 @@ end
     @unpack model, μ_true, L_true, n_dims, is_meanfield = modelstats
 
     @testset for ad in [
-        ADTypes.AutoForwardDiff(), ADTypes.AutoReverseDiff(), ADTypes.AutoZygote()
+        ADTypes.AutoForwardDiff(),
+        ADTypes.AutoReverseDiff(),
+        ADTypes.AutoZygote(),
+        ADTypes.AutoEnzyme(),
     ]
         q_true = MeanFieldGaussian(
             Vector{eltype(μ_true)}(μ_true), Diagonal(Vector{eltype(L_true)}(diag(L_true)))
@@ -44,7 +47,7 @@ end
         obj = RepGradELBO(10; entropy=StickingTheLandingEntropy())
         out = DiffResults.DiffResult(zero(eltype(params)), similar(params))
 
-        aux = (rng=rng, obj=obj, problem=model, restructure=re, q_stop=q_true)
+        aux = (rng=rng, obj=obj, problem=model, restructure=re, q_stop=q_true, adtype=ad)
         AdvancedVI.value_and_gradient!(
             ad, AdvancedVI.estimate_repgradelbo_ad_forward, params, aux, out
         )

test/models/normal.jl

@@ -35,9 +35,7 @@ function normal_meanfield(rng::Random.AbstractRNG, realtype::Type)
 
     σ0 = realtype(0.3)
     μ = Fill(realtype(5), n_dims)
-    #randn(rng, realtype, n_dims)
     σ = Fill(σ0, n_dims)
-    #log.(exp.(randn(rng, realtype, n_dims)) .+ 1)
 
     model = TestNormal(μ, Diagonal(σ .^ 2))
 

test/runtests.jl

@@ -20,7 +20,7 @@ using DistributionsAD
 using LogDensityProblems
 using Optimisers
 using ADTypes
-using Enzyme, ForwardDiff, ReverseDiff, Zygote
+using ForwardDiff, ReverseDiff, Zygote, Enzyme
 
 using AdvancedVI