Compatibility with DynamicPPL 0.38 + InitContext

HISTORY.md

@@ -1,5 +1,30 @@
 # 0.41.0
 
+## DynamicPPL 0.38
+
+Turing.jl v0.41 brings with it all the underlying changes in DynamicPPL 0.38.
+
+The only user-facing difference is that initial parameters for MCMC sampling must now be specified in a different form.
+You still need to use the `initial_params` keyword argument to `sample`, but the allowed values are different.
+For almost all samplers in Turing.jl (except `Emcee`) this should now be a `DynamicPPL.AbstractInitStrategy`.
+
+TODO LINK TO DPPL DOCS WHEN THIS IS LIVE
+
+There are three kinds of initialisation strategies provided out of the box with Turing.jl (they are exported so you can use these directly with `using Turing`):
+
+  - `InitFromPrior()`: Sample from the prior distribution. This is the default for most samplers in Turing.jl (if you don't specify `initial_params`).
+  - `InitFromUniform(a, b)`: Sample uniformly from `[a, b]` in linked space. This is the default for Hamiltonian samplers. If `a` and `b` are not specified it defaults to `[-2, 2]`, which preserves the behaviour in previous versions (and mimics that of Stan).
+  - `InitFromParams(p)`: Explicitly provide a set of initial parameters. **Note: `p` must be either a `NamedTuple` or a `Dict{<:VarName}`; it can no longer be a `Vector`.** Parameters must be provided in unlinked space, even if the sampler later performs linking.
+
+This change is made because Vectors are semantically ambiguous.
+It is not clear which element of the vector corresponds to which variable in the model, nor is it clear whether the parameters are in linked or unlinked space.
+Previously, both of these would depend on the internal structure of the VarInfo, which is an implementation detail.
+In contrast, the behaviour of `Dict`s and `NamedTuple`s is invariant to the ordering of variables and it is also easier for readers to understand which variable is being set to which value.
+
+If you were previously using `varinfo[:]` to extract a vector of initial parameters, you can now use `Dict(k => varinfo[k] for k in keys(varinfo)` to extract a Dict of initial parameters.
+
+## Initial step in MCMC sampling
+
 HMC and NUTS samplers no longer take an extra single step before starting the chain.
 This means that if you do not discard any samples at the start, the first sample will be the initial parameters (which may be user-provided).
 

Project.toml

@@ -45,7 +45,7 @@ Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 
 [extensions]
 TuringDynamicHMCExt = "DynamicHMC"
-TuringOptimExt = "Optim"
+TuringOptimExt = ["Optim", "AbstractPPL"]
 
 [compat]
 ADTypes = "1.9"
@@ -64,7 +64,7 @@ Distributions = "0.25.77"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.8, 0.9"
 DynamicHMC = "3.4"
-DynamicPPL = "0.37.2"
+DynamicPPL = "0.38"
 EllipticalSliceSampling = "0.5, 1, 2"
 ForwardDiff = "0.10.3, 1"
 Libtask = "0.9.3"
@@ -90,3 +90,6 @@ julia = "1.10.8"
 [extras]
 DynamicHMC = "bbc10e6e-7c05-544b-b16e-64fede858acb"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+
+[sources]
+DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "py/turing-fixes"}

docs/src/api.md

@@ -75,6 +75,16 @@ even though [`Prior()`](@ref) is actually defined in the `Turing.Inference` modu
 | `RepeatSampler`      | [`Turing.Inference.RepeatSampler`](@ref)      | A sampler that runs multiple times on the same variable             |
 | `externalsampler`    | [`Turing.Inference.externalsampler`](@ref)    | Wrap an external sampler for use in Turing                          |
 
+### Initialisation strategies
+
+Turing.jl provides several strategies to initialise parameters for models.
+
+| Exported symbol   | Documentation                           | Description                                                     |
+|:----------------- |:--------------------------------------- |:--------------------------------------------------------------- |
+| `InitFromPrior`   | [`DynamicPPL.InitFromPrior`](@extref)   | Obtain initial parameters from the prior distribution           |
+| `InitFromUniform` | [`DynamicPPL.InitFromUniform`](@extref) | Obtain initial parameters by sampling uniformly in linked space |
+| `InitFromParams`  | [`DynamicPPL.InitFromParams`](@extref)  | Manually specify (possibly a subset of) initial parameters      |
+
 ### Variational inference
 
 See the [docs of AdvancedVI.jl](https://turinglang.org/AdvancedVI.jl/stable/) for detailed usage and the [variational inference tutorial](https://turinglang.org/docs/tutorials/09-variational-inference/) for a basic walkthrough.

ext/TuringDynamicHMCExt.jl

@@ -44,10 +44,6 @@ struct DynamicNUTSState{L,V<:DynamicPPL.AbstractVarInfo,C,M,S}
     stepsize::S
 end
 
-function DynamicPPL.initialsampler(::DynamicPPL.Sampler{<:DynamicNUTS})
-    return DynamicPPL.SampleFromUniform()
-end
-
 function DynamicPPL.initialstep(
     rng::Random.AbstractRNG,
     model::DynamicPPL.Model,

ext/TuringOptimExt.jl

@@ -1,6 +1,7 @@
 module TuringOptimExt
 
 using Turing: Turing
+using AbstractPPL: AbstractPPL
 import Turing: DynamicPPL, NamedArrays, Accessors, Optimisation
 using Optim: Optim
 
@@ -186,7 +187,7 @@ function _optimize(
         f.ldf.model, f.ldf.getlogdensity, vi_optimum; adtype=f.ldf.adtype
     )
     vals_dict = Turing.Inference.getparams(f.ldf.model, vi_optimum)
-    iters = map(DynamicPPL.varname_and_value_leaves, keys(vals_dict), values(vals_dict))
+    iters = map(AbstractPPL.varname_and_value_leaves, keys(vals_dict), values(vals_dict))
     vns_vals_iter = mapreduce(collect, vcat, iters)
     varnames = map(Symbol ∘ first, vns_vals_iter)
     vals = map(last, vns_vals_iter)

src/Turing.jl

@@ -73,7 +73,10 @@ using DynamicPPL:
     conditioned,
     to_submodel,
     LogDensityFunction,
-    @addlogprob!
+    @addlogprob!,
+    InitFromPrior,
+    InitFromUniform,
+    InitFromParams
 using StatsBase: predict
 using OrderedCollections: OrderedDict
 
@@ -148,6 +151,10 @@ export
     fix,
     unfix,
     OrderedDict, # OrderedCollections
+    # Initialisation strategies for models
+    InitFromPrior,
+    InitFromUniform,
+    InitFromParams,
     # Point estimates - Turing.Optimisation
     # The MAP and MLE exports are only needed for the Optim.jl interface.
     maximum_a_posteriori,

src/mcmc/Inference.jl

@@ -24,8 +24,6 @@ using DynamicPPL:
     getdist,
     Model,
     Sampler,
-    SampleFromPrior,
-    SampleFromUniform,
     DefaultContext,
     set_flag!,
     unset_flag!
@@ -59,8 +57,6 @@ export InferenceAlgorithm,
     Hamiltonian,
     StaticHamiltonian,
     AdaptiveHamiltonian,
-    SampleFromUniform,
-    SampleFromPrior,
     MH,
     ESS,
     Emcee,
@@ -262,13 +258,13 @@ function _params_to_array(model::DynamicPPL.Model, ts::Vector)
     dicts = map(ts) do t
         # In general getparams returns a dict of VarName => values. We need to also
         # split it up into constituent elements using
-        # `DynamicPPL.varname_and_value_leaves` because otherwise MCMCChains.jl
+        # `AbstractPPL.varname_and_value_leaves` because otherwise MCMCChains.jl
         # won't understand it.
         vals = getparams(model, t)
         nms_and_vs = if isempty(vals)
             Tuple{VarName,Any}[]
         else
-            iters = map(DynamicPPL.varname_and_value_leaves, keys(vals), values(vals))
+            iters = map(AbstractPPL.varname_and_value_leaves, keys(vals), values(vals))
             mapreduce(collect, vcat, iters)
         end
         nms = map(first, nms_and_vs)
@@ -315,11 +311,10 @@ end
 getlogevidence(transitions, sampler, state) = missing
 
 # Default MCMCChains.Chains constructor.
-# This is type piracy (at least for SampleFromPrior).
 function AbstractMCMC.bundle_samples(
-    ts::Vector{<:Union{Transition,AbstractVarInfo}},
+    ts::Vector{<:Transition},
     model::AbstractModel,
-    spl::Union{Sampler{<:InferenceAlgorithm},SampleFromPrior,RepeatSampler},
+    spl::Union{Sampler{<:InferenceAlgorithm},RepeatSampler},
     state,
     chain_type::Type{MCMCChains.Chains};
     save_state=false,
@@ -378,11 +373,10 @@ function AbstractMCMC.bundle_samples(
     return sort_chain ? sort(chain) : chain
 end
 
-# This is type piracy (for SampleFromPrior).
 function AbstractMCMC.bundle_samples(
-    ts::Vector{<:Union{Transition,AbstractVarInfo}},
+    ts::Vector{<:Transition},
     model::AbstractModel,
-    spl::Union{Sampler{<:InferenceAlgorithm},SampleFromPrior,RepeatSampler},
+    spl::Union{Sampler{<:InferenceAlgorithm},RepeatSampler},
     state,
     chain_type::Type{Vector{NamedTuple}};
     kwargs...,

src/mcmc/abstractmcmc.jl

@@ -1,9 +1,9 @@
 # TODO: Implement additional checks for certain samplers, e.g.
 # HMC not supporting discrete parameters.
 function _check_model(model::DynamicPPL.Model)
-    # TODO(DPPL0.38/penelopeysm): use InitContext
-    spl_model = DynamicPPL.contextualize(model, DynamicPPL.SamplingContext(model.context))
-    return DynamicPPL.check_model(spl_model, VarInfo(); error_on_failure=true)
+    new_context = DynamicPPL.setleafcontext(model.context, DynamicPPL.InitContext())
+    new_model = DynamicPPL.contextualize(model, new_context)
+    return DynamicPPL.check_model(new_model, VarInfo(); error_on_failure=true)
 end
 function _check_model(model::DynamicPPL.Model, alg::InferenceAlgorithm)
     return _check_model(model)

src/mcmc/algorithm.jl

@@ -12,3 +12,7 @@ this wrapping occurs automatically.
 abstract type InferenceAlgorithm end
 
 DynamicPPL.default_chain_type(sampler::Sampler{<:InferenceAlgorithm}) = MCMCChains.Chains
+
+function DynamicPPL.init_strategy(sampler::Sampler{<:InferenceAlgorithm})
+    return DynamicPPL.InitFromPrior()
+end

src/mcmc/emcee.jl

@@ -31,12 +31,21 @@ struct EmceeState{V<:AbstractVarInfo,S}
     states::S
 end
 
+# Utility function to tetrieve the number of walkers
+_get_n_walkers(e::Emcee) = e.ensemble.n_walkers
+_get_n_walkers(spl::Sampler{<:Emcee}) = _get_n_walkers(spl.alg)
+
+# Because Emcee expects n_walkers initialisations, we need to override this
+function DynamicPPL.init_strategy(spl::Sampler{<:Emcee})
+    return fill(DynamicPPL.InitFromPrior(), _get_n_walkers(spl))
+end
+
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     model::Model,
     spl::Sampler{<:Emcee};
     resume_from=nothing,
-    initial_params=nothing,
+    initial_params,
     kwargs...,
 )
     if resume_from !== nothing
@@ -45,23 +54,19 @@ function AbstractMCMC.step(
     end
 
     # Sample from the prior
-    n = spl.alg.ensemble.n_walkers
-    vis = [VarInfo(rng, model, SampleFromPrior()) for _ in 1:n]
+    n = _get_n_walkers(spl)
+    vis = [VarInfo(rng, model) for _ in 1:n]
 
     # Update the parameters if provided.
-    if initial_params !== nothing
-        length(initial_params) == n ||
-            throw(ArgumentError("initial parameters have to be specified for each walker"))
-        vis = map(vis, initial_params) do vi, init
-            # TODO(DPPL0.38/penelopeysm) This whole thing can be replaced with init!!
-            vi = DynamicPPL.initialize_parameters!!(vi, init, model)
-
-            # Update log joint probability.
-            spl_model = DynamicPPL.contextualize(
-                model, DynamicPPL.SamplingContext(rng, SampleFromPrior(), model.context)
-            )
-            last(DynamicPPL.evaluate!!(spl_model, vi))
-        end
+    if !(
+        initial_params isa AbstractVector{<:DynamicPPL.AbstractInitStrategy} &&
+        length(initial_params) == n
+    )
+        err_msg = "initial_params for `Emcee` must be a vector of `DynamicPPL.AbstractInitStrategy`, with length equal to the number of walkers ($n)"
+        throw(ArgumentError(err_msg))
+    end
+    vis = map(vis, initial_params) do vi, strategy
+        last(DynamicPPL.init!!(rng, model, vi, strategy))
     end
 
     # Compute initial transition and states.

src/mcmc/ess.jl

@@ -82,23 +82,8 @@ EllipticalSliceSampling.isgaussian(::Type{<:ESSPrior}) = true
 
 # Only define out-of-place sampling
 function Base.rand(rng::Random.AbstractRNG, p::ESSPrior)
-    varinfo = p.varinfo
-    # TODO: Surely there's a better way of doing this now that we have `SamplingContext`?
-    # TODO(DPPL0.38/penelopeysm): This can be replaced with `init!!(p.model,
-    # p.varinfo, PriorInit())` after TuringLang/DynamicPPL.jl#984. The reason
-    # why we had to use the 'del' flag before this was because
-    # SampleFromPrior() wouldn't overwrite existing variables.
-    # The main problem I'm rather unsure about is ESS-within-Gibbs. The
-    # current implementation I think makes sure to only resample the variables
-    # that 'belong' to the current ESS sampler. InitContext on the other hand
-    # would resample all variables in the model (??) Need to think about this
-    # carefully.
-    vns = keys(varinfo)
-    for vn in vns
-        set_flag!(varinfo, vn, "del")
-    end
-    p.model(rng, varinfo)
-    return varinfo[:]
+    _, vi = DynamicPPL.init!!(rng, p.model, p.varinfo, DynamicPPL.InitFromPrior())
+    return vi[:]
 end
 
 # Mean of prior distribution

src/mcmc/external_sampler.jl

@@ -117,25 +117,25 @@ function AbstractMCMC.step(
     model::DynamicPPL.Model,
     sampler_wrapper::Sampler{<:ExternalSampler};
     initial_state=nothing,
-    initial_params=nothing,
+    initial_params=DynamicPPL.init_strategy(sampler_wrapper.alg.sampler),
     kwargs...,
 )
     alg = sampler_wrapper.alg
     sampler = alg.sampler
 
     # Initialise varinfo with initial params and link the varinfo if needed.
     varinfo = DynamicPPL.VarInfo(model)
+    _, varinfo = DynamicPPL.init!!(rng, model, varinfo, initial_params)
+
     if requires_unconstrained_space(alg)
-        if initial_params !== nothing
-            # If we have initial parameters, we need to set the varinfo before linking.
-            varinfo = DynamicPPL.link(DynamicPPL.unflatten(varinfo, initial_params), model)
-            # Extract initial parameters in unconstrained space.
-            initial_params = varinfo[:]
-        else
-            varinfo = DynamicPPL.link(varinfo, model)
-        end
+        varinfo = DynamicPPL.link(varinfo, model)
     end
 
+    # We need to extract the vectorised initial_params, because the later call to
+    # AbstractMCMC.step only sees a `LogDensityModel` which expects `initial_params`
+    # to be a vector.
+    initial_params_vector = varinfo[:]
+
     # Construct LogDensityFunction
     f = DynamicPPL.LogDensityFunction(
         model, DynamicPPL.getlogjoint_internal, varinfo; adtype=alg.adtype
@@ -144,15 +144,19 @@ function AbstractMCMC.step(
     # Then just call `AbstractMCMC.step` with the right arguments.
     if initial_state === nothing
         transition_inner, state_inner = AbstractMCMC.step(
-            rng, AbstractMCMC.LogDensityModel(f), sampler; initial_params, kwargs...
+            rng,
+            AbstractMCMC.LogDensityModel(f),
+            sampler;
+            initial_params=initial_params_vector,
+            kwargs...,
         )
     else
         transition_inner, state_inner = AbstractMCMC.step(
             rng,
             AbstractMCMC.LogDensityModel(f),
             sampler,
             initial_state;
-            initial_params,
+            initial_params=initial_params_vector,
             kwargs...,
         )
     end