Update Turing source + tests

Author: devmotion

test/Turing/Turing.jl

@@ -12,12 +12,11 @@ using Requires, Reexport, ForwardDiff
 using Bijectors, StatsFuns, SpecialFunctions
 using Statistics, LinearAlgebra, ProgressMeter
 using Markdown, Libtask, MacroTools
-using AbstractMCMC
+using AbstractMCMC: sample, psample
 @reexport using Distributions, MCMCChains, Libtask
 using Tracker: Tracker
 
 import Base: ~, ==, convert, hash, promote_rule, rand, getindex, setindex!
-import MCMCChains: AbstractChains, Chains
 import DynamicPPL: getspace, runmodel!
 
 const PROGRESS = Ref(true)
@@ -85,9 +84,6 @@ export  @model,                 # modelling
         SMC,
         CSMC,
         PG,
-        PIMH,
-        PMMH,
-        IPMCMC,
 
         vi,                    # variational inference
         ADVI,

test/Turing/contrib/inference/dynamichmc.jl

@@ -1,4 +1,4 @@
-using AbstractMCMC: init_callback, NoCallback
+using AbstractMCMC: NoCallback
 
 ###
 ### DynamicHMC backend - https://github.com/tpapp/DynamicHMC.jl
@@ -50,7 +50,7 @@ end
 
 getspace(::DynamicNUTS{<:Any, space}) where {space} = space
 
-function sample_init!(
+function AbstractMCMC.sample_init!(
     rng::AbstractRNG,
     model::Model,
     spl::Sampler{<:DynamicNUTS},
@@ -84,11 +84,12 @@ function sample_init!(
     spl.state.draws = results.chain
 end
 
-function step!(
+function AbstractMCMC.step!(
     rng::AbstractRNG,
     model::Model,
     spl::Sampler{<:DynamicNUTS},
-    N::Integer;
+    N::Integer,
+    transition;
     kwargs...
 )
     # Pop the next draw off the vector.

test/Turing/inference/AdvancedSMC.jl

@@ -7,23 +7,21 @@
 #######################
 
 """
-    ParticleTransition{T, F<:AbstractFloat} <: AbstractTransition
+    ParticleTransition{T, F<:AbstractFloat}
 
 Fields:
 - `θ`: The parameters for any given sample.
 - `lp`: The log pdf for the sample's parameters.
 - `le`: The log evidence retrieved from the particle.
 - `weight`: The weight of the particle the sample was retrieved from.
 """
-struct ParticleTransition{T, F<:AbstractFloat} <: AbstractTransition
+struct ParticleTransition{T, F<:AbstractFloat}
     θ::T
     lp::F
     le::F
     weight::F
 end
 
-transition_type(spl::Sampler{<:ParticleInference}) = ParticleTransition
-
 function additional_parameters(::Type{<:ParticleTransition})
     return [:lp,:le, :weight]
 end
@@ -76,9 +74,9 @@ function Sampler(alg::SMC, model::Model, s::Selector)
     return Sampler(alg, dict, s, state)
 end
 
-function sample_init!(
+function AbstractMCMC.sample_init!(
     ::AbstractRNG,
-    model::Turing.Model,
+    model::Model,
     spl::Sampler{<:SMC},
     N::Integer;
     kwargs...
@@ -105,12 +103,13 @@ function sample_init!(
     end
 end
 
-function step!(
+function AbstractMCMC.step!(
     ::AbstractRNG,
-    model::Turing.Model,
+    model::Model,
     spl::Sampler{<:SMC},
-    ::Integer;
-    iteration=-1,
+    ::Integer,
+    transition;
+    iteration = -1,
     kwargs...
 )
     # check that we received a real iteration number
@@ -182,11 +181,12 @@ function Sampler(alg::PG, model::Model, s::Selector)
     return Sampler(alg, info, s, state)
 end
 
-function step!(
+function AbstractMCMC.step!(
     ::AbstractRNG,
-    model::Turing.Model,
+    model::Model,
     spl::Sampler{<:PG},
-    ::Integer;
+    ::Integer,
+    transition;
     kwargs...
 )
     # obtain or create reference particle
@@ -232,12 +232,12 @@ function step!(
     return ParticleTransition(params, lp, pc.logE, 1.0)
 end
 
-function sample_end!(
+function AbstractMCMC.sample_end!(
     ::AbstractRNG,
     ::Model,
     spl::Sampler{<:ParticleInference},
     N::Integer,
-    ts::Vector{ParticleTransition};
+    ts::Vector{<:ParticleTransition};
     kwargs...
 )
     # Set the default for resuming the sampler.

test/Turing/inference/Inference.jl

@@ -12,20 +12,18 @@ using ProgressMeter, LinearAlgebra
 using ..Turing: PROGRESS, NamedDist, NoDist, Turing
 using StatsFuns: logsumexp
 using Random: GLOBAL_RNG, AbstractRNG, randexp
-using AbstractMCMC, DynamicPPL
+using DynamicPPL
 using Bijectors: _debug
 
 import MCMCChains: Chains
 import AdvancedHMC; const AHMC = AdvancedHMC
 import ..Core: getchunksize, getADtype
-import AbstractMCMC: AbstractTransition, sample, step!, sample_init!,
-    transitions_init, sample_end!, AbstractSampler, transition_type,
-    callback, init_callback, AbstractCallback, psample
+import AbstractMCMC
+using AbstractMCMC: AbstractModel, AbstractCallback, AbstractSampler
 import DynamicPPL: tilde, dot_tilde, getspace, get_matching_type
 
 export  InferenceAlgorithm,
         Hamiltonian,
-        AbstractGibbs,
         GibbsComponent,
         StaticHamiltonian,
         AdaptiveHamiltonian,
@@ -44,20 +42,8 @@ export  InferenceAlgorithm,
         SMC,
         CSMC,
         PG,
-        PIMH,
-        PMMH,
-        IPMCMC,  # particle-based sampling
         assume,
         observe,
-        step,
-        WelfordVar,
-        WelfordCovar,
-        NaiveCovar,
-        get_var,
-        get_covar,
-        add_sample!,
-        reset!,
-        step!,
         resume
 
 #######################
@@ -95,7 +81,7 @@ end
 # Default Transition #
 ######################
 
-struct Transition{T, F<:AbstractFloat} <: AbstractTransition
+struct Transition{T, F<:AbstractFloat}
     θ  :: T
     lp :: F
 end
@@ -147,19 +133,19 @@ function AbstractMCMC.sample(
     chain_type=Chains,
     kwargs...
 )
-    return sample(rng, model, Sampler(alg, model), N; progress=PROGRESS[], chain_type=chain_type, kwargs...)
+    return AbstractMCMC.sample(rng, model, Sampler(alg, model), N; progress=PROGRESS[], chain_type=chain_type, kwargs...)
 end
 
 function AbstractMCMC.sample(
-    model::AbstractModel,
+    model::Model,
     alg::InferenceAlgorithm,
     N::Integer;
     resume_from=nothing,
     chain_type=Chains,
     kwargs...
 )
     if resume_from === nothing
-        return sample(model, Sampler(alg, model), N; progress=PROGRESS[], chain_type=chain_type, kwargs...)
+        return AbstractMCMC.sample(model, Sampler(alg, model), N; progress=PROGRESS[], chain_type=chain_type, kwargs...)
     else
         return resume(resume_from, N)
     end
@@ -174,7 +160,7 @@ function AbstractMCMC.psample(
     chain_type=Chains,
     kwargs...
 )
-    return psample(GLOBAL_RNG, model, alg, N, n_chains; progress=false, chain_type=chain_type, kwargs...)
+    return AbstractMCMC.psample(GLOBAL_RNG, model, alg, N, n_chains; progress=false, chain_type=chain_type, kwargs...)
 end
 
 function AbstractMCMC.psample(
@@ -186,7 +172,7 @@ function AbstractMCMC.psample(
     chain_type=Chains,
     kwargs...
 )
-    return psample(rng, model, Sampler(alg, model), N, n_chains; progress=false, chain_type=chain_type, kwargs...)
+    return AbstractMCMC.psample(rng, model, Sampler(alg, model), N, n_chains; progress=false, chain_type=chain_type, kwargs...)
 end
 
 function AbstractMCMC.sample_init!(
@@ -206,9 +192,9 @@ end
 function AbstractMCMC.sample_end!(
     ::AbstractRNG,
     ::Model,
-    ::AbstractSampler,
+    ::Sampler,
     ::Integer,
-    ::Vector{<:AbstractTransition};
+    ::Vector;
     kwargs...
 )
     # Silence the default API function.
@@ -244,7 +230,7 @@ end
 # Chain making utilities #
 ##########################
 
-function _params_to_array(ts::Vector{<:AbstractTransition}, spl::Sampler)
+function _params_to_array(ts::Vector, spl::Sampler)
     names_set = Set{String}()
     # Extract the parameter names and values from each transition.
     dicts = map(ts) do t
@@ -276,7 +262,7 @@ function flatten_namedtuple(nt::NamedTuple)
     return [vn[1] for vn in names_vals], [vn[2] for vn in names_vals]
 end
 
-function get_transition_extras(ts::Vector{<:AbstractTransition})
+function get_transition_extras(ts::Vector)
     # Get the extra field names from the sampler state type.
     # This handles things like :lp or :weight.
     extra_params = additional_parameters(eltype(ts))
@@ -322,8 +308,8 @@ function AbstractMCMC.bundle_samples(
     model::AbstractModel,
     spl::Sampler,
     N::Integer,
-    ts::Vector{<:AbstractTransition},
-    ct::Type{Chains};
+    ts::Vector,
+    ::Type{Chains};
     discard_adapt::Bool=true,
     save_state=true,
     kwargs...
@@ -384,7 +370,7 @@ function resume(c::Chains, n_iter::Int; chain_type=Chains, kwargs...)
     @assert !isempty(c.info) "[Turing] cannot resume from a chain without state info"
 
     # Sample a new chain.
-    newchain = sample(
+    newchain = AbstractMCMC.sample(
         c.info[:range],
         c.info[:model],
         c.info[:spl],
@@ -432,13 +418,12 @@ include("is.jl")
 include("AdvancedSMC.jl")
 include("gibbs.jl")
 include("../contrib/inference/sghmc.jl")
-include("../contrib/inference/AdvancedSMCExtensions.jl")
 
 ################
 # Typing tools #
 ################
 
-for alg in (:SMC, :PG, :PMMH, :IPMCMC, :MH, :IS, :ESS, :Gibbs)
+for alg in (:SMC, :PG, :MH, :IS, :ESS, :Gibbs)
     @eval getspace(::$alg{space}) where {space} = space
 end
 for alg in (:HMC, :HMCDA, :NUTS, :SGLD, :SGHMC)
@@ -494,7 +479,6 @@ end
 ## Fallback functions
 
 alg_str(spl::Sampler) = string(nameof(typeof(spl.alg)))
-transition_type(spl::Sampler) = typeof(Transition(spl))
 
 # utility funcs for querying sampler information
 require_gradient(spl::Sampler) = false

test/Turing/inference/ess.jl

@@ -54,16 +54,23 @@ isgaussian(::NormalCanon) = true
 isgaussian(::AbstractMvNormal) = true
 
 # always accept in the first step
-function step!(::AbstractRNG, model::Model, spl::Sampler{<:ESS}, ::Integer; kwargs...)
+function AbstractMCMC.step!(
+    ::AbstractRNG,
+    ::Model,
+    spl::Sampler{<:ESS},
+    ::Integer,
+    ::Nothing;
+    kwargs...
+)
     return Transition(spl)
 end
 
-function step!(
+function AbstractMCMC.step!(
     rng::AbstractRNG,
     model::Model,
     spl::Sampler{<:ESS},
     ::Integer,
-    ::Transition;
+    transition;
     kwargs...
 )
     # obtain mean of distribution