Fix some tests

Author: penelopeysm

src/sampler.jl

@@ -59,8 +59,8 @@ function AbstractMCMC.step(
 )
     vi = VarInfo()
     strategy = sampler isa SampleFromPrior ? PriorInit() : UniformInit()
-    DynamicPPL.init!!(rng, model, vi, strategy)
-    return vi, nothing
+    _, new_vi = DynamicPPL.init!!(rng, model, vi, strategy)
+    return new_vi, nothing
 end
 
 """

test/model.jl

@@ -495,7 +495,11 @@ const GDEMO_DEFAULT = DynamicPPL.TestUtils.demo_assume_observe_literal()
 
             # Construct a chain with 'sampled values' of β
             ground_truth_β = 2
-            β_chain = MCMCChains.Chains(rand(Normal(ground_truth_β, 0.002), 1000), [:β])
+            β_chain = MCMCChains.Chains(
+                rand(Normal(ground_truth_β, 0.002), 1000),
+                [:β];
+                info=(; varname_to_symbol=Dict(@varname(β) => :β)),
+            )
 
             # Generate predictions from that chain
             xs_test = [10 + 0.1, 10 + 2 * 0.1]
@@ -541,7 +545,9 @@ const GDEMO_DEFAULT = DynamicPPL.TestUtils.demo_assume_observe_literal()
             @testset "prediction from multiple chains" begin
                 # Normal linreg model
                 multiple_β_chain = MCMCChains.Chains(
-                    reshape(rand(Normal(ground_truth_β, 0.002), 1000, 2), 1000, 1, 2), [:β]
+                    reshape(rand(Normal(ground_truth_β, 0.002), 1000, 2), 1000, 1, 2),
+                    [:β];
+                    info=(; varname_to_symbol=Dict(@varname(β) => :β)),
                 )
                 predictions = DynamicPPL.predict(m_lin_reg_test, multiple_β_chain)
                 @test size(multiple_β_chain, 3) == size(predictions, 3)

test/sampler.jl

@@ -25,8 +25,8 @@
         @test length(chains) == N
 
         # `m` is Gaussian, i.e. no transformation is used, so it
-        # should have a mean equal to its prior, i.e. 2.
-        @test mean(vi[@varname(m)] for vi in chains) ≈ 2 atol = 0.1
+        # will be drawn from U[-2, 2] and its mean should be 0.
+        @test mean(vi[@varname(m)] for vi in chains) ≈ 0.0 atol = 0.1
 
         # Expected value of ``exp(X)`` where ``X ~ U[-2, 2]`` is ≈ 1.8.
         @test mean(vi[@varname(s)] for vi in chains) ≈ 1.8 atol = 0.1
@@ -81,10 +81,8 @@
             model = coinflip()
             sampler = Sampler(alg)
             lptrue = logpdf(Binomial(25, 0.2), 10)
-            let inits = (; p=0.2)
-                chain = sample(
-                    model, sampler, 1; initial_params=ParamsInit(inits), progress=false
-                )
+            let inits = ParamsInit((; p=0.2))
+                chain = sample(model, sampler, 1; initial_params=inits, progress=false)
                 @test chain[1].metadata.p.vals == [0.2]
                 @test getlogjoint(chain[1]) == lptrue
 
@@ -111,10 +109,8 @@
             end
             model = twovars()
             lptrue = logpdf(InverseGamma(2, 3), 4) + logpdf(Normal(0, 2), -1)
-            for inits in ([4, -1], (; s=4, m=-1))
-                chain = sample(
-                    model, sampler, 1; initial_params=ParamsInit(inits), progress=false
-                )
+            let inits = ParamsInit((; s=4, m=-1))
+                chain = sample(model, sampler, 1; initial_params=inits, progress=false)
                 @test chain[1].metadata.s.vals == [4]
                 @test chain[1].metadata.m.vals == [-1]
                 @test getlogjoint(chain[1]) == lptrue
@@ -126,7 +122,7 @@
                     MCMCThreads(),
                     1,
                     10;
-                    initial_params=fill(ParamsInit(inits), 10),
+                    initial_params=fill(inits, 10),
                     progress=false,
                 )
                 for c in chains
@@ -137,10 +133,8 @@
             end
 
             # set only m = -1
-            for inits in ((; s=missing, m=-1), (; m=-1))
-                chain = sample(
-                    model, sampler, 1; initial_params=ParamsInit(inits), progress=false
-                )
+            for inits in (ParamsInit((; s=missing, m=-1)), ParamsInit((; m=-1)))
+                chain = sample(model, sampler, 1; initial_params=inits, progress=false)
                 @test !ismissing(chain[1].metadata.s.vals[1])
                 @test chain[1].metadata.m.vals == [-1]