Support Mixing AD Frameworks for LogDensityProblems and the objective

Project.toml

@@ -5,6 +5,7 @@ version = "0.5.0"
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -20,31 +21,44 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [weakdeps]
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 
 [extensions]
-AdvancedVIBijectorsExt = "Bijectors"
+AdvancedVIBijectorsExt = ["Bijectors", "Optimisers"]
+AdvancedVIEnzymeExt = ["Enzyme", "ChainRulesCore"]
+AdvancedVIMooncakeExt = ["Mooncake", "ChainRulesCore"]
+AdvancedVIReverseDiffExt = ["ReverseDiff", "ChainRulesCore"]
 
 [compat]
 ADTypes = "1"
 Accessors = "0.1"
 Bijectors = "0.13, 0.14, 0.15"
+ChainRulesCore = "1"
 DiffResults = "1"
 DifferentiationInterface = "0.6, 0.7"
 Distributions = "0.25.111"
 DocStringExtensions = "0.8, 0.9"
+Enzyme = "0.13"
 FillArrays = "1.3"
 Functors = "0.4, 0.5"
 LinearAlgebra = "1"
 LogDensityProblems = "2"
+Mooncake = "0.4"
 Optimisers = "0.2.16, 0.3, 0.4"
 ProgressMeter = "1.6"
 Random = "1"
+ReverseDiff = "1"
 StatsBase = "0.32, 0.33, 0.34"
 julia = "1.10, 1.11.2"
 
 [extras]
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]

bench/benchmarks.jl

@@ -47,10 +47,15 @@ begin
             ],
             (adname, adtype) in [
                 ("Zygote", AutoZygote()),
-                ("ForwardDiff", AutoForwardDiff()),
-                ("ReverseDiff", AutoReverseDiff()),
                 ("Mooncake", AutoMooncake(; config=Mooncake.Config())),
-                # ("Enzyme", AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=Enzyme.Const)),
+                ("ReverseDiff", AutoReverseDiff()),
+                (
+                    "Enzyme",
+                    AutoEnzyme(;
+                        mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
+                        function_annotation=Enzyme.Const,
+                    ),
+                ),
             ],
             (familyname, family) in [
                 ("meanfield", MeanFieldGaussian(zeros(T, d), Diagonal(ones(T, d)))),

bench/normallognormal.jl

@@ -12,12 +12,19 @@ function LogDensityProblems.logdensity(model::NormalLogNormal, θ)
     return log_density_x + log_density_y
 end
 
+function LogDensityProblems.logdensity_and_gradient(model::NormalLogNormal, θ)
+    return (
+        LogDensityProblems.logdensity(model, θ),
+        ForwardDiff.gradient(Base.Fix1(LogDensityProblems.logdensity, model), θ),
+    )
+end
+
 function LogDensityProblems.dimension(model::NormalLogNormal)
     return length(model.μ_y) + 1
 end
 
 function LogDensityProblems.capabilities(::Type{<:NormalLogNormal})
-    return LogDensityProblems.LogDensityOrder{0}()
+    return LogDensityProblems.LogDensityOrder{1}()
 end
 
 function Bijectors.bijector(model::NormalLogNormal)

docs/Project.toml

@@ -7,6 +7,7 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
+LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 QuasiMonteCarlo = "8a4e6c94-4038-4cdc-81c3-7e6ffdb2a71b"

docs/src/examples.md

@@ -51,41 +51,51 @@ model = NormalLogNormal(μ_x, σ_x, μ_y, Diagonal(σ_y .^ 2));
 nothing
 ```
 
-Since the `y` follows a log-normal prior, its support is bounded to be the positive half-space ``\mathbb{R}_+``.
-Thus, we will use [Bijectors](https://github.com/TuringLang/Bijectors.jl) to match the support of our target posterior and the variational approximation.
+Some of the VI algorithms require gradients of the target log-density.
+In this example, we will use `KLMinRepGradDescent`, which requires first-order differentiation [capability](https://www.tamaspapp.eu/LogDensityProblems.jl/stable/#LogDensityProblems.capabilities).
+For this, we can rely on `LogDensityProblemsAD`:
 
 ```@example elboexample
-using Bijectors
-
-function Bijectors.bijector(model::NormalLogNormal)
-    (; μ_x, σ_x, μ_y, Σ_y) = model
-    return Bijectors.Stacked(
-        Bijectors.bijector.([LogNormal(μ_x, σ_x), MvNormal(μ_y, Σ_y)]),
-        [1:1, 2:(1 + length(μ_y))],
-    )
-end
+using LogDensityProblemsAD
+using ADTypes, ReverseDiff
 
-b = Bijectors.bijector(model);
-binv = inverse(b)
+model_ad = ADgradient(AutoReverseDiff(), model)
 nothing
 ```
 
 Let's now load `AdvancedVI`.
-Since BBVI relies on automatic differentiation (AD), we need to load an AD library, *before* loading `AdvancedVI`.
-Also, the selected AD framework needs to be communicated to `AdvancedVI` using the [ADTypes](https://github.com/SciML/ADTypes.jl) interface.
+In addition to gradients of the target log-density, `KLMinRepGradDescent` internally uses automatic differentiation.
+Therefore, we have to select an AD framework to be used within `KLMinRepGradDescent`.
+(This does not need to be the same as the backend used by `model_ad`.)
+The selected AD framework needs to be communicated to `AdvancedVI` using the [ADTypes](https://github.com/SciML/ADTypes.jl) interface.
 Here, we will use `ForwardDiff`, which can be selected by later passing `ADTypes.AutoForwardDiff()`.
 
 ```@example elboexample
 using Optimisers
-using ADTypes, ForwardDiff
 using AdvancedVI
+
+alg = KLMinRepGradDescent(AutoReverseDiff());
+nothing
 ```
 
-We now need to select 1. a variational objective, and 2. a variational family.
-Here, we will use the [`RepGradELBO` objective](@ref repgradelbo), which expects an object implementing the [`LogDensityProblems`](https://github.com/tpapp/LogDensityProblems.jl) interface, and the inverse bijector.
+Now, `KLMinRepGradDescent` requires the variational approximation and the target log-density to have the same support.
+Since `y` follows a log-normal prior, its support is bounded to be the positive half-space ``\mathbb{R}_+``.
+Thus, we will use [Bijectors](https://github.com/TuringLang/Bijectors.jl) to match the support of our target posterior and the variational approximation.
 
 ```@example elboexample
-alg = KLMinRepGradDescent(AutoForwardDiff())
+using Bijectors
+
+function Bijectors.bijector(model::NormalLogNormal)
+    (; μ_x, σ_x, μ_y, Σ_y) = model
+    return Bijectors.Stacked(
+        Bijectors.bijector.([LogNormal(μ_x, σ_x), MvNormal(μ_y, Σ_y)]),
+        [1:1, 2:(1 + length(μ_y))],
+    )
+end
+
+b = Bijectors.bijector(model);
+binv = inverse(b)
+nothing
 ```
 
 For the variational family, we will use the classic mean-field Gaussian family.
@@ -109,7 +119,9 @@ Passing `objective` and the initial variational approximation `q` to `optimize`
 
 ```@example elboexample
 n_max_iter = 10^4
-q_out, info, _ = AdvancedVI.optimize(alg, n_max_iter, model, q0_trans; show_progress=false);
+q_out, info, _ = AdvancedVI.optimize(
+    alg, n_max_iter, model_ad, q0_trans; show_progress=false
+);
 nothing
 ```
 

docs/src/families.md

@@ -135,10 +135,10 @@ using ADTypes
 using AdvancedVI
 using Distributions
 using LinearAlgebra
-using LogDensityProblems
+using LogDensityProblems, LogDensityProblemsAD
 using Optimisers
 using Plots
-using ReverseDiff
+using ForwardDiff, ReverseDiff
 
 struct Target{D}
     dist::D
@@ -163,6 +163,7 @@ D_true     = Diagonal(log.(1 .+ exp.(randn(n_dims))))
 Σsqrt_true = sqrt(Σ_true)
 μ_true     = randn(n_dims)
 model      = Target(MvNormal(μ_true, Σ_true));
+model_ad   = ADgradient(AutoForwardDiff(), model)
 
 d  = LogDensityProblems.dimension(model);
 μ  = zeros(d);
@@ -188,19 +189,19 @@ function callback(; params, averaged_params, restructure, kwargs...)
 end
 
 _, info_fr, _ = AdvancedVI.optimize(
-    alg, max_iter, model, q0_fr;
+    alg, max_iter, model_ad, q0_fr;
     show_progress = false,
     callback      = callback,
 ); 
 
 _, info_mf, _ = AdvancedVI.optimize(
-    alg, max_iter, model, q0_mf;
+    alg, max_iter, model_ad, q0_mf;
     show_progress = false,
     callback      = callback,
 ); 
 
 _, info_lr, _ = AdvancedVI.optimize(
-    alg, max_iter, model, q0_lr;
+    alg, max_iter, model_ad, q0_lr;
     show_progress = false,
     callback      = callback,
 ); 

docs/src/paramspacesgd/repgradelbo.md

@@ -122,12 +122,12 @@ For example, if ``q_{\lambda}`` is a Gaussian with a full-rank covariance, a bac
 using Bijectors
 using FillArrays
 using LinearAlgebra
-using LogDensityProblems
+using LogDensityProblems, LogDensityProblemsAD
 using Plots
 using Random
 
 using Optimisers
-using ADTypes, ForwardDiff
+using ADTypes, ForwardDiff, ReverseDiff
 using AdvancedVI
 
 struct NormalLogNormal{MX,SX,MY,SY}
@@ -150,12 +150,13 @@ function LogDensityProblems.capabilities(::Type{<:NormalLogNormal})
     LogDensityProblems.LogDensityOrder{0}()
 end
 
-n_dims = 10
-μ_x    = 2.0
-σ_x    = 0.3
-μ_y    = Fill(2.0, n_dims)
-σ_y    = Fill(1.0, n_dims)
-model  = NormalLogNormal(μ_x, σ_x, μ_y, Diagonal(σ_y.^2));
+n_dims   = 10
+μ_x      = 2.0
+σ_x      = 0.3
+μ_y      = Fill(2.0, n_dims)
+σ_y      = Fill(1.0, n_dims)
+model    = NormalLogNormal(μ_x, σ_x, μ_y, Diagonal(σ_y.^2));
+model_ad = ADgradient(AutoForwardDiff(), model)
 
 d  = LogDensityProblems.dimension(model);
 μ  = zeros(d);
@@ -185,7 +186,7 @@ binv = inverse(b)
 q0_trans = Bijectors.TransformedDistribution(q0, binv)
 
 cfe = KLMinRepGradDescent(
-    AutoForwardDiff(); entropy=ClosedFormEntropy(), optimizer=Adam(1e-2)
+    AutoReverseDiff(); entropy=ClosedFormEntropy(), optimizer=Adam(1e-2)
 )
 nothing
 ```
@@ -194,7 +195,7 @@ The repgradelbo estimator can instead be created as follows:
 
 ```@example repgradelbo
 stl = KLMinRepGradDescent(
-    AutoForwardDiff(); entropy=StickingTheLandingEntropy(), optimizer=Adam(1e-2)
+    AutoReverseDiff(); entropy=StickingTheLandingEntropy(), optimizer=Adam(1e-2)
 )
 nothing
 ```
@@ -212,7 +213,7 @@ end
 _, info_cfe, _ = AdvancedVI.optimize(
     cfe,
     max_iter,
-    model,
+    model_ad,
     q0_trans;
     show_progress = false,
     callback      = callback,
@@ -221,7 +222,16 @@ _, info_cfe, _ = AdvancedVI.optimize(
 _, info_stl, _ = AdvancedVI.optimize(
     stl,
     max_iter,
-    model,
+    model_ad,
+    q0_trans;
+    show_progress = false,
+    callback      = callback,
+); 
+
+_, info_stl, _ = AdvancedVI.optimize(
+    stl,
+    max_iter,
+    model_ad,
     q0_trans;
     show_progress = false,
     callback      = callback,
@@ -302,9 +312,9 @@ nothing
 
 ```@setup repgradelbo
 _, info_qmc, _ = AdvancedVI.optimize(
-    KLMinRepGradDescent(AutoForwardDiff(); n_samples=n_montecarlo, optimizer=Adam(1e-2)),
+    KLMinRepGradDescent(AutoReverseDiff(); n_samples=n_montecarlo, optimizer=Adam(1e-2)),
     max_iter,
-    model,
+    model_ad,
     q0_trans;
     show_progress = false,
     callback      = callback,

ext/AdvancedVIEnzymeExt.jl

@@ -0,0 +1,13 @@
+module AdvancedVIEnzymeExt
+
+using AdvancedVI
+using LogDensityProblems
+using Enzyme
+
+Enzyme.@import_rrule(
+    typeof(LogDensityProblems.logdensity),
+    AdvancedVI.MixedADLogDensityProblem,
+    AbstractVector
+)
+
+end

ext/AdvancedVIMooncakeExt.jl

@@ -0,0 +1,14 @@
+module AdvancedVIMooncakeExt
+
+using AdvancedVI
+using Base: IEEEFloat
+using LogDensityProblems
+using Mooncake
+
+Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{
+    typeof(LogDensityProblems.logdensity),
+    AdvancedVI.MixedADLogDensityProblem,
+    Array{<:IEEEFloat,1},
+}
+
+end

ext/AdvancedVIReverseDiffExt.jl

@@ -0,0 +1,11 @@
+module AdvancedVIReverseDiffExt
+
+using AdvancedVI
+using LogDensityProblems
+using ReverseDiff
+
+ReverseDiff.@grad_from_chainrules LogDensityProblems.logdensity(
+    prob::AdvancedVI.MixedADLogDensityProblem, x::ReverseDiff.TrackedArray
+)
+
+end

src/AdvancedVI.jl

@@ -18,6 +18,7 @@ using LogDensityProblems
 using ADTypes
 using DiffResults
 using DifferentiationInterface
+using ChainRulesCore
 
 using FillArrays
 
@@ -95,6 +96,8 @@ This is an indirection for handling the type stability of `restructure`, as some
 """
 restructure_ad_forward(::ADTypes.AbstractADType, restructure, params) = restructure(params)
 
+include("mixedad_logdensity.jl")
+
 # Variational Families
 export MvLocationScale, MeanFieldGaussian, FullRankGaussian
 

src/algorithms/paramspacesgd/abstractobjective.jl

@@ -59,7 +59,7 @@ function estimate_objective end
 export estimate_objective
 
 """
-    estimate_gradient!(rng, obj, adtype, out, prob, params, restructure, obj_state)
+    estimate_gradient!(rng, obj, adtype, out, params, restructure, obj_state)
 
 Estimate (possibly stochastic) gradients of the variational objective `obj` targeting `prob` with respect to the variational parameters `λ`
 
@@ -68,7 +68,6 @@ Estimate (possibly stochastic) gradients of the variational objective `obj` targ
 - `obj::AbstractVariationalObjective`: Variational objective.
 - `adtype::ADTypes.AbstractADType`: Automatic differentiation backend. 
 - `out::DiffResults.MutableDiffResult`: Buffer containing the objective value and gradient estimates. 
-- `prob`: The target log-joint likelihood implementing the `LogDensityProblem` interface.
 - `params`: Variational parameters to evaluate the gradient on.
 - `restructure`: Function that reconstructs the variational approximation from `params`.
 - `obj_state`: Previous state of the objective.