Fix more tests

Author: penelopeysm

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

@@ -36,7 +36,9 @@ _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
-DynamicPPL.init_strategy(spl::Sampler{<:Emcee}) = fill(InitFromPrior(), _get_n_walkers(spl))
+function DynamicPPL.init_strategy(spl::Sampler{<:Emcee})
+    return fill(DynamicPPL.InitFromPrior(), _get_n_walkers(spl))
+end
 
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,

src/mcmc/hmc.jl

@@ -90,6 +90,7 @@ function AbstractMCMC.sample(
     N::Integer;
     chain_type=DynamicPPL.default_chain_type(sampler),
     resume_from=nothing,
+    initial_params=DynamicPPL.init_strategy(sampler),
     initial_state=DynamicPPL.loadstate(resume_from),
     progress=PROGRESS[],
     nadapts=sampler.alg.n_adapts,

src/mcmc/is.jl

@@ -37,7 +37,7 @@ function AbstractMCMC.step(
 )
     model = DynamicPPL.setleafcontext(model, ISContext(rng))
     _, vi = DynamicPPL.evaluate!!(model, DynamicPPL.VarInfo())
-    vi = DynamicPPL.typed_varinfo(vi, model)
+    vi = DynamicPPL.typed_varinfo(vi)
     return Transition(model, vi, nothing), nothing
 end
 

src/mcmc/repeat_sampler.jl

@@ -81,3 +81,58 @@ function AbstractMCMC.step_warmup(
     end
     return transition, state
 end
+
+# Need some extra leg work to make RepeatSampler work seamlessly with DynamicPPL models +
+# samplers, instead of generic AbstractMCMC samplers.
+
+function DynamicPPL.init_strategy(spl::RepeatSampler{<:Sampler})
+    return DynamicPPL.init_strategy(spl.sampler)
+end
+
+function AbstractMCMC.sample(
+    rng::AbstractRNG,
+    model::DynamicPPL.Model,
+    sampler::RepeatSampler{<:Sampler},
+    N::Integer;
+    initial_params=DynamicPPL.init_strategy(sampler),
+    chain_type=MCMCChains.Chains,
+    progress=PROGRESS[],
+    kwargs...,
+)
+    return AbstractMCMC.mcmcsample(
+        rng,
+        model,
+        sampler,
+        N;
+        initial_params=initial_params,
+        chain_type=chain_type,
+        progress=progress,
+        kwargs...,
+    )
+end
+
+function AbstractMCMC.sample(
+    rng::AbstractRNG,
+    model::DynamicPPL.Model,
+    sampler::RepeatSampler{<:Sampler},
+    ensemble::AbstractMCMC.AbstractMCMCEnsemble,
+    N::Integer,
+    n_chains::Integer;
+    initial_params=fill(DynamicPPL.init_strategy(sampler), n_chains),
+    chain_type=MCMCChains.Chains,
+    progress=PROGRESS[],
+    kwargs...,
+)
+    return AbstractMCMC.mcmcsample(
+        rng,
+        model,
+        sampler,
+        ensemble,
+        N,
+        n_chains;
+        initial_params=initial_params,
+        chain_type=chain_type,
+        progress=progress,
+        kwargs...,
+    )
+end

test/mcmc/repeat_sampler.jl

@@ -2,8 +2,8 @@ module RepeatSamplerTests
 
 using ..Models: gdemo_default
 using DynamicPPL: Sampler
-using MCMCChains: Chains
-using StableRNGs: StableRNG
+using MCMCChains: MCMCChains
+using Random: Xoshiro
 using Test: @test, @testset
 using Turing
 
@@ -14,10 +14,12 @@ using Turing
     num_samples = 10
     num_chains = 2
 
-    rng = StableRNG(0)
+    # Use Xoshiro instead of StableRNGs as the output should always be
+    # similar regardless of what kind of random seed is used (as long
+    # as there is a random seed).
     for sampler in [MH(), Sampler(HMC(0.01, 4))]
         chn1 = sample(
-            copy(rng),
+            Xoshiro(0),
             gdemo_default,
             sampler,
             MCMCThreads(),
@@ -27,15 +29,16 @@ using Turing
         )
         repeat_sampler = RepeatSampler(sampler, num_repeats)
         chn2 = sample(
-            copy(rng),
+            Xoshiro(0),
             gdemo_default,
             repeat_sampler,
             MCMCThreads(),
             num_samples,
-            num_chains;
-            chain_type=Chains,
+            num_chains,
         )
         # isequal to avoid comparing `missing`s in chain stats
+        @test chn1 isa MCMCChains.Chains
+        @test chn2 isa MCMCChains.Chains
         @test isequal(chn1.value, chn2.value)
     end
 end