Update to AbstractMCMC 2 (#1428)

Co-authored-by: Cameron Pfiffer <[email protected]>

Author: devmotion

Project.toml

@@ -1,12 +1,13 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.14.12"
+version = "0.15.0"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
 AdvancedMH = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170"
 AdvancedVI = "b5ca4192-6429-45e5-a2d9-87aec30a685c"
+BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
@@ -31,16 +32,17 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
-AbstractMCMC = "1"
+AbstractMCMC = "2.1"
 AdvancedHMC = "0.2.24"
-AdvancedMH = "0.5.1"
+AdvancedMH = "0.5.2"
 AdvancedVI = "0.1"
+BangBang = "0.3"
 Bijectors = "0.8"
 Distributions = "0.23.3"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.8"
-DynamicPPL = "0.9.5"
-EllipticalSliceSampling = "0.2, 0.3"
+DynamicPPL = "0.10.0"
+EllipticalSliceSampling = "0.3"
 ForwardDiff = "0.10.3"
 Libtask = "0.4, 0.5"
 LogDensityProblems = "^0.9, 0.10"

src/Turing.jl

@@ -75,8 +75,7 @@ end
 # Exports #
 ###########
 # `using` statements for stuff to re-export
-using DynamicPPL: pointwise_loglikelihoods, elementwise_loglikelihoods,
-    generated_quantities, logprior, logjoint
+using DynamicPPL: pointwise_loglikelihoods, generated_quantities, logprior, logjoint
 using StatsBase: predict
 
 # Turing essentials - modelling macros and inference algorithms
@@ -137,8 +136,4 @@ export  @model,                 # modelling
         generated_quantities,
         logprior,
         logjoint
-
-# deprecations
-include("deprecations.jl")
-
 end

src/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!(
+DynamicPPL.initialsampler(::Sampler{<:DynamicNUTS}) = SampleFromUniform()
+
+function DynamicPPL.initialstep(
     rng::AbstractRNG,
     model::Model,
     spl::Sampler{<:DynamicNUTS},
-    N::Integer;
+    vi::AbstractVarInfo;
+    N::Int,
     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!(
+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,
@@ -131,9 +131,9 @@ end
     end
     if resume_from === nothing
         return AbstractMCMC.sample(rng, model, Sampler(alg, model), N;
-                                   chain_type=chain_type, progress=false, kwargs...)
+                                   chain_type=chain_type, progress=false, N=N, kwargs...)
     else
-        return resume(resume_from, N; chain_type=chain_type, progress=false, kwargs...)
+        return resume(resume_from, N; chain_type=chain_type, progress=false, N=N, kwargs...)
     end
 end
 
@@ -152,5 +152,5 @@ function AbstractMCMC.sample(
         @warn "[HMC] Progress logging in Turing is disabled since DynamicHMC provides its own progress meter"
     end
     return AbstractMCMC.sample(rng, model, Sampler(alg, model), parallel, N, n_chains;
-                               chain_type=chain_type, progress=false, kwargs...)
+                               chain_type=chain_type, progress=false, N=N, kwargs...)
 end

src/core/Core.jl

@@ -19,7 +19,6 @@ import ZygoteRules
 
 include("container.jl")
 include("ad.jl")
-include("deprecations.jl")
 
 function __init__()
     @require ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" begin