Try #150:

Author: bors[bot]

Project.toml

@@ -12,7 +12,7 @@ NaturalSort = "c020b1a1-e9b0-503a-9c33-f039bfc54a85"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
 [compat]
-AbstractMCMC = "1"
+AbstractMCMC = "2"
 Bijectors = "0.5.2, 0.6, 0.7, 0.8"
 ChainRulesCore = "0.9.7"
 Distributions = "0.23.8"

src/DynamicPPL.jl

@@ -61,7 +61,6 @@ export  AbstractVarInfo,
         Sample,
         init,
         vectorize,
-        set_resume!,
 # Model
         Model,
         getmissings,

src/model.jl

@@ -109,9 +109,6 @@ See also: [`evaluate_threadsafe`](@ref)
 """
 function evaluate_threadunsafe(rng, model, varinfo, sampler, context)
     resetlogp!(varinfo)
-    if has_eval_num(sampler)
-        sampler.state.eval_num += 1
-    end
     return _evaluate(rng, model, varinfo, sampler, context)
 end
 
@@ -128,9 +125,6 @@ See also: [`evaluate_threadunsafe`](@ref)
 """
 function evaluate_threadsafe(rng, model, varinfo, sampler, context)
     resetlogp!(varinfo)
-    if has_eval_num(sampler)
-        sampler.state.eval_num += 1
-    end
     wrapper = ThreadSafeVarInfo(varinfo)
     result = _evaluate(rng, model, wrapper, sampler, context)
     setlogp!(varinfo, getlogp(wrapper))

src/prob_macro.jl

@@ -136,7 +136,7 @@ function logprior(
 
     # When all of model args are on the lhs of |, this is also equal to the logjoint.
     model = make_prior_model(left, right, _model)
-    vi = _vi === nothing ? VarInfo(deepcopy(model), PriorContext()) : _vi
+    vi = _vi === nothing ? VarInfo(deepcopy(model), SampleFromPrior(), PriorContext()) : _vi
     foreach(keys(vi.metadata)) do n
         @assert n in keys(left) "Variable $n is not defined."
     end

src/sampler.jl

@@ -1,3 +1,5 @@
+# TODO: Make `UniformSampling` and `Prior` algs + just use `Sampler`
+# That would let us use all defaults for Sampler, combine it with other samplers etc.
 """
 Robust initialization method for model parameters in Hamiltonian samplers.
 """
@@ -17,55 +19,93 @@ function init(rng, dist, ::SampleFromUniform, n::Int)
     return istransformable(dist) ? inittrans(rng, dist, n) : rand(rng, dist, n)
 end
 
-"""
-    has_eval_num(spl::AbstractSampler)
-
-Check whether `spl` has a field called `eval_num` in its state variables or not.
-"""
-has_eval_num(spl::SampleFromUniform) = false
-has_eval_num(spl::SampleFromPrior) = false
-has_eval_num(spl::AbstractSampler) = :eval_num in fieldnames(typeof(spl.state))
-
-"""
-An abstract type that mutable sampler state structs inherit from.
-"""
-abstract type AbstractSamplerState end
-
 """
     Sampler{T}
 
-Generic interface for implementing inference algorithms.
-An implementation of an algorithm should include the following:
-
-1. A type specifying the algorithm and its parameters, derived from InferenceAlgorithm
-2. A method of `sample` function that produces results of inference, which is where actual inference happens.
-
-DynamicPPL translates models to chunks that call the modelling functions at specified points.
-The dispatch is based on the value of a `sampler` variable.
-To include a new inference algorithm implements the requirements mentioned above in a separate file,
-then include that file at the end of this one.
+Generic sampler type for inference algorithms in DynamicPPL.
 """
-mutable struct Sampler{T, S<:AbstractSamplerState} <: AbstractSampler
-    alg      ::  T
-    info     ::  Dict{Symbol, Any} # sampler infomation
-    selector ::  Selector
-    state    ::  S
+struct Sampler{T} <: AbstractSampler
+    alg::T
+    # TODO: remove selector & add space
+    selector::Selector
 end
 Sampler(alg) = Sampler(alg, Selector())
 Sampler(alg, model::Model) = Sampler(alg, model, Selector())
-Sampler(alg, model::Model, s::Selector) = Sampler(alg, model, s)
+Sampler(alg, model::Model, s::Selector) = Sampler(alg, s)
 
 # AbstractMCMC interface for SampleFromUniform and SampleFromPrior
-
-function AbstractMCMC.step!(
+function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     model::Model,
     sampler::Union{SampleFromUniform,SampleFromPrior},
-    ::Integer,
-    transition;
+    state = nothing;
     kwargs...
 )
     vi = VarInfo()
-    model(vi, sampler)
-    return vi
+    model(rng, vi, sampler)
+    return vi, nothing
+end
+
+# initial step: general interface for resuming and
+function AbstractMCMC.step(
+    rng::Random.AbstractRNG,
+    model::Model,
+    spl::Sampler;
+    resume_from = nothing,
+    kwargs...
+)
+    if resume_from !== nothing
+        state = loadstate(resume_from)
+        return AbstractMCMC.step(rng, model, spl, state; kwargs...)
+    end
+
+    # Sample initial values.
+    _spl = initialsampler(spl)
+    vi = VarInfo(rng, model, _spl)
+
+    # Update the parameters if provided.
+    if haskey(kwargs, :init_params)
+        initialize_parameters!(vi, kwargs[:init_params], spl)
+
+        # Update joint log probability.
+        model(rng, vi, _spl)
+    end
+
+    return initialstep(rng, model, spl, vi; kwargs...)
+end
+
+function loadstate end
+
+initialsampler(spl::Sampler) = SampleFromPrior()
+
+function initialstep end
+
+function initialize_parameters!(vi::AbstractVarInfo, init_params, spl::Sampler)
+    @debug "Using passed-in initial variable values" init_params
+
+    # Flatten parameters.
+    init_theta = mapreduce(vcat, init_params) do x
+        vec([x;])
+    end
+
+    # Get all values.
+    linked = islinked(vi, spl)
+    linked && invlink!(vi, spl)
+    theta = vi[spl]
+    length(theta) == length(init_theta_flat) ||
+        error("Provided initial value doesn't match the dimension of the model")
+
+    # Update values that are provided.
+    for i in 1:length(init_theta)
+        x = init_theta[i]
+        if x !== missing
+            theta[i] = x
+        end
+    end
+
+    # Update in `vi`.
+    vi[spl] = theta
+    linked && link!(vi, spl)
+
+    return
 end

src/varinfo.jl

@@ -105,21 +105,30 @@ end
 const UntypedVarInfo = VarInfo{<:Metadata}
 const TypedVarInfo = VarInfo{<:NamedTuple}
 
-function VarInfo(model::Model, ctx = DefaultContext())
-    vi = VarInfo()
-    model(vi, SampleFromPrior(), ctx)
-    return TypedVarInfo(vi)
-end
-
 function VarInfo(old_vi::UntypedVarInfo, spl, x::AbstractVector)
     new_vi = deepcopy(old_vi)
     new_vi[spl] = x
     return new_vi
 end
+
 function VarInfo(old_vi::TypedVarInfo, spl, x::AbstractVector)
     md = newmetadata(old_vi.metadata, Val(getspace(spl)), x)
     VarInfo(md, Base.RefValue{eltype(x)}(getlogp(old_vi)), Ref(get_num_produce(old_vi)))
 end
+
+VarInfo(model::Model, args...) = VarInfo(Random.GLOBAL_RNG, model, args...)
+
+function VarInfo(
+    rng::Random.AbstractRNG,
+    model::Model,
+    sampler::AbstractSampler = SampleFromPrior(),
+    context::AbstractContext = DefaultContext(),
+)
+    varinfo = VarInfo()
+    model(rng, varinfo, sampler, context)
+    return TypedVarInfo(varinfo)
+end
+
 @generated function newmetadata(metadata::NamedTuple{names}, ::Val{space}, x) where {names, space}
     exprs = []
     offset = :(0)
@@ -1000,7 +1009,6 @@ from a distribution `dist` to `VarInfo` `vi`.
 The sampler is passed here to invalidate its cache where defined.
 """
 function push!(vi::AbstractVarInfo, vn::VarName, r, dist::Distribution, spl::Sampler)
-    spl.info[:cache_updated] = CACHERESET
     return push!(vi, vn, r, dist, spl.selector)
 end
 function push!(vi::AbstractVarInfo, vn::VarName, r, dist::Distribution, spl::AbstractSampler)

test/Project.toml

@@ -2,6 +2,7 @@
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
 AdvancedMH = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170"
+BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -32,14 +33,15 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-AbstractMCMC = "1.0.1"
+AbstractMCMC = "2"
 AdvancedHMC = "0.2.25"
-AdvancedMH = "0.5.1"
+AdvancedMH = "0.5.2"
+BangBang = "0.3"
 Bijectors = "0.8.2"
 Distributions = "0.23.8"
 DistributionsAD = "0.6.3"
 DocStringExtensions = "0.8.2"
-EllipticalSliceSampling = "0.2.2"
+EllipticalSliceSampling = "0.3"
 ForwardDiff = "0.10.12"
 Libtask = "0.4.1"
 LogDensityProblems = "0.10.3"

test/Turing/contrib/inference/dynamichmc.jl

@@ -41,82 +41,82 @@ function DynamicNUTS{AD}(space::Symbol...) where AD
     DynamicNUTS{AD, space}()
 end
 
-mutable struct DynamicNUTSState{V<:VarInfo, D} <: AbstractSamplerState
+struct DynamicNUTSState{V<:AbstractVarInfo,D}
     vi::V
     draws::Vector{D}
 end
 
 DynamicPPL.getspace(::DynamicNUTS{<:Any, space}) where {space} = space
 
-function AbstractMCMC.sample_init!(
+# initial step: general interface for resuming and
+DynamicPPL.initialsampler(::Sampler{<:DynamicNUTS}) = SampleFromUniform()
+
+function DynamicPPL.initialstep(
     rng::AbstractRNG,
     model::Model,
     spl::Sampler{<:DynamicNUTS},
-    N::Integer;
+    vi::AbstractVarInfo;
     kwargs...
 )
     # Set up lp function.
     function _lp(x)
-        gradient_logp(x, spl.state.vi, model, spl)
+        gradient_logp(x, vi, model, spl)
     end
 
-    # Set the parameters to a starting value.
-    initialize_parameters!(spl; kwargs...)
-
-    model(spl.state.vi, SampleFromUniform())
-    link!(spl.state.vi, spl)
-    l, dl = _lp(spl.state.vi[spl])
+    link!(vi, spl)
+    l, dl = _lp(vi[spl])
     while !isfinite(l) || !isfinite(dl)
-        model(spl.state.vi, SampleFromUniform())
-        link!(spl.state.vi, spl)
-        l, dl = _lp(spl.state.vi[spl])
+        model(vi, SampleFromUniform())
+        link!(vi, spl)
+        l, dl = _lp(vi[spl])
     end
 
-    if spl.selector.tag == :default && !islinked(spl.state.vi, spl)
-        link!(spl.state.vi, spl)
-        model(spl.state.vi, spl)
+    if spl.selector.tag == :default && !islinked(vi, spl)
+        link!(vi, spl)
+        model(vi, spl)
     end
 
     results = mcmc_with_warmup(
         rng,
         FunctionLogDensity(
-            length(spl.state.vi[spl]),
+            length(vi[spl]),
             _lp
         ),
         N
     )
+    draws = results.chain
 
-    spl.state.draws = results.chain
+    # Compute first transition and state.
+    draw = popfirst!(draws)
+    vi[spl] = draw
+    transition = Transition(vi)
+    state = DynamicNUTSState(vi, draws)
+
+    return transition, state
 end
 
 function AbstractMCMC.step!(
     rng::AbstractRNG,
     model::Model,
     spl::Sampler{<:DynamicNUTS},
-    N::Integer,
-    transition;
+    state::DynamicNUTSState;
     kwargs...
 )
+    # Extract VarInfo object.
+    vi = state.vi
+
     # Pop the next draw off the vector.
-    draw = popfirst!(spl.state.draws)
-    spl.state.vi[spl] = draw
-    return Transition(spl)
-end
+    draw = popfirst!(state.draws)
+    vi[spl] = draw
 
-function Sampler(
-    alg::DynamicNUTS,
-    model::Model,
-    s::Selector=Selector()
-)
-    # Construct a state, using a default function.
-    state = DynamicNUTSState(VarInfo(model), [])
+    # Compute next transition.
+    transition = Transition(vi)
 
-    # Return a new sampler.
-    return Sampler(alg, Dict{Symbol,Any}(), s, state)
+    return transition, state
 end
 
- # Disable the progress logging for DynamicHMC, since it has its own progress meter.
- function AbstractMCMC.sample(
+# Disable the progress logging for DynamicHMC, since it has its own progress meter.
+function AbstractMCMC.sample(
     rng::AbstractRNG,
     model::AbstractModel,
     alg::DynamicNUTS,