Compatibility with new DPPL version (#1900)

* use unflatten in evaluation of LogDensityFunction

* make AD-related functions able to take AbstractVarInfo

* use unflatten where appropriate

* updated Gibbs

* updated HMC

* move to using BangBang versions of link and invlink

* use link!!

* update tests to be compatible with new DynamicPPL.TestUtils

* updated deps for tests

* fixed tests for ESS

* upper-bound distributions in tests because otherwise depwarns will
cause timeouts

* replace link! with link!!, etc.

* added Setfield and updated optimization stuff

* updated the contrib to use link!!, etc.

* updated AD tests

* updated DPPL versions

* removed usage of deprecated inv

* made some function signatures more restrictive

* Update src/inference/mh.jl

* fixed MH sampler

* increase atol for certain tests to make them pass on MacOS

* reduce atol for a MH test

* disable emcee tests for now

* Update Project.toml

Co-authored-by: David Widmann <[email protected]>

* further reductions in atol to make tests pass

* Update test/runtests.jl

* Update mh.jl

* restrict ForwardDiff for tests to avoid issue with cholesky

* increased number of samples and lowered atol for MH tests

Co-authored-by: David Widmann <[email protected]>
Co-authored-by: Hong Ge <[email protected]>

Author: 3 people

Project.toml

@@ -27,6 +27,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
@@ -45,7 +46,7 @@ DataStructures = "0.18"
 Distributions = "0.23.3, 0.24, 0.25"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.8, 0.9"
-DynamicPPL = "0.20"
+DynamicPPL = "0.21"
 EllipticalSliceSampling = "0.5, 1"
 ForwardDiff = "0.10.3"
 Libtask = "0.6.7, 0.7"
@@ -55,6 +56,7 @@ NamedArrays = "0.9"
 Reexport = "0.2, 1"
 Requires = "0.5, 1.0"
 SciMLBase = "1.37.1"
+Setfield = "0.8"
 SpecialFunctions = "0.7.2, 0.8, 0.9, 0.10, 1, 2"
 StatsBase = "0.32, 0.33"
 StatsFuns = "0.8, 0.9, 1"

src/Turing.jl

@@ -35,7 +35,8 @@ struct LogDensityFunction{V,M,S,C}
 end
 
 function (f::LogDensityFunction)(θ::AbstractVector)
-    return getlogp(last(DynamicPPL.evaluate!!(f.model, VarInfo(f.varinfo, f.sampler, θ), f.sampler, f.context)))
+    vi_new = DynamicPPL.unflatten(f.varinfo, f.sampler, θ)
+    return getlogp(last(DynamicPPL.evaluate!!(f.model, vi_new, f.sampler, f.context)))
 end
 
 # LogDensityProblems interface

src/contrib/inference/dynamichmc.jl

@@ -60,8 +60,8 @@ function DynamicPPL.initialstep(
 )
     # Ensure that initial sample is in unconstrained space.
     if !DynamicPPL.islinked(vi, spl)
-        DynamicPPL.link!(vi, spl)
-        model(rng, vi, spl)
+        vi = DynamicPPL.link!!(vi, spl, model)
+        vi = last(DynamicPPL.evaluate!!(model, vi, DynamicPPL.SamplingContext(rng, spl)))
     end
 
     # Define log-density function.
@@ -79,8 +79,8 @@ function DynamicPPL.initialstep(
     Q, _ = DynamicHMC.mcmc_next_step(steps, results.final_warmup_state.Q)
 
     # Update the variables.
-    vi[spl] = Q.q
-    DynamicPPL.setlogp!!(vi, Q.ℓq)
+    vi = DynamicPPL.setindex!!(vi, Q.q, spl)
+    vi = DynamicPPL.setlogp!!(vi, Q.ℓq)
 
     # Create first sample and state.
     sample = Transition(vi)
@@ -109,8 +109,8 @@ function AbstractMCMC.step(
     Q, _ = DynamicHMC.mcmc_next_step(steps, state.cache)
 
     # Update the variables.
-    vi[spl] = Q.q
-    DynamicPPL.setlogp!!(vi, Q.ℓq)
+    vi = DynamicPPL.setindex!!(vi, Q.q, spl)
+    vi = DynamicPPL.setlogp!!(vi, Q.ℓq)
 
     # Create next sample and state.
     sample = Transition(vi)

src/contrib/inference/sghmc.jl

@@ -56,8 +56,8 @@ function DynamicPPL.initialstep(
 )
     # Transform the samples to unconstrained space and compute the joint log probability.
     if !DynamicPPL.islinked(vi, spl)
-        DynamicPPL.link!(vi, spl)
-        model(rng, vi, spl)
+        vi = DynamicPPL.link!!(vi, spl, model)
+        vi = last(DynamicPPL.evaluate!!(model, vi, DynamicPPL.SamplingContext(rng, spl)))
     end
 
     # Compute initial sample and state.
@@ -90,8 +90,8 @@ function AbstractMCMC.step(
     newv = (1 - α) .* v .+ η .* grad .+ sqrt(2 * η * α) .* randn(rng, eltype(v), length(v))
 
     # Save new variables and recompute log density.
-    vi[spl] = θ
-    model(rng, vi, spl)
+    vi = DynamicPPL.setindex!!(vi, θ, spl)
+    vi = last(DynamicPPL.evaluate!!(model, vi, DynamicPPL.SamplingContext(rng, spl)))
 
     # Compute next sample and state.
     sample = Transition(vi)
@@ -209,8 +209,8 @@ function DynamicPPL.initialstep(
 )
     # Transform the samples to unconstrained space and compute the joint log probability.
     if !DynamicPPL.islinked(vi, spl)
-        DynamicPPL.link!(vi, spl)
-        model(rng, vi, spl)
+        vi = DynamicPPL.link!!(vi, spl, model)
+        vi = last(DynamicPPL.evaluate!!(model, vi, DynamicPPL.SamplingContext(rng, spl)))
     end
 
     # Create first sample and state.
@@ -238,8 +238,8 @@ function AbstractMCMC.step(
     θ .+= (stepsize / 2) .* grad .+ sqrt(stepsize) .* randn(rng, eltype(θ), length(θ))
 
     # Save new variables and recompute log density.
-    vi[spl] = θ
-    model(rng, vi, spl)
+    vi = DynamicPPL.setindex!!(vi, θ, spl)
+    vi = last(DynamicPPL.evaluate!!(model, vi, DynamicPPL.SamplingContext(rng, spl)))
 
     # Compute next sample and state.
     sample = SGLDTransition(vi, stepsize)

src/inference/emcee.jl

@@ -43,7 +43,7 @@ function AbstractMCMC.step(
             ArgumentError("initial parameters have to be specified for each walker")
         )
         vis = map(vis, init_params) do vi, init
-            vi = DynamicPPL.initialize_parameters!!(vi, init, spl)
+            vi = DynamicPPL.initialize_parameters!!(vi, init, spl, model)
 
             # Update log joint probability.
             last(DynamicPPL.evaluate!!(model, rng, vi, SampleFromPrior()))
@@ -57,7 +57,7 @@ function AbstractMCMC.step(
     state = EmceeState(
         vis[1],
         map(vis) do vi
-            DynamicPPL.link!(vi, spl)
+            vi = DynamicPPL.link!!(vi, spl, model)
             AMH.Transition(vi[spl], getlogp(vi))
         end
     )
@@ -82,9 +82,9 @@ function AbstractMCMC.step(
     # Compute the next transition and state.
     transition = map(states) do _state
         vi = setindex!!(vi, _state.params, spl)
-        DynamicPPL.invlink!(vi, spl)
+        vi = DynamicPPL.invlink!!(vi, spl, model)
         t = Transition(tonamedtuple(vi), _state.lp)
-        DynamicPPL.link!(vi, spl)
+        vi = DynamicPPL.link!!(vi, spl, model)
         return t
     end
     newstate = EmceeState(vi, states)

src/inference/gibbs.jl

@@ -199,7 +199,7 @@ function DynamicPPL.initialstep(
     states = map(samplers) do local_spl
         # Recompute `vi.logp` if needed.
         if local_spl.selector.rerun
-            model(rng, vi, local_spl)
+            vi = last(DynamicPPL.evaluate!!(model, vi, DynamicPPL.SamplingContext(rng, local_spl)))
         end
 
         # Compute initial state.

src/inference/hmc.jl

@@ -150,8 +150,7 @@ function DynamicPPL.initialstep(
     kwargs...
 )
     # Transform the samples to unconstrained space and compute the joint log probability.
-    link!(vi, spl)
-    vi = last(DynamicPPL.evaluate!!(model, rng, vi, spl))
+    vi = link!!(vi, spl, model)
 
     # Extract parameters.
     theta = vi[spl]
@@ -173,8 +172,8 @@ function DynamicPPL.initialstep(
     # and its gradient are finite.
     if init_params === nothing
         while !isfinite(z)
+            # NOTE: This will sample in the unconstrained space.
             vi = last(DynamicPPL.evaluate!!(model, rng, vi, SampleFromUniform()))
-            link!(vi, spl)
             theta = vi[spl]
 
             hamiltonian = AHMC.Hamiltonian(metric, logπ, ∂logπ∂θ)
@@ -210,10 +209,10 @@ function DynamicPPL.initialstep(
 
     # Update `vi` based on acceptance
     if t.stat.is_accept
-        vi = setindex!!(vi, t.z.θ, spl)
+        vi = DynamicPPL.unflatten(vi, spl, t.z.θ)
         vi = setlogp!!(vi, t.stat.log_density)
     else
-        vi = setindex!!(vi, theta, spl)
+        vi = DynamicPPL.unflatten(vi, spl, theta)
         vi = setlogp!!(vi, log_density_old)
     end
 
@@ -252,7 +251,7 @@ function AbstractMCMC.step(
     # Update variables
     vi = state.vi
     if t.stat.is_accept
-        vi = setindex!!(vi, t.z.θ, spl)
+        vi = DynamicPPL.unflatten(vi, spl, t.z.θ)
         vi = setlogp!!(vi, t.stat.log_density)
     end
 
@@ -532,8 +531,9 @@ function HMCState(
     kwargs...
 )
     # Link everything if needed.
-    if !islinked(vi, spl)
-        link!(vi, spl)
+    waslinked = islinked(vi, spl)
+    if !waslinked
+        vi = link!!(vi, spl, model)
     end
 
     # Get the initial log pdf and gradient functions.
@@ -562,8 +562,10 @@ function HMCState(
     # Generate a phasepoint. Replaced during sample_init!
     h, t = AHMC.sample_init(rng, h, θ_init) # this also ensure AHMC has the same dim as θ.
 
-    # Unlink everything.
-    invlink!(vi, spl)
+    # Unlink everything, if it was indeed linked before.
+    if waslinked
+        vi = invlink!!(vi, spl, model)
+    end
 
     return HMCState(vi, 0, 0, kernel.τ, h, AHMCAdaptor(spl.alg, metric; ϵ=ϵ), t.z)
 end

src/inference/mh.jl

@@ -197,7 +197,11 @@ end
 
 Places the values of a `NamedTuple` into the relevant places of a `VarInfo`.
 """
-function set_namedtuple!(vi::VarInfo, nt::NamedTuple)
+function set_namedtuple!(vi::DynamicPPL.VarInfoOrThreadSafeVarInfo, nt::NamedTuple)
+    # TODO: Replace this with something like
+    # for vn in keys(vi)
+    #     vi = DynamicPPL.setindex!!(vi, get(nt, vn))
+    # end
     for (n, vals) in pairs(nt)
         vns = vi.metadata[n].vns
         nvns = length(vns)
@@ -245,6 +249,7 @@ This variant uses the  `set_namedtuple!` function to update the `VarInfo`.
 const MHLogDensityFunction{M<:Model,S<:Sampler{<:MH},V<:AbstractVarInfo} = Turing.LogDensityFunction{V,M,S,DynamicPPL.DefaultContext}
 
 function (f::MHLogDensityFunction)(x::NamedTuple)
+    # TODO: Make this work with immutable `f.varinfo` too.
     sampler = f.sampler
     vi = f.varinfo
 
@@ -286,14 +291,14 @@ function reconstruct(
 end
 
 """
-    dist_val_tuple(spl::Sampler{<:MH}, vi::AbstractVarInfo)
+    dist_val_tuple(spl::Sampler{<:MH}, vi::VarInfo)
 
 Return two `NamedTuples`.
 
 The first `NamedTuple` has symbols as keys and distributions as values.
 The second `NamedTuple` has model symbols as keys and their stored values as values.
 """
-function dist_val_tuple(spl::Sampler{<:MH}, vi::AbstractVarInfo)
+function dist_val_tuple(spl::Sampler{<:MH}, vi::DynamicPPL.VarInfoOrThreadSafeVarInfo)
     vns = _getvns(vi, spl)
     dt = _dist_tuple(spl.alg.proposals, vi, vns)
     vt = _val_tuple(vi, vns)
@@ -349,15 +354,12 @@ function should_link(
     return true
 end
 
-function maybe_link!(varinfo, sampler, proposal)
-    if should_link(varinfo, sampler, proposal)
-        link!(varinfo, sampler)
-    end
-    return nothing
+function maybe_link!!(varinfo, sampler, proposal, model)
+    return should_link(varinfo, sampler, proposal) ? link!!(varinfo, sampler, model) : varinfo
 end
 
 # Make a proposal if we don't have a covariance proposal matrix (the default).
-function propose!(
+function propose!!(
     rng::AbstractRNG,
     vi::AbstractVarInfo,
     model::Model,
@@ -378,13 +380,11 @@ function propose!(
     # TODO: Make this compatible with immutable `VarInfo`.
     # Update the values in the VarInfo.
     set_namedtuple!(vi, trans.params)
-    setlogp!!(vi, trans.lp)
-
-    return vi
+    return setlogp!!(vi, trans.lp)
 end
 
 # Make a proposal if we DO have a covariance proposal matrix.
-function propose!(
+function propose!!(
     rng::AbstractRNG,
     vi::AbstractVarInfo,
     model::Model,
@@ -403,12 +403,7 @@ function propose!(
     densitymodel = AMH.DensityModel(Turing.LogDensityFunction(vi, model, spl, DynamicPPL.DefaultContext()))
     trans, _ = AbstractMCMC.step(rng, densitymodel, mh_sampler, prev_trans)
 
-    # TODO: Make this compatible with immutable `VarInfo`.
-    # Update the values in the VarInfo.
-    setindex!!(vi, trans.params, spl)
-    setlogp!!(vi, trans.lp)
-
-    return vi
+    return setlogp!!(DynamicPPL.unflatten(vi, spl, trans.params), trans.lp)
 end
 
 function DynamicPPL.initialstep(
@@ -420,7 +415,7 @@ function DynamicPPL.initialstep(
 )
     # If we're doing random walk with a covariance matrix,
     # just link everything before sampling.
-    maybe_link!(vi, spl, spl.alg.proposals)
+    vi = maybe_link!!(vi, spl, spl.alg.proposals, model)
 
     return Transition(vi), vi
 end
@@ -435,7 +430,7 @@ function AbstractMCMC.step(
     # Cases:
     # 1. A covariance proposal matrix
     # 2. A bunch of NamedTuples that specify the proposal space
-    propose!(rng, vi, model, spl, spl.alg.proposals)
+    vi = propose!!(rng, vi, model, spl, spl.alg.proposals)
 
     return Transition(vi), vi
 end