Fix a bunch of tests; I'll let CI tell me what's still broken...

Author: penelopeysm

Project.toml

@@ -92,4 +92,4 @@ DynamicHMC = "bbc10e6e-7c05-544b-b16e-64fede858acb"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 
 [sources]
-DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "breaking"}
+DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "py/setleafcontext-model"}

src/mcmc/hmc.jl

@@ -161,7 +161,7 @@ function find_initial_params(
             @warn "failed to find valid initial parameters in $(attempts) tries; consider providing a different initialisation strategy with the `initial_params` keyword"
 
         # Resample and try again.
-        varinfo = DynamicPPL.init!!(rng, model, varinfo, init_strategy)
+        _, varinfo = DynamicPPL.init!!(rng, model, varinfo, init_strategy)
     end
 
     # if we failed to find valid initial parameters, error

src/mcmc/prior.jl

@@ -13,7 +13,7 @@ function AbstractMCMC.step(
     kwargs...,
 )
     vi = DynamicPPL.setaccs!!(
-        vi,
+        DynamicPPL.VarInfo(),
         (
             DynamicPPL.ValuesAsInModelAccumulator(true),
             DynamicPPL.LogPriorAccumulator(),

test/Project.toml

@@ -53,7 +53,6 @@ Combinatorics = "1"
 Distributions = "0.25"
 DistributionsAD = "0.6.3"
 DynamicHMC = "2.1.6, 3.0"
-DynamicPPL = "0.37.2"
 FiniteDifferences = "0.10.8, 0.11, 0.12"
 ForwardDiff = "0.10.12 - 0.10.32, 0.10, 1"
 HypothesisTests = "0.11"
@@ -77,3 +76,6 @@ StatsBase = "0.33, 0.34"
 StatsFuns = "0.9.5, 1"
 TimerOutputs = "0.5"
 julia = "1.10"
+
+[sources]
+DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "py/setleafcontext-model"}

test/essential/container.jl

@@ -22,7 +22,7 @@ using Turing
         vi = DynamicPPL.setacc!!(vi, Turing.Inference.ProduceLogLikelihoodAccumulator())
         sampler = Sampler(PG(10))
         model = test()
-        trace = AdvancedPS.Trace(model, sampler, vi, AdvancedPS.TracedRNG())
+        trace = AdvancedPS.Trace(model, vi, AdvancedPS.TracedRNG(), false)
 
         # Make sure the backreference from taped_globals to the trace is in place.
         @test trace.model.ctask.taped_globals.other === trace
@@ -48,7 +48,7 @@ using Turing
         sampler = Sampler(PG(10))
         model = normal()
 
-        trace = AdvancedPS.Trace(model, sampler, vi, AdvancedPS.TracedRNG())
+        trace = AdvancedPS.Trace(model, vi, AdvancedPS.TracedRNG(), false)
 
         newtrace = AdvancedPS.forkr(trace)
         # Catch broken replay mechanism

test/mcmc/ess.jl

@@ -108,8 +108,12 @@ using Turing
         spl_x = Gibbs(@varname(z) => NUTS(), @varname(x) => ESS())
         spl_xy = Gibbs(@varname(z) => NUTS(), (@varname(x), @varname(y)) => ESS())
 
-        @test sample(StableRNG(23), xy(), spl_xy, num_samples).value ≈
-            sample(StableRNG(23), x12(), spl_x, num_samples).value
+        chn1 = sample(StableRNG(23), xy(), spl_xy, num_samples)
+        chn2 = sample(StableRNG(23), x12(), spl_x, num_samples)
+
+        @test mean(chn1[:z]) ≈ mean(chn2[:z]) atol = 0.05
+        @test mean(chn1[:x]) ≈ mean(chn2["x[1]"]) atol = 0.05
+        @test mean(chn1[:y]) ≈ mean(chn2["x[2]"]) atol = 0.05
     end
 end
 

test/mcmc/external_sampler.jl

@@ -45,6 +45,8 @@ using Turing.Inference: AdvancedHMC
         rng::Random.AbstractRNG,
         model::AbstractMCMC.LogDensityModel,
         sampler::MySampler;
+        # This initial_params should be an AbstractVector because the model is just a
+        # LogDensityModel, not a DynamicPPL.Model
         initial_params::AbstractVector,
         kwargs...,
     )
@@ -82,7 +84,10 @@ using Turing.Inference: AdvancedHMC
     model = test_external_sampler()
     a, b = 0.5, 0.0
 
-    chn = sample(model, externalsampler(MySampler()), 10; initial_params=[a, b])
+    # This `initial_params` should be an InitStrategy
+    chn = sample(
+        model, externalsampler(MySampler()), 10; initial_params=InitFromParams((a=a, b=b))
+    )
     @test chn isa MCMCChains.Chains
     @test all(chn[:a] .== a)
     @test all(chn[:b] .== b)
@@ -167,9 +172,7 @@ function initialize_mh_with_prior_proposal(model)
     )
 end
 
-function test_initial_params(
-    model, sampler, initial_params=DynamicPPL.VarInfo(model)[:]; kwargs...
-)
+function test_initial_params(model, sampler, initial_params=InitFromPrior(); kwargs...)
     # Execute the transition with two different RNGs and check that the resulting
     # parameter values are the same.
     rng1 = Random.MersenneTwister(42)
@@ -204,7 +207,7 @@ end
                 n_adapts=1_000,
                 discard_initial=1_000,
                 # FIXME: Remove this once we can run `test_initial_params` above.
-                initial_params=DynamicPPL.VarInfo(model)[:],
+                initial_params=InitFromPrior(),
             )
 
             @testset "inference" begin

test/mcmc/gibbs.jl

@@ -693,13 +693,9 @@ end
                 num_chains = 4
 
                 # Determine initial parameters to make comparison as fair as possible.
+                # posterior_mean returns a NamedTuple so we can plug it in directly.
                 posterior_mean = DynamicPPL.TestUtils.posterior_mean(model)
-                initial_params = DynamicPPL.TestUtils.update_values!!(
-                    DynamicPPL.VarInfo(model),
-                    posterior_mean,
-                    DynamicPPL.TestUtils.varnames(model),
-                )[:]
-                initial_params = fill(initial_params, num_chains)
+                initial_params = fill(InitFromParams(initial_params), num_chains)
 
                 # Sampler to use for Gibbs components.
                 hmc = HMC(0.1, 32)

test/mcmc/hmc.jl

@@ -177,7 +177,11 @@ using Turing
         @testset "$spl_name" for (spl_name, spl) in
                                  (("HMC", HMC(0.1, 10)), ("NUTS", NUTS()))
             chain = sample(
-                demo_norm(), spl, 5; discard_adapt=false, initial_params=(x=init_x,)
+                demo_norm(),
+                spl,
+                5;
+                discard_adapt=false,
+                initial_params=InitFromParams((x=init_x,)),
             )
             @test chain[:x][1] == init_x
             chain = sample(
@@ -187,7 +191,7 @@ using Turing
                 5,
                 5;
                 discard_adapt=false,
-                initial_params=(fill((x=init_x,), 5)),
+                initial_params=(fill(InitFromParams((x=init_x,)), 5)),
             )
             @test all(chain[:x][1, :] .== init_x)
         end
@@ -202,12 +206,11 @@ using Turing
             end
         end
 
-        @test_logs (
-            :warn,
-            "failed to find valid initial parameters in 10 tries; consider providing explicit initial parameters using the `initial_params` keyword",
-        ) (:info,) match_mode = :any begin
-            sample(demo_warn_initial_params(), NUTS(), 5)
-        end
+        # verbose=false to suppress the initial step size notification, which messes with
+        # the test
+        @test_logs (:warn, r"consider providing a different initialisation strategy") sample(
+            demo_warn_initial_params(), NUTS(), 5; verbose=false
+        )
     end
 
     @testset "error for impossible model" begin
@@ -253,7 +256,8 @@ using Turing
         model = buggy_model()
         num_samples = 1_000
 
-        chain = sample(model, NUTS(), num_samples; initial_params=[0.5, 1.75, 1.0])
+        initial_params = InitFromParams((lb=0.5, ub=1.75, x=1.0))
+        chain = sample(model, NUTS(), num_samples; initial_params=initial_params)
         chain_prior = sample(model, Prior(), num_samples)
 
         # Extract the `x` like this because running `generated_quantities` was how
@@ -275,7 +279,11 @@ using Turing
             # Construct a HMC state by taking a single step
             spl = Sampler(alg)
             hmc_state = DynamicPPL.initialstep(
-                Random.default_rng(), gdemo_default, spl, DynamicPPL.VarInfo(gdemo_default)
+                Random.default_rng(),
+                gdemo_default,
+                spl,
+                DynamicPPL.VarInfo(gdemo_default);
+                initial_params=InitFromUniform(),
             )[2]
             # Check that we can obtain the current step size
             @test Turing.Inference.getstepsize(spl, hmc_state) isa Float64

test/mcmc/mh.jl

@@ -49,7 +49,7 @@ GKernel(var) = (x) -> Normal(x, sqrt.(var))
         # Set the initial parameters, because if we get unlucky with the initial state,
         # these chains are too short to converge to reasonable numbers.
         discard_initial = 1_000
-        initial_params = [1.0, 1.0]
+        initial_params = InitFromParams((s=1.0, m=1.0))
 
         @testset "gdemo_default" begin
             alg = MH()
@@ -81,13 +81,16 @@ GKernel(var) = (x) -> Normal(x, sqrt.(var))
                 @varname(mu1) => MH((:mu1, GKernel(1))),
                 @varname(mu2) => MH((:mu2, GKernel(1))),
             )
+            initial_params = InitFromParams((
+                mu1=1.0, mu2=1.0, z1=0.0, z2=0.0, z3=1.0, z4=1.0
+            ))
             chain = sample(
                 StableRNG(seed),
                 MoGtest_default,
                 gibbs,
                 500;
                 discard_initial=100,
-                initial_params=[1.0, 1.0, 0.0, 0.0, 1.0, 4.0],
+                initial_params=initial_params,
             )
             check_MoGtest_default(chain; atol=0.2)
         end