Add code for AbstractMCMC Callbacks support

Project.toml

@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.42.7"
+version = "0.42.8"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/mcmc/Inference.jl

@@ -115,6 +115,26 @@ function mh_accept(logp_current::Real, logp_proposal::Real, log_proposal_ratio::
     return log(rand()) + logp_current ≤ logp_proposal + log_proposal_ratio
 end
 
+# Directly overload the constructor of `AbstractMCMC.ParamsWithStats` so that we don't
+# hit the default method, which uses `getparams(state)` and `getstats(state)`. For Turing's
+# MCMC samplers, the state might contain results that are in linked space. Using the
+# outputs of the transition here ensures that parameters and logprobs are provided in
+# user space (similar to chains output).
+function AbstractMCMC.ParamsWithStats(
+    model,
+    sampler,
+    transition::DynamicPPL.ParamsWithStats,
+    state;
+    params::Bool=true,
+    stats::Bool=false,
+    extras::Bool=false,
+)
+    p = params ? [string(k) => v for (k, v) in transition.params] : nothing
+    s = stats ? transition.stats : NamedTuple()
+    e = extras ? NamedTuple() : NamedTuple()
+    return AbstractMCMC.ParamsWithStats(p, s, e)
+end
+
 #######################################
 # Concrete algorithm implementations. #
 #######################################

test/Project.toml

@@ -40,7 +40,7 @@ TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [compat]
 ADTypes = "1"
-AbstractMCMC = "5.9"
+AbstractMCMC = "5.13"
 AbstractPPL = "0.11, 0.12, 0.13"
 AdvancedMH = "0.8.9"
 AdvancedPS = "0.7.2"

test/mcmc/callbacks.jl

@@ -0,0 +1,65 @@
+module CallbacksTests
+
+using Test, Turing, AbstractMCMC, Random, Distributions, LinearAlgebra
+
+@model function test_normals()
+    x ~ Normal()
+    return y ~ MvNormal(zeros(3), I)
+end
+
+@testset "AbstractMCMC Callbacks Interface" begin
+    rng = Random.default_rng()
+    model = test_normals()
+
+    samplers = [
+        ("NUTS", NUTS(10, 0.65)),
+        ("HMC", HMC(0.1, 5)),
+        ("MH", MH()),
+        ("ESS", ESS()),
+        ("Gibbs", Gibbs(:x => HMC(0.1, 5), :y => MH())),
+        ("SGHMC", SGHMC(; learning_rate=0.01, momentum_decay=1e-2)),
+        ("PG", PG(10)),
+    ]
+
+    for (name, sampler) in samplers
+        @testset "$name" begin
+            t1, s1 = AbstractMCMC.step(
+                rng, model, sampler; initial_params=Turing.Inference.init_strategy(sampler)
+            )
+
+            # ParamsWithStats returns named params (not θ[i])
+            pws = AbstractMCMC.ParamsWithStats(
+                model, sampler, t1, s1; params=true, stats=true
+            )
+            pairs_dict = Dict(k => v for (k, v) in Base.pairs(pws))
+            # Keys are Symbols since ParamsWithStats stores NamedTuple internally
+            @test haskey(pairs_dict, Symbol("x"))
+            @test haskey(pairs_dict, Symbol("y"))
+            @test pairs_dict[Symbol("y")] isa AbstractVector
+            @test length(pairs_dict[Symbol("y")]) == 3
+
+            # Check stats contain lp
+            @test haskey(pairs_dict, :lp) || haskey(pairs_dict, :logjoint)
+        end
+    end
+
+    # NUTS second step has full AHMC transition metrics
+    @testset "NUTS Transition Metrics" begin
+        sampler = NUTS(10, 0.65)
+        t1, s1 = AbstractMCMC.step(
+            rng, model, sampler; initial_params=Turing.Inference.init_strategy(sampler)
+        )
+        t2, s2 = AbstractMCMC.step(rng, model, sampler, s1)
+
+        pws = AbstractMCMC.ParamsWithStats(model, sampler, t2, s2; params=true, stats=true)
+        pairs_dict = Dict(k => v for (k, v) in Base.pairs(pws))
+
+        # Keys are Symbols from NamedTuple
+        @test haskey(pairs_dict, :tree_depth)
+        @test haskey(pairs_dict, :n_steps)
+        @test haskey(pairs_dict, :acceptance_rate)
+        @test haskey(pairs_dict, :hamiltonian_energy)
+    end
+end
+
+end

test/runtests.jl

@@ -44,6 +44,7 @@ end
 
     @testset "samplers (without AD)" verbose = true begin
         @timeit_include("mcmc/abstractmcmc.jl")
+        @timeit_include("mcmc/callbacks.jl")
         @timeit_include("mcmc/particle_mcmc.jl")
         @timeit_include("mcmc/emcee.jl")
         @timeit_include("mcmc/ess.jl")