Merge pull request #34 from TuringLang/cpfiffer-patch-1

Updated AbstractMCMC compat

Author: yebai

Project.toml

@@ -10,7 +10,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 
 [compat]
-AbstractMCMC = "~0.1"
+AbstractMCMC = "0.3"
 AdvancedHMC = "0.2.20"
 Bijectors = "0.5.2"
 Distributions = "0.22"

test/Turing/inference/Inference.jl

@@ -144,20 +144,22 @@ function AbstractMCMC.sample(
     model::AbstractModel,
     alg::InferenceAlgorithm,
     N::Integer;
+    chain_type=Chains,
     kwargs...
 )
-    return sample(rng, model, Sampler(alg, model), N; progress=PROGRESS[], kwargs...)
+    return sample(rng, model, Sampler(alg, model), N; progress=PROGRESS[], chain_type=chain_type, kwargs...)
 end
 
 function AbstractMCMC.sample(
     model::AbstractModel,
     alg::InferenceAlgorithm,
     N::Integer;
     resume_from=nothing,
+    chain_type=Chains,
     kwargs...
 )
     if resume_from === nothing
-        return sample(model, Sampler(alg, model), N; progress=PROGRESS[], kwargs...)
+        return sample(model, Sampler(alg, model), N; progress=PROGRESS[], chain_type=chain_type, kwargs...)
     else
         return resume(resume_from, N)
     end
@@ -169,9 +171,10 @@ function AbstractMCMC.psample(
     alg::InferenceAlgorithm,
     N::Integer,
     n_chains::Integer;
+    chain_type=Chains,
     kwargs...
 )
-    return psample(GLOBAL_RNG, model, alg, N, n_chains; progress=false, kwargs...)
+    return psample(GLOBAL_RNG, model, alg, N, n_chains; progress=false, chain_type=chain_type, kwargs...)
 end
 
 function AbstractMCMC.psample(
@@ -180,9 +183,10 @@ function AbstractMCMC.psample(
     alg::InferenceAlgorithm,
     N::Integer,
     n_chains::Integer;
+    chain_type=Chains,
     kwargs...
 )
-    return psample(rng, model, Sampler(alg, model), N, n_chains; progress=false, kwargs...)
+    return psample(rng, model, Sampler(alg, model), N, n_chains; progress=false, chain_type=chain_type, kwargs...)
 end
 
 function AbstractMCMC.sample_init!(
@@ -318,7 +322,8 @@ function AbstractMCMC.bundle_samples(
     model::AbstractModel,
     spl::Sampler,
     N::Integer,
-    ts::Vector{<:AbstractTransition};
+    ts::Vector{<:AbstractTransition},
+    ct::Type{Chains};
     discard_adapt::Bool=true,
     save_state=true,
     kwargs...
@@ -375,7 +380,7 @@ function save(c::Chains, spl::AbstractSampler, model, vi, samples)
     return setinfo(c, merge(nt, c.info))
 end
 
-function resume(c::Chains, n_iter::Int; kwargs...)
+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.
@@ -386,6 +391,7 @@ function resume(c::Chains, n_iter::Int; kwargs...)
         n_iter;
         resume_from=c,
         reuse_spl_n=n_iter,
+        chain_type=chain_type,
         kwargs...
     )