Progress towards compat with DPPL v0.35

Author: mhauru

Project.toml

@@ -63,7 +63,7 @@ Distributions = "0.23.3, 0.24, 0.25"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.8, 0.9"
 DynamicHMC = "3.4"
-DynamicPPL = "0.34.1"
+DynamicPPL = "0.35"
 EllipticalSliceSampling = "0.5, 1, 2"
 ForwardDiff = "0.10.3"
 Libtask = "0.8.8"

ext/TuringDynamicHMCExt.jl

@@ -25,22 +25,15 @@ To use it, make sure you have DynamicHMC package (version >= 2) loaded:
 using DynamicHMC
 ```
 """
-struct DynamicNUTS{AD,space,T<:DynamicHMC.NUTS} <: Turing.Inference.Hamiltonian
+struct DynamicNUTS{AD,T<:DynamicHMC.NUTS} <: Turing.Inference.Hamiltonian
     sampler::T
     adtype::AD
 end
 
-function DynamicNUTS(
-    spl::DynamicHMC.NUTS=DynamicHMC.NUTS(),
-    space::Tuple=();
-    adtype::ADTypes.AbstractADType=Turing.DEFAULT_ADTYPE,
-)
-    return DynamicNUTS{typeof(adtype),space,typeof(spl)}(spl, adtype)
-end
+DynamicNUTS() = DynamicNUTS(DynamicHMC.NUTS())
+DynamicNUTS(spl) = DynamicNUTS(spl, Turing.DEFAULT_ADTYPE)
 Turing.externalsampler(spl::DynamicHMC.NUTS) = DynamicNUTS(spl)
 
-DynamicPPL.getspace(::DynamicNUTS{<:Any,space}) where {space} = space
-
 """
     DynamicNUTSState
 
@@ -70,8 +63,8 @@ function DynamicPPL.initialstep(
     kwargs...,
 )
     # Ensure that initial sample is in unconstrained space.
-    if !DynamicPPL.islinked(vi, spl)
-        vi = DynamicPPL.link!!(vi, spl, model)
+    if !DynamicPPL.islinked(vi)
+        vi = DynamicPPL.link!!(vi, model)
         vi = last(DynamicPPL.evaluate!!(model, vi, DynamicPPL.SamplingContext(rng, spl)))
     end
 
@@ -82,13 +75,13 @@ function DynamicPPL.initialstep(
 
     # Perform initial step.
     results = DynamicHMC.mcmc_keep_warmup(
-        rng, ℓ, 0; initialization=(q=vi[spl],), reporter=DynamicHMC.NoProgressReport()
+        rng, ℓ, 0; initialization=(q=vi[:],), reporter=DynamicHMC.NoProgressReport()
     )
     steps = DynamicHMC.mcmc_steps(results.sampling_logdensity, results.final_warmup_state)
     Q, _ = DynamicHMC.mcmc_next_step(steps, results.final_warmup_state.Q)
 
     # Update the variables.
-    vi = DynamicPPL.setindex!!(vi, Q.q, spl)
+    vi = DynamicPPL.unflatten(vi, Q.q)
     vi = DynamicPPL.setlogp!!(vi, Q.ℓq)
 
     # Create first sample and state.
@@ -112,7 +105,7 @@ function AbstractMCMC.step(
     Q, _ = DynamicHMC.mcmc_next_step(steps, state.cache)
 
     # Update the variables.
-    vi = DynamicPPL.setindex!!(vi, Q.q, spl)
+    vi = DynamicPPL.unflatten(vi, Q.q)
     vi = DynamicPPL.setlogp!!(vi, Q.ℓq)
 
     # Create next sample and state.

src/Turing.jl

@@ -9,7 +9,7 @@ using Compat: pkgversion
 
 using AdvancedVI: AdvancedVI
 using DynamicPPL: DynamicPPL, LogDensityFunction
-import DynamicPPL: getspace, NoDist, NamedDist
+import DynamicPPL: NoDist, NamedDist
 using LogDensityProblems: LogDensityProblems
 using NamedArrays: NamedArrays
 using Accessors: Accessors

src/essential/container.jl

@@ -39,7 +39,7 @@ function AdvancedPS.advance!(
 end
 
 function AdvancedPS.delete_retained!(trace::TracedModel)
-    DynamicPPL.set_retained_vns_del_by_spl!(trace.varinfo, trace.sampler)
+    DynamicPPL.set_retained_vns_del!(trace.varinfo)
     return trace
 end
 

src/mcmc/Inference.jl

@@ -5,14 +5,16 @@ using DynamicPPL:
     Metadata,
     VarInfo,
     TypedVarInfo,
+    # TODO(mhauru) all_varnames_grouped_by_symbol isn't exported by DPPL. Either export it
+    # or use something else.
+    all_varnames_grouped_by_symbol,
     islinked,
     setindex!!,
     push!!,
     setlogp!!,
     getlogp,
     VarName,
     getsym,
-    _getvns,
     getdist,
     Model,
     Sampler,
@@ -22,9 +24,7 @@ using DynamicPPL:
     PriorContext,
     LikelihoodContext,
     set_flag!,
-    unset_flag!,
-    getspace,
-    inspace
+    unset_flag!
 using Distributions, Libtask, Bijectors
 using DistributionsAD: VectorOfMultivariate
 using LinearAlgebra
@@ -75,9 +75,7 @@ export InferenceAlgorithm,
     RepeatSampler,
     Prior,
     assume,
-    dot_assume,
     observe,
-    dot_observe,
     predict,
     externalsampler
 
@@ -299,7 +297,7 @@ function AbstractMCMC.sample(
     kwargs...,
 )
     check_model && _check_model(model, alg)
-    return AbstractMCMC.sample(rng, model, Sampler(alg, model), N; kwargs...)
+    return AbstractMCMC.sample(rng, model, Sampler(alg), N; kwargs...)
 end
 
 function AbstractMCMC.sample(
@@ -326,9 +324,7 @@ function AbstractMCMC.sample(
     kwargs...,
 )
     check_model && _check_model(model, alg)
-    return AbstractMCMC.sample(
-        rng, model, Sampler(alg, model), ensemble, N, n_chains; kwargs...
-    )
+    return AbstractMCMC.sample(rng, model, Sampler(alg), ensemble, N, n_chains; kwargs...)
 end
 
 function AbstractMCMC.sample(
@@ -583,11 +579,6 @@ end
 # Utilities  #
 ##############
 
-# TODO(mhauru) Remove this once DynamicPPL has removed all its Selector stuff.
-DynamicPPL.getspace(::InferenceAlgorithm) = ()
-DynamicPPL.getspace(spl::Sampler) = getspace(spl.alg)
-DynamicPPL.inspace(vn::VarName, spl::Sampler) = inspace(vn, getspace(spl.alg))
-
 """
 
     transitions_from_chain(

src/mcmc/emcee.jl

@@ -67,7 +67,7 @@ function AbstractMCMC.step(
     state = EmceeState(
         vis[1],
         map(vis) do vi
-            vi = DynamicPPL.link!!(vi, spl, model)
+            vi = DynamicPPL.link!!(vi, model)
             AMH.Transition(vi[spl], getlogp(vi), false)
         end,
     )
@@ -89,7 +89,7 @@ function AbstractMCMC.step(
 
     # Compute the next transition and state.
     transition = map(states) do _state
-        vi = setindex!!(vi, _state.params, spl)
+        vi = DynamicPPL.unflatten(vi, _state.params)
         t = Transition(getparams(model, vi), _state.lp)
         return t
     end

src/mcmc/ess.jl

@@ -136,17 +136,3 @@ function DynamicPPL.tilde_observe(
 )
     return DynamicPPL.tilde_observe(ctx, SampleFromPrior(), right, left, vi)
 end
-
-function DynamicPPL.dot_tilde_assume(
-    rng::Random.AbstractRNG, ::DefaultContext, ::Sampler{<:ESS}, right, left, vns, vi
-)
-    return DynamicPPL.dot_tilde_assume(
-        rng, LikelihoodContext(), SampleFromPrior(), right, left, vns, vi
-    )
-end
-
-function DynamicPPL.dot_tilde_observe(
-    ctx::DefaultContext, sampler::Sampler{<:ESS}, right, left, vi
-)
-    return DynamicPPL.dot_tilde_observe(ctx, SampleFromPrior(), right, left, vi)
-end

src/mcmc/gibbs.jl

@@ -198,58 +198,6 @@ function DynamicPPL.tilde_assume(
     end
 end
 
-# Like the above tilde_assume methods, but with dot_tilde_assume and broadcasting of logpdf.
-# See comments there for more details.
-function DynamicPPL.dot_tilde_assume(context::GibbsContext, right, left, vns, vi)
-    child_context = DynamicPPL.childcontext(context)
-    return if is_target_varname(context, vns)
-        DynamicPPL.dot_tilde_assume(child_context, right, left, vns, vi)
-    elseif has_conditioned_gibbs(context, vns)
-        value, lp, _ = DynamicPPL.dot_tilde_assume(
-            child_context, right, left, vns, get_global_varinfo(context)
-        )
-        value, lp, vi
-    else
-        value, lp, new_global_vi = DynamicPPL.dot_tilde_assume(
-            child_context,
-            DynamicPPL.SampleFromPrior(),
-            right,
-            left,
-            vns,
-            get_global_varinfo(context),
-        )
-        set_global_varinfo!(context, new_global_vi)
-        value, lp, vi
-    end
-end
-
-# As above but with an RNG.
-function DynamicPPL.dot_tilde_assume(
-    rng::Random.AbstractRNG, context::GibbsContext, sampler, right, left, vns, vi
-)
-    child_context = DynamicPPL.childcontext(context)
-    return if is_target_varname(context, vns)
-        DynamicPPL.dot_tilde_assume(rng, child_context, sampler, right, left, vns, vi)
-    elseif has_conditioned_gibbs(context, vns)
-        value, lp, _ = DynamicPPL.dot_tilde_assume(
-            child_context, right, left, vns, get_global_varinfo(context)
-        )
-        value, lp, vi
-    else
-        value, lp, new_global_vi = DynamicPPL.dot_tilde_assume(
-            rng,
-            child_context,
-            DynamicPPL.SampleFromPrior(),
-            right,
-            left,
-            vns,
-            get_global_varinfo(context),
-        )
-        set_global_varinfo!(context, new_global_vi)
-        value, lp, vi
-    end
-end
-
 """
     make_conditional(model, target_variables, varinfo)
 
@@ -281,16 +229,8 @@ function make_conditional(
     return DynamicPPL.contextualize(model, gibbs_context), gibbs_context_inner
 end
 
-# All samplers are given the same Selector, so that they will sample all variables
-# given to them by the Gibbs sampler. This avoids conflicts between the new and the old way
-# of choosing which sampler to use.
-function set_selector(x::DynamicPPL.Sampler)
-    return DynamicPPL.Sampler(x.alg, DynamicPPL.Selector(0))
-end
-function set_selector(x::RepeatSampler)
-    return RepeatSampler(set_selector(x.sampler), x.num_repeat)
-end
-set_selector(x::InferenceAlgorithm) = DynamicPPL.Sampler(x, DynamicPPL.Selector(0))
+wrap_in_sampler(x::AbstractMCMC.AbstractSampler) = x
+wrap_in_sampler(x::InferenceAlgorithm) = DynamicPPL.Sampler(x)
 
 to_varname_list(x::Union{VarName,Symbol}) = [VarName(x)]
 # Any other value is assumed to be an iterable of VarNames and Symbols.
@@ -343,9 +283,7 @@ struct Gibbs{N,V<:NTuple{N,AbstractVector{<:VarName}},A<:NTuple{N,Any}} <:
             end
         end
 
-        # Ensure that samplers have the same selector, and that varnames are lists of
-        # VarNames.
-        samplers = tuple(map(set_selector, samplers)...)
+        samplers = tuple(map(wrap_in_sampler, samplers)...)
         varnames = tuple(map(to_varname_list, varnames)...)
         return new{length(samplers),typeof(varnames),typeof(samplers)}(varnames, samplers)
     end

src/mcmc/hmc.jl

@@ -148,10 +148,10 @@ function DynamicPPL.initialstep(
     kwargs...,
 )
     # Transform the samples to unconstrained space and compute the joint log probability.
-    vi = DynamicPPL.link(vi_original, spl, model)
+    vi = DynamicPPL.link(vi_original, model)
 
     # Extract parameters.
-    theta = vi[spl]
+    theta = vi[:]
 
     # Create a Hamiltonian.
     metricT = getmetricT(spl.alg)
@@ -189,7 +189,7 @@ function DynamicPPL.initialstep(
 
             # NOTE: This will sample in the unconstrained space.
             vi = last(DynamicPPL.evaluate!!(model, rng, vi, SampleFromUniform()))
-            theta = vi[spl]
+            theta = vi[:]
 
             hamiltonian = AHMC.Hamiltonian(metric, logπ, ∂logπ∂θ)
             z = AHMC.phasepoint(rng, theta, hamiltonian)
@@ -226,10 +226,10 @@ function DynamicPPL.initialstep(
 
     # Update `vi` based on acceptance
     if t.stat.is_accept
-        vi = DynamicPPL.unflatten(vi, spl, t.z.θ)
+        vi = DynamicPPL.unflatten(vi, t.z.θ)
         vi = setlogp!!(vi, t.stat.log_density)
     else
-        vi = DynamicPPL.unflatten(vi, spl, theta)
+        vi = DynamicPPL.unflatten(vi, theta)
         vi = setlogp!!(vi, log_density_old)
     end
 
@@ -274,7 +274,7 @@ function AbstractMCMC.step(
     # Update variables
     vi = state.vi
     if t.stat.is_accept
-        vi = DynamicPPL.unflatten(vi, spl, t.z.θ)
+        vi = DynamicPPL.unflatten(vi, t.z.θ)
         vi = setlogp!!(vi, t.stat.log_density)
     end
 
@@ -493,45 +493,15 @@ end
 #### Compiler interface, i.e. tilde operators.
 ####
 function DynamicPPL.assume(
-    rng, spl::Sampler{<:Hamiltonian}, dist::Distribution, vn::VarName, vi
+    rng, ::Sampler{<:Hamiltonian}, dist::Distribution, vn::VarName, vi
 )
     return DynamicPPL.assume(dist, vn, vi)
 end
 
-function DynamicPPL.dot_assume(
-    rng,
-    spl::Sampler{<:Hamiltonian},
-    dist::MultivariateDistribution,
-    vns::AbstractArray{<:VarName},
-    var::AbstractMatrix,
-    vi,
-)
-    return DynamicPPL.dot_assume(dist, var, vns, vi)
-end
-function DynamicPPL.dot_assume(
-    rng,
-    spl::Sampler{<:Hamiltonian},
-    dists::Union{Distribution,AbstractArray{<:Distribution}},
-    vns::AbstractArray{<:VarName},
-    var::AbstractArray,
-    vi,
-)
-    return DynamicPPL.dot_assume(dists, var, vns, vi)
-end
-
-function DynamicPPL.observe(spl::Sampler{<:Hamiltonian}, d::Distribution, value, vi)
+function DynamicPPL.observe(::Sampler{<:Hamiltonian}, d::Distribution, value, vi)
     return DynamicPPL.observe(d, value, vi)
 end
 
-function DynamicPPL.dot_observe(
-    spl::Sampler{<:Hamiltonian},
-    ds::Union{Distribution,AbstractArray{<:Distribution}},
-    value::AbstractArray,
-    vi,
-)
-    return DynamicPPL.dot_observe(ds, value, vi)
-end
-
 ####
 #### Default HMC stepsize and mass matrix adaptor
 ####