(some) Compatibility with DynamicPPL 0.39

Author: penelopeysm

Project.toml

@@ -92,4 +92,4 @@ DynamicHMC = "bbc10e6e-7c05-544b-b16e-64fede858acb"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 
 [sources]
-DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "py/ldf"}
+DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "py/not-experimental"}

src/mcmc/Inference.jl

@@ -6,7 +6,6 @@ using DynamicPPL:
     Metadata,
     VarInfo,
     LogDensityFunction,
-    SimpleVarInfo,
     AbstractVarInfo,
     # TODO(mhauru) all_varnames_grouped_by_symbol isn't exported by DPPL, because it is only
     # implemented for NTVarInfo. It is used by mh.jl. Either refactor mh.jl to not use it
@@ -92,9 +91,6 @@ function DynamicPPL.unflatten(vi::DynamicPPL.NTVarInfo, θ::NamedTuple)
     set_namedtuple!(deepcopy(vi), θ)
     return vi
 end
-function DynamicPPL.unflatten(vi::SimpleVarInfo, θ::NamedTuple)
-    return SimpleVarInfo(θ, vi.logp, vi.transformation)
-end
 
 """
     mh_accept(logp_current::Real, logp_proposal::Real, log_proposal_ratio::Real)

src/mcmc/emcee.jl

@@ -86,11 +86,13 @@ function AbstractMCMC.step(
     densitymodel = AMH.DensityModel(Base.Fix1(LogDensityProblems.logdensity, state.ldf))
 
     # Compute the next states.
-    t, states = AbstractMCMC.step(rng, densitymodel, spl.ensemble, state.states)
+    _, states = AbstractMCMC.step(rng, densitymodel, spl.ensemble, state.states)
 
     # Compute the next transition and state.
     transition = map(states) do _state
-        return DynamicPPL.ParamsWithStats(_state.params, state.ldf, t)
+        return DynamicPPL.ParamsWithStats(
+            _state.params, state.ldf, AbstractMCMC.getstats(_state)
+        )
     end
     newstate = EmceeState(state.ldf, states)
 

src/mcmc/ess.jl

@@ -31,7 +31,7 @@ function Turing.Inference.initialstep(
         EllipticalSliceSampling.isgaussian(typeof(dist)) ||
             error("ESS only supports Gaussian prior distributions")
     end
-    return Transition(model, vi, nothing), vi
+    return DynamicPPL.ParamsWithStats(vi, model), vi
 end
 
 function AbstractMCMC.step(
@@ -56,7 +56,7 @@ function AbstractMCMC.step(
     vi = DynamicPPL.unflatten(vi, sample)
     vi = DynamicPPL.setloglikelihood!!(vi, state.loglikelihood)
 
-    return Transition(model, vi, nothing), vi
+    return DynamicPPL.ParamsWithStats(vi, model), vi
 end
 
 # Prior distribution of considered random variable

src/mcmc/external_sampler.jl

@@ -122,12 +122,12 @@ function externalsampler(
 end
 
 # TODO(penelopeysm): Can't we clean this up somehow?
-struct TuringState{S,V1,M,V}
+struct TuringState{S,V,L<:DynamicPPL.LogDensityFunction}
     state::S
-    # Note that this varinfo must have the correct parameters set; but logp
-    # does not matter as it will be re-evaluated
-    varinfo::V1
-    ldf::DynamicPPL.LogDensityFunction{M,V}
+    # Note that this varinfo is used only for structure. Its parameters and other info do
+    # not need to be accurate
+    varinfo::V
+    ldf::L
 end
 
 # get_varinfo should return something from which the correct parameters can be
@@ -185,11 +185,10 @@ function AbstractMCMC.step(
     end
 
     new_parameters = AbstractMCMC.getparams(f.model, state_inner)
-    new_vi = DynamicPPL.unflatten(f.varinfo, new_parameters)
     new_stats = AbstractMCMC.getstats(state_inner)
     return (
-        DynamicPPL.ParamsWithStats(new_vi, f.model, new_stats),
-        TuringState(state_inner, new_vi, f),
+        DynamicPPL.ParamsWithStats(new_parameters, f, new_stats),
+        TuringState(state_inner, varinfo, f),
     )
 end
 
@@ -209,10 +208,9 @@ function AbstractMCMC.step(
     )
 
     new_parameters = AbstractMCMC.getparams(f.model, state_inner)
-    new_vi = DynamicPPL.unflatten(f.varinfo, new_parameters)
     new_stats = AbstractMCMC.getstats(state_inner)
     return (
-        DynamicPPL.ParamsWithStats(new_vi, f.model, new_stats),
-        TuringState(state_inner, new_vi, f),
+        DynamicPPL.ParamsWithStats(new_parameters, f, new_stats),
+        TuringState(state_inner, state.varinfo, f),
     )
 end

src/mcmc/gibbs.jl

@@ -488,18 +488,12 @@ function setparams_varinfo!!(
 end
 
 function setparams_varinfo!!(
-    model::DynamicPPL.Model,
-    sampler::ExternalSampler,
-    state::TuringState,
-    params::AbstractVarInfo,
+    ::DynamicPPL.Model, ::ExternalSampler, state::TuringState, params::AbstractVarInfo
 )
-    logdensity = DynamicPPL.LogDensityFunction(
-        model, DynamicPPL.getlogjoint_internal, state.ldf.varinfo; adtype=sampler.adtype
-    )
     new_inner_state = AbstractMCMC.setparams!!(
-        AbstractMCMC.LogDensityModel(logdensity), state.state, params[:]
+        AbstractMCMC.LogDensityModel(state.ldf), state.state, params[:]
     )
-    return TuringState(new_inner_state, params, logdensity)
+    return TuringState(new_inner_state, params, state.ldf)
 end
 
 function setparams_varinfo!!(
@@ -513,11 +507,11 @@ function setparams_varinfo!!(
     z = state.z
     resize!(z.θ, length(θ_new))
     z.θ .= θ_new
-    return HMCState(params, state.i, state.kernel, hamiltonian, z, state.adaptor)
+    return HMCState(params, state.i, state.kernel, hamiltonian, z, state.adaptor, state.ldf)
 end
 
 function setparams_varinfo!!(
-    model::DynamicPPL.Model, sampler::PG, state::PGState, params::AbstractVarInfo
+    ::DynamicPPL.Model, ::PG, state::PGState, params::AbstractVarInfo
 )
     return PGState(params, state.rng)
 end

src/mcmc/mh.jl

@@ -179,33 +179,33 @@ get_varinfo(s::MHState) = s.varinfo
 #####################
 
 """
-    set_namedtuple!(vi::VarInfo, nt::NamedTuple)
+    OldLogDensityFunction
 
-Places the values of a `NamedTuple` into the relevant places of a `VarInfo`.
+This is a clone of pre-0.39 DynamicPPL.LogDensityFunction. It is needed for MH because MH
+doesn't actually obey the LogDensityProblems.jl interface: it evaluates
+'LogDensityFunctions' with a NamedTuple(!!)
+
+This means that we can't _really_ use DynamicPPL's LogDensityFunction, since that only
+promises to obey the interface of being called with a vector.
+
+In particular, because `set_namedtuple!` acts on a VarInfo, we need to store the VarInfo
+inside this struct (which DynamicPPL's LogDensityFunction no longer does).
+
+This SHOULD really be refactored to remove this requirement.
 """
-function set_namedtuple!(vi::DynamicPPL.VarInfoOrThreadSafeVarInfo, nt::NamedTuple)
-    for (n, vals) in pairs(nt)
-        vns = vi.metadata[n].vns
-        if vals isa AbstractVector
-            vals = unvectorize(vals)
-        end
-        if length(vns) == 1
-            # Only one variable, assign the values to it
-            DynamicPPL.setindex!(vi, vals, vns[1])
-        else
-            # Spread the values across the variables
-            length(vns) == length(vals) || error("Unequal number of variables and values")
-            for (vn, val) in zip(vns, vals)
-                DynamicPPL.setindex!(vi, val, vn)
-            end
-        end
-    end
+struct OldLogDensityFunction{M<:DynamicPPL.Model,V<:DynamicPPL.AbstractVarInfo}
+    model::M
+    varinfo::V
+end
+function (f::OldLogDensityFunction)(x::AbstractVector)
+    vi = DynamicPPL.unflatten(f.varinfo, x)
+    _, vi = DynamicPPL.evaluate!!(f.model, vi)
+    return DynamicPPL.getlogjoint_internal(vi)
 end
-
 # NOTE(penelopeysm): MH does not conform to the usual LogDensityProblems
 # interface in that it gets evaluated with a NamedTuple. Hence we need this
 # method just to deal with MH.
-function LogDensityProblems.logdensity(f::LogDensityFunction, x::NamedTuple)
+function (f::OldLogDensityFunction)(x::NamedTuple)
     vi = deepcopy(f.varinfo)
     # Note that the NamedTuple `x` does NOT conform to the structure required for
     # `InitFromParams`. In particular, for models that look like this:
@@ -223,8 +223,31 @@ function LogDensityProblems.logdensity(f::LogDensityFunction, x::NamedTuple)
     set_namedtuple!(vi, x)
     # Update log probability.
     _, vi_new = DynamicPPL.evaluate!!(f.model, vi)
-    lj = f.getlogdensity(vi_new)
-    return lj
+    return DynamicPPL.getlogjoint_internal(vi_new)
+end
+
+"""
+    set_namedtuple!(vi::VarInfo, nt::NamedTuple)
+
+Places the values of a `NamedTuple` into the relevant places of a `VarInfo`.
+"""
+function set_namedtuple!(vi::DynamicPPL.VarInfoOrThreadSafeVarInfo, nt::NamedTuple)
+    for (n, vals) in pairs(nt)
+        vns = vi.metadata[n].vns
+        if vals isa AbstractVector
+            vals = unvectorize(vals)
+        end
+        if length(vns) == 1
+            # Only one variable, assign the values to it
+            DynamicPPL.setindex!(vi, vals, vns[1])
+        else
+            # Spread the values across the variables
+            length(vns) == length(vals) || error("Unequal number of variables and values")
+            for (vn, val) in zip(vns, vals)
+                DynamicPPL.setindex!(vi, val, vn)
+            end
+        end
+    end
 end
 
 # unpack a vector if possible
@@ -335,12 +358,7 @@ function propose!!(rng::AbstractRNG, prev_state::MHState, model::Model, spl::MH,
 
     # Make a new transition.
     model = DynamicPPL.setleafcontext(model, MHContext(rng))
-    densitymodel = AMH.DensityModel(
-        Base.Fix1(
-            LogDensityProblems.logdensity,
-            DynamicPPL.LogDensityFunction(model, DynamicPPL.getlogjoint_internal, vi),
-        ),
-    )
+    densitymodel = AMH.DensityModel(OldLogDensityFunction(model, vi))
     trans, _ = AbstractMCMC.step(rng, densitymodel, mh_sampler, prev_trans)
     # trans.params isa NamedTuple
     set_namedtuple!(vi, trans.params)
@@ -370,12 +388,7 @@ function propose!!(
 
     # Make a new transition.
     model = DynamicPPL.setleafcontext(model, MHContext(rng))
-    densitymodel = AMH.DensityModel(
-        Base.Fix1(
-            LogDensityProblems.logdensity,
-            DynamicPPL.LogDensityFunction(model, DynamicPPL.getlogjoint_internal, vi),
-        ),
-    )
+    densitymodel = AMH.DensityModel(OldLogDensityFunction(model, vi))
     trans, _ = AbstractMCMC.step(rng, densitymodel, mh_sampler, prev_trans)
     # trans.params isa AbstractVector
     vi = DynamicPPL.unflatten(vi, trans.params)

src/mcmc/prior.jl

@@ -17,6 +17,6 @@ function AbstractMCMC.step(
         DynamicPPL.LogPriorAccumulator(),
         DynamicPPL.LogLikelihoodAccumulator(),
     ))
-    _, vi = DynamicPPL.fast_evaluate!!(rng, model, InitFromPrior(), accs)
+    _, vi = DynamicPPL.fast_evaluate!!(rng, model, DynamicPPL.InitFromPrior(), accs)
     return DynamicPPL.ParamsWithStats(vi), nothing
 end

test/Project.toml

@@ -79,4 +79,4 @@ TimerOutputs = "0.5"
 julia = "1.10"
 
 [sources]
-DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "py/ldf"}
+DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "py/not-experimental"}

test/mcmc/Inference.jl

@@ -174,19 +174,6 @@ using Turing
             @test mean(chains, :m) ≈ 0 atol = 0.1
         end
 
-        @testset "Vector chain_type" begin
-            chains = sample(
-                StableRNG(seed), gdemo_d(), Prior(), N; chain_type=Vector{NamedTuple}
-            )
-            @test chains isa Vector{<:NamedTuple}
-            @test length(chains) == N
-            @test all(haskey(x, :lp) for x in chains)
-            @test all(haskey(x, :logprior) for x in chains)
-            @test all(haskey(x, :loglikelihood) for x in chains)
-            @test mean(x[:s][1] for x in chains) ≈ 3 atol = 0.11
-            @test mean(x[:m][1] for x in chains) ≈ 0 atol = 0.1
-        end
-
         @testset "accumulators are set correctly" begin
             # Prior() does not reevaluate the model when constructing a
             # `DynamicPPL.ParamsWithStats`, so we had better make sure that it does capture
@@ -638,32 +625,6 @@ using Turing
         )
     end
 
-    @testset "getparams" begin
-        @model function e(x=1.0)
-            return x ~ Normal()
-        end
-        evi = DynamicPPL.VarInfo(e())
-        @test isempty(Turing.Inference.getparams(e(), evi))
-
-        @model function f()
-            return x ~ Normal()
-        end
-        fvi = DynamicPPL.VarInfo(f())
-        fparams = Turing.Inference.getparams(f(), fvi)
-        @test fparams[@varname(x)] == fvi[@varname(x)]
-        @test length(fparams) == 1
-
-        @model function g()
-            x ~ Normal()
-            return y ~ Poisson()
-        end
-        gvi = DynamicPPL.VarInfo(g())
-        gparams = Turing.Inference.getparams(g(), gvi)
-        @test gparams[@varname(x)] == gvi[@varname(x)]
-        @test gparams[@varname(y)] == gvi[@varname(y)]
-        @test length(gparams) == 2
-    end
-
     @testset "empty model" begin
         @model function e(x=1.0)
             return x ~ Normal()