Merge branch 'breaking' into py/turing-fixes

Author: penelopeysm

.github/workflows/CI.yml

@@ -80,8 +80,3 @@ jobs:
           files: lcov.info
           token: ${{ secrets.CODECOV_TOKEN }}
           fail_ci_if_error: true
-
-      - uses: coverallsapp/github-action@v2
-        with:
-          github-token: ${{ secrets.GITHUB_TOKEN }}
-          path-to-lcov: lcov.info

HISTORY.md

@@ -71,6 +71,22 @@ Their APIs are the same.
 The implementation of the `varname_leaves` and `varname_and_value_leaves` functions have been moved to AbstractPPL.jl.
 Their behaviour is otherwise identical, and they are still accessible from the DynamicPPL module (though still not exported).
 
+## 0.37.5
+
+A minor optimisation for Enzyme AD on DynamicPPL models.
+
+## 0.37.4
+
+An extension for MarginalLogDensities.jl has been added.
+
+Loading DynamicPPL and MarginalLogDensities now provides the `DynamicPPL.marginalize` function to marginalise out variables from a model.
+This is useful for averaging out random effects or nuisance parameters while improving inference on fixed effects/parameters of interest.
+The `marginalize` function returns a `MarginalLogDensities.MarginalLogDensity`, a function-like callable struct that returns the approximate log-density of a subset of the parameters after integrating out the rest of them.
+By default, this uses the Laplace approximation and sparse AD, making the marginalisation computationally very efficient.
+Note that the Laplace approximation relies on the model being differentiable with respect to the marginalised variables, and that their posteriors are unimodal and approximately Gaussian.
+
+Please see [the MarginalLogDensities documentation](https://eloceanografo.github.io/MarginalLogDensities.jl/stable) and the [new Marginalisation section of the DynamicPPL documentation](https://turinglang.org/DynamicPPL.jl/v0.37/api/#Marginalisation) for further information.
+
 ## 0.37.3
 
 Prevents inlining of `DynamicPPL.istrans` with Enzyme, which allows Enzyme to differentiate models where `VarName`s have the same symbol but different types.

Project.toml

@@ -33,6 +33,7 @@ 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]
@@ -41,6 +42,7 @@ DynamicPPLEnzymeCoreExt = ["EnzymeCore"]
 DynamicPPLForwardDiffExt = ["ForwardDiff"]
 DynamicPPLJETExt = ["JET"]
 DynamicPPLMCMCChainsExt = ["MCMCChains"]
+DynamicPPLMarginalLogDensitiesExt = ["MarginalLogDensities"]
 DynamicPPLMooncakeExt = ["Mooncake"]
 
 [compat]
@@ -66,6 +68,7 @@ LinearAlgebra = "1.6"
 LogDensityProblems = "2"
 MCMCChains = "6, 7"
 MacroTools = "0.5.6"
+MarginalLogDensities = "0.4.3"
 Mooncake = "0.4.147"
 OrderedCollections = "1"
 Printf = "1.10"

README.md

@@ -5,7 +5,6 @@
 [![CI](https://github.com/TuringLang/DynamicPPL.jl/workflows/CI/badge.svg?branch=main)](https://github.com/TuringLang/DynamicPPL.jl/actions?query=workflow%3ACI+branch%3Amain)
 [![JuliaPre](https://github.com/TuringLang/DynamicPPL.jl/workflows/JuliaPre/badge.svg?branch=main)](https://github.com/TuringLang/DynamicPPL.jl/actions?query=workflow%3AJuliaPre+branch%3Amain)
 [![IntegrationTest](https://github.com/TuringLang/DynamicPPL.jl/workflows/IntegrationTest/badge.svg?branch=main)](https://github.com/TuringLang/DynamicPPL.jl/actions?query=workflow%3AIntegrationTest+branch%3Amain)
-[![Coverage Status](https://coveralls.io/repos/github/TuringLang/DynamicPPL.jl/badge.svg?branch=main)](https://coveralls.io/github/TuringLang/DynamicPPL.jl?branch=main)
 [![Codecov](https://codecov.io/gh/TuringLang/DynamicPPL.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/TuringLang/DynamicPPL.jl)
 [![Code Style: Blue](https://img.shields.io/badge/code%20style-blue-4495d1.svg)](https://github.com/invenia/BlueStyle)
 [![ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://img.shields.io/badge/ColPrac-Contributor%27s%20Guide-blueviolet)](https://docs.sciml.ai/ColPrac/stable/)

benchmarks/Project.toml

@@ -7,6 +7,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
@@ -23,9 +24,10 @@ ADTypes = "1.14.0"
 BenchmarkTools = "1.6.0"
 Distributions = "0.25.117"
 DynamicPPL = "0.38"
+Enzyme = "0.13"
 ForwardDiff = "0.10.38, 1"
 LogDensityProblems = "2.1.2"
 Mooncake = "0.4"
 PrettyTables = "3"
 ReverseDiff = "1.15.3"
-StableRNGs = "1"
+StableRNGs = "1"

benchmarks/benchmarks.jl

@@ -43,6 +43,7 @@ chosen_combinations = [
     ("Smorgasbord", smorgasbord_instance, :simple_dict, :forwarddiff, true),
     ("Smorgasbord", smorgasbord_instance, :typed, :reversediff, true),
     ("Smorgasbord", smorgasbord_instance, :typed, :mooncake, true),
+    ("Smorgasbord", smorgasbord_instance, :typed, :enzyme, true),
     ("Loop univariate 1k", loop_univariate1k, :typed, :mooncake, true),
     ("Multivariate 1k", multivariate1k, :typed, :mooncake, true),
     ("Loop univariate 10k", loop_univariate10k, :typed, :mooncake, true),

benchmarks/src/DynamicPPLBenchmarks.jl

@@ -13,7 +13,7 @@ using StableRNGs: StableRNG
 
 include("./Models.jl")
 using .Models: Models
-
+using Enzyme: Enzyme
 export Models, make_suite, model_dimension
 
 """
@@ -37,6 +37,10 @@ const SYMBOL_TO_BACKEND = Dict(
     :reversediff => ADTypes.AutoReverseDiff(; compile=false),
     :reversediff_compiled => ADTypes.AutoReverseDiff(; compile=true),
     :mooncake => ADTypes.AutoMooncake(; config=nothing),
+    :enzyme => ADTypes.AutoEnzyme(;
+        mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
+        function_annotation=Enzyme.Const,
+    ),
 )
 
 to_backend(x) = error("Unknown backend: $x")

docs/Project.toml

@@ -10,6 +10,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+MarginalLogDensities = "f0c3360a-fb8d-11e9-1194-5521fd7ee392"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 
 [compat]
@@ -23,5 +24,6 @@ FillArrays = "0.13, 1"
 ForwardDiff = "0.10, 1"
 JET = "0.9, 0.10"
 LogDensityProblems = "2"
+MarginalLogDensities = "0.4"
 MCMCChains = "5, 6, 7"
 StableRNGs = "1"

docs/make.jl

@@ -11,6 +11,10 @@ using Distributions
 using DocumenterMermaid
 # load MCMCChains package extension to make `predict` available
 using MCMCChains
+using MarginalLogDensities: MarginalLogDensities
+
+# Need this to document a method which uses a type inside the extension...
+DPPLMLDExt = Base.get_extension(DynamicPPL, :DynamicPPLMarginalLogDensitiesExt)
 
 # Doctest setup
 DocMeta.setdocmeta!(
@@ -24,7 +28,11 @@ makedocs(;
     format=Documenter.HTML(;
         size_threshold=2^10 * 400, mathengine=Documenter.HTMLWriter.MathJax3()
     ),
-    modules=[DynamicPPL, Base.get_extension(DynamicPPL, :DynamicPPLMCMCChainsExt)],
+    modules=[
+        DynamicPPL,
+        Base.get_extension(DynamicPPL, :DynamicPPLMCMCChainsExt),
+        Base.get_extension(DynamicPPL, :DynamicPPLMarginalLogDensitiesExt),
+    ],
     pages=[
         "Home" => "index.md", "API" => "api.md", "Internals" => ["internals/varinfo.md"]
     ],

docs/src/api.md

@@ -136,6 +136,22 @@ When using `predict` with `MCMCChains.Chains`, you can control which variables a
   - `include_all=false` (default): Include only newly predicted variables
   - `include_all=true`: Include both parameters from the original chain and predicted variables
 
+## Marginalisation
+
+DynamicPPL provides the `marginalize` function to marginalise out variables from a model.
+This requires `MarginalLogDensities.jl` to be loaded in your environment.
+
+```@docs
+marginalize
+```
+
+A `MarginalLogDensity` object acts as a function which maps non-marginalised parameter values to a marginal log-probability.
+To retrieve a VarInfo object from it, you can use:
+
+```@docs
+VarInfo(::MarginalLogDensities.MarginalLogDensity{<:DPPLMLDExt.LogDensityFunctionWrapper}, ::Union{AbstractVector,Nothing})
+```
+
 ## Models within models
 
 One can include models and call another model inside the model function with `left ~ to_submodel(model)`.