Transfer MarginalLogDensities extension from Turing

See: TuringLang/Turing.jl#2664

Co-authored-by: Sam Urmy <[email protected]>

Author: penelopeysm

Project.toml

@@ -34,13 +34,15 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+MarginalLogDensities = "f0c3360a-fb8d-11e9-1194-5521fd7ee392"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 
 [extensions]
 DynamicPPLChainRulesCoreExt = ["ChainRulesCore"]
 DynamicPPLEnzymeCoreExt = ["EnzymeCore"]
 DynamicPPLForwardDiffExt = ["ForwardDiff"]
 DynamicPPLJETExt = ["JET"]
+DynamicPPLMarginalLogDensitiesExt = ["MarginalLogDensities"]
 DynamicPPLMCMCChainsExt = ["MCMCChains"]
 DynamicPPLMooncakeExt = ["Mooncake"]
 

ext/DynamicPPLMarginalLogDensitiesExt.jl

@@ -0,0 +1,36 @@
+module DynamicPPLMarginalLogDensitiesExt
+
+using DynamicPPL: DynamicPPL, LogDensityProblems, VarName
+using MarginalLogDensities: MarginalLogDensities
+
+_to_varname(n::Symbol) = VarName{n}()
+_to_varname(n::VarName) = n
+
+function DynamicPPL.marginalize(
+    model::DynamicPPL.Model,
+    varnames::AbstractVector{<:Union{Symbol,<:VarName}},
+    getlogprob=DynamicPPL.getlogjoint,
+    method::MarginalLogDensities.AbstractMarginalizer=MarginalLogDensities.LaplaceApprox();
+    kwargs...,
+)
+    # Determine the indices for the variables to marginalise out.
+    varinfo = DynamicPPL.typed_varinfo(model)
+    vns = map(_to_varname, varnames)
+    varindices = reduce(vcat, DynamicPPL.vector_getranges(varinfo, vns))
+    # Construct the marginal log-density model.
+    # Use linked `varinfo` to that we're working in unconstrained space
+    varinfo_linked = DynamicPPL.link(varinfo, model)
+
+    f = DynamicPPL.LogDensityFunction(model, getlogprob, varinfo_linked)
+    mdl = MarginalLogDensities.MarginalLogDensity(
+        (x, _) -> LogDensityProblems.logdensity(f, x),
+        varinfo_linked[:],
+        varindices,
+        (),
+        method;
+        kwargs...,
+    )
+    return mdl
+end
+
+end

src/DynamicPPL.jl

@@ -122,6 +122,7 @@ export AbstractVarInfo,
     fix,
     unfix,
     predict,
+    marginalize,
     prefix,
     returned,
     to_submodel,
@@ -199,10 +200,6 @@ include("test_utils.jl")
 include("experimental.jl")
 include("deprecated.jl")
 
-if !isdefined(Base, :get_extension)
-    using Requires
-end
-
 # Better error message if users forget to load JET
 if isdefined(Base.Experimental, :register_error_hint)
     function __init__()
@@ -247,4 +244,7 @@ end
 # Ref: https://www.stochasticlifestyle.com/improved-forwarddiff-jl-stacktraces-with-package-tags/
 struct DynamicPPLTag end
 
+# Extended in MarginalLogDensitiesExt
+function marginalize end
+
 end # module

test/Project.toml

@@ -18,6 +18,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+MarginalLogDensities = "f0c3360a-fb8d-11e9-1194-5521fd7ee392"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

test/ext/DynamicPPLMarginalLogDensitiesExt.jl

@@ -0,0 +1,27 @@
+module MarginalLogDensitiesExtTests
+
+using DynamicPPL, Distributions, Test
+using MarginalLogDensities
+using ADTypes: AutoForwardDiff
+
+@testset "MarginalLogDensities" begin
+    # Simple test case.
+    @model function demo()
+        x ~ MvNormal(zeros(2), [1, 1])
+        return y ~ Normal(0, 1)
+    end
+    model = demo()
+    # Marginalize out `x`.
+
+    for vn in [@varname(x), :x]
+        for getlogprob in [DynamicPPL.getlogprior, DynamicPPL.getlogjoint]
+            marginalized = marginalize(
+                model, [vn], getlogprob; hess_adtype=AutoForwardDiff()
+            )
+            # Compute the marginal log-density of `y = 0.0`.
+            @test marginalized([0.0]) ≈ logpdf(Normal(0, 1), 0.0) atol = 1e-5
+        end
+    end
+end
+
+end

test/runtests.jl

@@ -80,6 +80,7 @@ include("test_util.jl")
         @testset "extensions" begin
             include("ext/DynamicPPLMCMCChainsExt.jl")
             include("ext/DynamicPPLJETExt.jl")
+            include("ext/DynamicPPLMarginalLogDensitiesExt.jl")
         end
         @testset "ad" begin
             include("ext/DynamicPPLForwardDiffExt.jl")