From 7b2869f8fe55255689917b244f957079216c3e68 Mon Sep 17 00:00:00 2001
From: Markus Hauru <markus@mhauru.org>
Date: Sun, 23 Jun 2024 21:23:23 +0200
Subject: [PATCH 1/4] Add Base.get method for ModeResult (#2269)

* Add Base.get method for ModeResult

* Make get(::ModeResult, itr) work for any iterator

* Fix array type in get(::ModeResult, ...)
---
 src/optimisation/Optimisation.jl  | 31 +++++++++++++++++++++++++++++++
 test/optimisation/Optimisation.jl | 27 ++++++++++++++++++++++++++-
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/src/optimisation/Optimisation.jl b/src/optimisation/Optimisation.jl
index 682e664a66..82bda7da79 100644
--- a/src/optimisation/Optimisation.jl
+++ b/src/optimisation/Optimisation.jl
@@ -273,6 +273,37 @@ StatsBase.params(m::ModeResult) = StatsBase.coefnames(m)
 StatsBase.vcov(m::ModeResult) = inv(StatsBase.informationmatrix(m))
 StatsBase.loglikelihood(m::ModeResult) = m.lp
 
+"""
+    Base.get(m::ModeResult, var_symbol::Symbol)
+    Base.get(m::ModeResult, var_symbols)
+
+Return the values of all the variables with the symbol(s) `var_symbol` in the mode result
+`m`. The return value is a `NamedTuple` with `var_symbols` as the key(s). The second
+argument should be either a `Symbol` or an iterator of `Symbol`s.
+"""
+function Base.get(m::ModeResult, var_symbols)
+    log_density = m.f
+    # Get all the variable names in the model. This is the same as the list of keys in
+    # m.values, but they are more convenient to filter when they are VarNames rather than
+    # Symbols.
+    varnames = collect(
+        map(first, Turing.Inference.getparams(log_density.model, log_density.varinfo))
+    )
+    # For each symbol s in var_symbols, pick all the values from m.values for which the
+    # variable name has that symbol.
+    et = eltype(m.values)
+    value_vectors = Vector{et}[]
+    for s in var_symbols
+        push!(
+            value_vectors,
+            [m.values[Symbol(vn)] for vn in varnames if DynamicPPL.getsym(vn) == s],
+        )
+    end
+    return (; zip(var_symbols, value_vectors)...)
+end
+
+Base.get(m::ModeResult, var_symbol::Symbol) = get(m, (var_symbol,))
+
 """
     ModeResult(log_density::OptimLogDensity, solution::SciMLBase.OptimizationSolution)
 
diff --git a/test/optimisation/Optimisation.jl b/test/optimisation/Optimisation.jl
index 5e6144e578..76d3a940d6 100644
--- a/test/optimisation/Optimisation.jl
+++ b/test/optimisation/Optimisation.jl
@@ -4,13 +4,14 @@ using ..Models: gdemo, gdemo_default
 using Distributions
 using Distributions.FillArrays: Zeros
 using DynamicPPL: DynamicPPL
-using LinearAlgebra: I
+using LinearAlgebra: Diagonal, I
 using Random: Random
 using Optimization
 using Optimization: Optimization
 using OptimizationBBO: OptimizationBBO
 using OptimizationNLopt: OptimizationNLopt
 using OptimizationOptimJL: OptimizationOptimJL
+using ReverseDiff: ReverseDiff
 using StatsBase: StatsBase
 using StatsBase: coef, coefnames, coeftable, informationmatrix, stderror, vcov
 using Test: @test, @testset, @test_throws
@@ -591,6 +592,30 @@ using Turing
         @test result.values[:x] ≈ 0 atol = 1e-1
         @test result.values[:y] ≈ 100 atol = 1e-1
     end
+
+    @testset "get ModeResult" begin
+        @model function demo_model(N)
+            half_N = N ÷ 2
+            a ~ arraydist(LogNormal.(fill(0, half_N), 1))
+            b ~ arraydist(LogNormal.(fill(0, N - half_N), 1))
+            covariance_matrix = Diagonal(vcat(a, b))
+            x ~ MvNormal(covariance_matrix)
+            return nothing
+        end
+
+        N = 12
+        m = demo_model(N) | (x=randn(N),)
+        result = maximum_a_posteriori(m)
+        get_a = get(result, :a)
+        get_b = get(result, :b)
+        get_ab = get(result, [:a, :b])
+        @assert keys(get_a) == (:a,)
+        @assert keys(get_b) == (:b,)
+        @assert keys(get_ab) == (:a, :b)
+        @assert get_b[:b] == get_ab[:b]
+        @assert vcat(get_a[:a], get_b[:b]) == result.values.array
+        @assert get(result, :c) == (; :c => Array{Float64}[])
+    end
 end
 
 end

From a0db647693f69e9c50d0eb60a4c425e83fd61b3e Mon Sep 17 00:00:00 2001
From: Markus Hauru <markus@mhauru.org>
Date: Tue, 25 Jun 2024 17:52:59 +0200
Subject: [PATCH 2/4] Use new style kwarg constructor for AutoReverseDiff
 (#2273)

---
 benchmarks/benchmarks_suite.jl |  2 +-
 test/essential/ad.jl           | 12 +++++++-----
 test/mcmc/Inference.jl         |  2 +-
 test/mcmc/gibbs.jl             |  2 +-
 test/mcmc/gibbs_conditional.jl |  2 +-
 test/mcmc/hmc.jl               |  2 +-
 test/mcmc/sghmc.jl             |  4 ++--
 7 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/benchmarks/benchmarks_suite.jl b/benchmarks/benchmarks_suite.jl
index 7f8db0980c..9d7bbfa554 100644
--- a/benchmarks/benchmarks_suite.jl
+++ b/benchmarks/benchmarks_suite.jl
@@ -84,5 +84,5 @@ BenchmarkSuite["mnormal"]["forwarddiff"] = @benchmarkable sample(
 
 # ReverseDiff
 BenchmarkSuite["mnormal"]["reversediff"] = @benchmarkable sample(
-    $(mdemo(d, 1)), $(HMC(0.1, 5; adtype=AutoReverseDiff(false))), 5000
+    $(mdemo(d, 1)), $(HMC(0.1, 5; adtype=AutoReverseDiff(; compile=false))), 5000
 )
diff --git a/test/essential/ad.jl b/test/essential/ad.jl
index 6583ed911f..943b4eadc6 100644
--- a/test/essential/ad.jl
+++ b/test/essential/ad.jl
@@ -154,22 +154,22 @@ end
             return theta ~ Dirichlet(1 ./ fill(4, 4))
         end
         sample(dir(), HMC(0.01, 1; adtype=AutoZygote()), 1000)
-        sample(dir(), HMC(0.01, 1; adtype=AutoReverseDiff(false)), 1000)
-        sample(dir(), HMC(0.01, 1; adtype=AutoReverseDiff(true)), 1000)
+        sample(dir(), HMC(0.01, 1; adtype=AutoReverseDiff(; compile=false)), 1000)
+        sample(dir(), HMC(0.01, 1; adtype=AutoReverseDiff(; compile=true)), 1000)
     end
     @testset "PDMatDistribution AD" begin
         @model function wishart()
             return theta ~ Wishart(4, Matrix{Float64}(I, 4, 4))
         end
 
-        sample(wishart(), HMC(0.01, 1; adtype=AutoReverseDiff(false)), 1000)
+        sample(wishart(), HMC(0.01, 1; adtype=AutoReverseDiff(; compile=false)), 1000)
         sample(wishart(), HMC(0.01, 1; adtype=AutoZygote()), 1000)
 
         @model function invwishart()
             return theta ~ InverseWishart(4, Matrix{Float64}(I, 4, 4))
         end
 
-        sample(invwishart(), HMC(0.01, 1; adtype=AutoReverseDiff(false)), 1000)
+        sample(invwishart(), HMC(0.01, 1; adtype=AutoReverseDiff(; compile=false)), 1000)
         sample(invwishart(), HMC(0.01, 1; adtype=AutoZygote()), 1000)
     end
     @testset "Hessian test" begin
@@ -231,7 +231,9 @@ end
         for i in 1:5
             d = Normal(0.0, i)
             data = rand(d, N)
-            chn = sample(demo(data), NUTS(0.65; adtype=AutoReverseDiff(true)), 1000)
+            chn = sample(
+                demo(data), NUTS(0.65; adtype=AutoReverseDiff(; compile=true)), 1000
+            )
             @test mean(Array(chn[:sigma])) ≈ std(data) atol = 0.5
         end
     end
diff --git a/test/mcmc/Inference.jl b/test/mcmc/Inference.jl
index f7601b2e15..f1ad4a621d 100644
--- a/test/mcmc/Inference.jl
+++ b/test/mcmc/Inference.jl
@@ -14,7 +14,7 @@ import ReverseDiff
 using Test: @test, @test_throws, @testset
 using Turing
 
-@testset "Testing inference.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(false))
+@testset "Testing inference.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false))
     # Only test threading if 1.3+.
     if VERSION > v"1.2"
         @testset "threaded sampling" begin
diff --git a/test/mcmc/gibbs.jl b/test/mcmc/gibbs.jl
index 4d2053c14e..5159e022b9 100644
--- a/test/mcmc/gibbs.jl
+++ b/test/mcmc/gibbs.jl
@@ -13,7 +13,7 @@ using Turing: Inference
 using Turing.RandomMeasures: ChineseRestaurantProcess, DirichletProcess
 
 @testset "Testing gibbs.jl with $adbackend" for adbackend in (
-    AutoForwardDiff(; chunksize=0), AutoReverseDiff(false)
+    AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false)
 )
     @testset "gibbs constructor" begin
         N = 500
diff --git a/test/mcmc/gibbs_conditional.jl b/test/mcmc/gibbs_conditional.jl
index 3ba2fdbedd..6110bbdb18 100644
--- a/test/mcmc/gibbs_conditional.jl
+++ b/test/mcmc/gibbs_conditional.jl
@@ -15,7 +15,7 @@ using Test: @test, @testset
 using Turing
 
 @testset "Testing gibbs conditionals.jl with $adbackend" for adbackend in (
-    AutoForwardDiff(; chunksize=0), AutoReverseDiff(false)
+    AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false)
 )
     Random.seed!(1000)
     rng = StableRNG(123)
diff --git a/test/mcmc/hmc.jl b/test/mcmc/hmc.jl
index 89dc4a3149..c802f8d9e7 100644
--- a/test/mcmc/hmc.jl
+++ b/test/mcmc/hmc.jl
@@ -16,7 +16,7 @@ using StatsFuns: logistic
 using Test: @test, @test_logs, @testset
 using Turing
 
-@testset "Testing hmc.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(false))
+@testset "Testing hmc.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false))
     # Set a seed
     rng = StableRNG(123)
     @testset "constrained bounded" begin
diff --git a/test/mcmc/sghmc.jl b/test/mcmc/sghmc.jl
index f65ca1c372..a5829eb180 100644
--- a/test/mcmc/sghmc.jl
+++ b/test/mcmc/sghmc.jl
@@ -10,7 +10,7 @@ using StableRNGs: StableRNG
 using Test: @test, @testset
 using Turing
 
-@testset "Testing sghmc.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(false))
+@testset "Testing sghmc.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false))
     @testset "sghmc constructor" begin
         alg = SGHMC(; learning_rate=0.01, momentum_decay=0.1, adtype=adbackend)
         @test alg isa SGHMC
@@ -36,7 +36,7 @@ using Turing
     end
 end
 
-@testset "Testing sgld.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(false))
+@testset "Testing sgld.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false))
     @testset "sgld constructor" begin
         alg = SGLD(; stepsize=PolynomialStepsize(0.25), adtype=adbackend)
         @test alg isa SGLD

From 927abcd979eb49adc76cdfea13db4929fde41cd7 Mon Sep 17 00:00:00 2001
From: Markus Hauru <markus@mhauru.org>
Date: Wed, 26 Jun 2024 15:33:53 +0200
Subject: [PATCH 3/4] Bump DynamicPPL to v0.28 (#2276)

* Bump DynamicPPL to v0.28

* Also bump DPPL to v0.28 in tests

* Qualify use of getcontext
---
 Project.toml                     | 2 +-
 src/mcmc/Inference.jl            | 2 +-
 src/mcmc/mh.jl                   | 2 +-
 src/optimisation/Optimisation.jl | 4 +++-
 test/Project.toml                | 2 +-
 5 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/Project.toml b/Project.toml
index 3ee30eb953..5bf3336f31 100644
--- a/Project.toml
+++ b/Project.toml
@@ -63,7 +63,7 @@ Distributions = "0.23.3, 0.24, 0.25"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.8, 0.9"
 DynamicHMC = "3.4"
-DynamicPPL = "0.27.1"
+DynamicPPL = "0.28"
 Compat = "4.15.0"
 EllipticalSliceSampling = "0.5, 1, 2"
 ForwardDiff = "0.10.3"
diff --git a/src/mcmc/Inference.jl b/src/mcmc/Inference.jl
index 7a0c541da2..4f1f423c1f 100644
--- a/src/mcmc/Inference.jl
+++ b/src/mcmc/Inference.jl
@@ -149,7 +149,7 @@ getADType(::DynamicPPL.IsParent, ctx::DynamicPPL.AbstractContext) = getADType(Dy
 getADType(alg::Hamiltonian) = alg.adtype
 
 function LogDensityProblemsAD.ADgradient(ℓ::DynamicPPL.LogDensityFunction)
-    return LogDensityProblemsAD.ADgradient(getADType(ℓ.context), ℓ)
+    return LogDensityProblemsAD.ADgradient(getADType(DynamicPPL.getcontext(ℓ)), ℓ)
 end
 
 function LogDensityProblems.logdensity(
diff --git a/src/mcmc/mh.jl b/src/mcmc/mh.jl
index 9b7f3ff090..ffc064eb12 100644
--- a/src/mcmc/mh.jl
+++ b/src/mcmc/mh.jl
@@ -271,7 +271,7 @@ function LogDensityProblems.logdensity(f::MHLogDensityFunction, x::NamedTuple)
 
     x_old, lj_old = vi[sampler], getlogp(vi)
     set_namedtuple!(vi, x)
-    vi_new = last(DynamicPPL.evaluate!!(f.model, vi, f.context))
+    vi_new = last(DynamicPPL.evaluate!!(f.model, vi, DynamicPPL.getcontext(f)))
     lj = getlogp(vi_new)
 
     # Reset old `vi`.
diff --git a/src/optimisation/Optimisation.jl b/src/optimisation/Optimisation.jl
index 82bda7da79..e069897920 100644
--- a/src/optimisation/Optimisation.jl
+++ b/src/optimisation/Optimisation.jl
@@ -142,7 +142,9 @@ required by Optimization.jl.
 """
 function (f::OptimLogDensity)(z::AbstractVector)
     varinfo = DynamicPPL.unflatten(f.varinfo, z)
-    return -DynamicPPL.getlogp(last(DynamicPPL.evaluate!!(f.model, varinfo, f.context)))
+    return -DynamicPPL.getlogp(
+        last(DynamicPPL.evaluate!!(f.model, varinfo, DynamicPPL.getcontext(f)))
+    )
 end
 
 (f::OptimLogDensity)(z, _) = f(z)
diff --git a/test/Project.toml b/test/Project.toml
index a7538afbfc..67292d2af5 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -45,7 +45,7 @@ Clustering = "0.14, 0.15"
 Distributions = "0.25"
 DistributionsAD = "0.6.3"
 DynamicHMC = "2.1.6, 3.0"
-DynamicPPL = "0.27"
+DynamicPPL = "0.28"
 FiniteDifferences = "0.10.8, 0.11, 0.12"
 ForwardDiff = "0.10.12 - 0.10.32, 0.10"
 HypothesisTests = "0.11"

From cbd5d7949483180f2aa8d37d3f5fccd6fcd81369 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 28 Jun 2024 22:05:26 +0100
Subject: [PATCH 4/4] Check model by default (#2218)

* check model by default

* removed check_model kwargs from non-leaf method

* uncomment tests

* removed incorrect usage of check_model

* fixed IS tests

* relax gibbs tests

* Give the MH inference tests some burn-in to see if that can help

* made the MH inference tests a bit more predictable by providing
initial params

* Relaxed HMC tests a bit

---------

Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com>
---
 src/mcmc/Inference.jl  | 13 +++++++++++++
 test/mcmc/Inference.jl | 22 ++++++++++++++++++++++
 test/mcmc/gibbs.jl     |  4 +++-
 test/mcmc/hmc.jl       |  2 +-
 test/mcmc/is.jl        |  2 +-
 test/mcmc/mh.jl        | 21 ++++++++++++++++-----
 6 files changed, 56 insertions(+), 8 deletions(-)

diff --git a/src/mcmc/Inference.jl b/src/mcmc/Inference.jl
index 4f1f423c1f..cc43451eab 100644
--- a/src/mcmc/Inference.jl
+++ b/src/mcmc/Inference.jl
@@ -238,6 +238,15 @@ DynamicPPL.getlogp(t::Transition) = t.lp
 # Metadata of VarInfo object
 metadata(vi::AbstractVarInfo) = (lp = getlogp(vi),)
 
+# TODO: Implement additional checks for certain samplers, e.g.
+# HMC not supporting discrete parameters.
+function _check_model(model::DynamicPPL.Model)
+    return DynamicPPL.check_model(model; error_on_failure=true)
+end
+function _check_model(model::DynamicPPL.Model, alg::InferenceAlgorithm)
+    return _check_model(model)
+end
+
 #########################################
 # Default definitions for the interface #
 #########################################
@@ -256,8 +265,10 @@ function AbstractMCMC.sample(
     model::AbstractModel,
     alg::InferenceAlgorithm,
     N::Integer;
+    check_model::Bool=true,
     kwargs...
 )
+    check_model && _check_model(model, alg)
     return AbstractMCMC.sample(rng, model, Sampler(alg, model), N; kwargs...)
 end
 
@@ -280,8 +291,10 @@ function AbstractMCMC.sample(
     ensemble::AbstractMCMC.AbstractMCMCEnsemble,
     N::Integer,
     n_chains::Integer;
+    check_model::Bool=true,
     kwargs...
 )
+    check_model && _check_model(model, alg)
     return AbstractMCMC.sample(rng, model, Sampler(alg, model), ensemble, N, n_chains;
                                kwargs...)
 end
diff --git a/test/mcmc/Inference.jl b/test/mcmc/Inference.jl
index f1ad4a621d..64a5e95df0 100644
--- a/test/mcmc/Inference.jl
+++ b/test/mcmc/Inference.jl
@@ -559,6 +559,28 @@ using Turing
         @test all(xs[:, 1] .=== [1, missing, 3])
         @test all(xs[:, 2] .=== [missing, 2, 4])
     end
+
+    @testset "check model" begin
+        @model function demo_repeated_varname()
+            x ~ Normal(0, 1)
+            x ~ Normal(x, 1)
+        end
+
+        @test_throws ErrorException sample(
+            demo_repeated_varname(), NUTS(), 1000; check_model=true
+        )
+        # Make sure that disabling the check also works.
+        @test (sample(
+            demo_repeated_varname(), Prior(), 10; check_model=false
+        ); true)
+
+        @model function demo_incorrect_missing(y)
+            y[1:1] ~ MvNormal(zeros(1), 1)
+        end
+        @test_throws ErrorException sample(
+            demo_incorrect_missing([missing]), NUTS(), 1000; check_model=true
+        )
+    end
 end
 
 end
diff --git a/test/mcmc/gibbs.jl b/test/mcmc/gibbs.jl
index 5159e022b9..f30dc0f777 100644
--- a/test/mcmc/gibbs.jl
+++ b/test/mcmc/gibbs.jl
@@ -50,7 +50,9 @@ using Turing.RandomMeasures: ChineseRestaurantProcess, DirichletProcess
         Random.seed!(100)
         alg = Gibbs(CSMC(15, :s), HMC(0.2, 4, :m; adtype=adbackend))
         chain = sample(gdemo(1.5, 2.0), alg, 10_000)
-        check_numerical(chain, [:s, :m], [49 / 24, 7 / 6]; atol=0.15)
+        check_numerical(chain, [:m], [7 / 6]; atol=0.15)
+        # Be more relaxed with the tolerance of the variance.
+        check_numerical(chain, [:s], [49 / 24]; atol=0.35)
 
         Random.seed!(100)
 
diff --git a/test/mcmc/hmc.jl b/test/mcmc/hmc.jl
index c802f8d9e7..968f24d7b7 100644
--- a/test/mcmc/hmc.jl
+++ b/test/mcmc/hmc.jl
@@ -319,7 +319,7 @@ using Turing
 
         # The discrepancies in the chains are in the tails, so we can't just compare the mean, etc.
         # KS will compare the empirical CDFs, which seems like a reasonable thing to do here.
-        @test pvalue(ApproximateTwoSampleKSTest(vec(results), vec(results_prior))) > 0.01
+        @test pvalue(ApproximateTwoSampleKSTest(vec(results), vec(results_prior))) > 0.001
     end
 end
 
diff --git a/test/mcmc/is.jl b/test/mcmc/is.jl
index bd3186cd93..47b20cc736 100644
--- a/test/mcmc/is.jl
+++ b/test/mcmc/is.jl
@@ -46,7 +46,7 @@ using Turing
         ref = reference(n)
 
         Random.seed!(seed)
-        chain = sample(model, alg, n)
+        chain = sample(model, alg, n; check_model=false)
         sampled = get(chain, [:a, :b, :lp])
 
         @test vec(sampled.a) == ref.as
diff --git a/test/mcmc/mh.jl b/test/mcmc/mh.jl
index 0e3cc91f6f..a01d3dc253 100644
--- a/test/mcmc/mh.jl
+++ b/test/mcmc/mh.jl
@@ -44,21 +44,26 @@ GKernel(var) = (x) -> Normal(x, sqrt.(var))
         # c6 = sample(gdemo_default, s6, N)
     end
     @testset "mh inference" begin
+        # 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 = 1000
+        initial_params = [1.0, 1.0]
+
         Random.seed!(125)
         alg = MH()
-        chain = sample(gdemo_default, alg, 10_000)
+        chain = sample(gdemo_default, alg, 10_000; discard_initial, initial_params)
         check_gdemo(chain; atol=0.1)
 
         Random.seed!(125)
         # MH with Gaussian proposal
         alg = MH((:s, InverseGamma(2, 3)), (:m, GKernel(1.0)))
-        chain = sample(gdemo_default, alg, 10_000)
+        chain = sample(gdemo_default, alg, 10_000; discard_initial, initial_params)
         check_gdemo(chain; atol=0.1)
 
         Random.seed!(125)
         # MH within Gibbs
         alg = Gibbs(MH(:m), MH(:s))
-        chain = sample(gdemo_default, alg, 10_000)
+        chain = sample(gdemo_default, alg, 10_000; discard_initial, initial_params)
         check_gdemo(chain; atol=0.1)
 
         Random.seed!(125)
@@ -66,8 +71,14 @@ GKernel(var) = (x) -> Normal(x, sqrt.(var))
         gibbs = Gibbs(
             CSMC(15, :z1, :z2, :z3, :z4), MH((:mu1, GKernel(1)), (:mu2, GKernel(1)))
         )
-        chain = sample(MoGtest_default, gibbs, 500)
-        check_MoGtest_default(chain; atol=0.15)
+        chain = sample(
+            MoGtest_default,
+            gibbs,
+            500;
+            discard_initial=100,
+            initial_params=[1.0, 1.0, 0.0, 0.0, 1.0, 4.0],
+        )
+        check_MoGtest_default(chain; atol=0.2)
     end
 
     # Test MH shape passing.