Merge branch 'main' into breaking

penelopeysm · web-flow · commit 1e1cd94382c2 · 2025-09-15T16:55:37.000+01:00
diff --git a/.github/workflows/Benchmarking.yml b/.github/workflows/Benchmarking.yml
@@ -9,7 +9,7 @@ jobs:
 
     steps:
       - name: Checkout Repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v5
         with:
           ref: ${{ github.event.pull_request.head.sha }}
 
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -56,7 +56,7 @@ jobs:
           - Group2
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
 
       - uses: julia-actions/setup-julia@v2
         with:
diff --git a/.github/workflows/DocTest.yml b/.github/workflows/DocTest.yml
@@ -26,7 +26,7 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
 
       - uses: julia-actions/setup-julia@v2
         with:
diff --git a/.github/workflows/DocsNav.yml b/.github/workflows/DocsNav.yml
@@ -17,7 +17,7 @@ jobs:
 
     steps:
       - name: Checkout gh-pages branch
-        uses: actions/checkout@v4
+        uses: actions/checkout@v5
         with:
           ref: gh-pages
 
diff --git a/.github/workflows/DocsPreviewCleanup.yml b/.github/workflows/DocsPreviewCleanup.yml
@@ -9,7 +9,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout gh-pages branch
-        uses: actions/checkout@v4
+        uses: actions/checkout@v5
         with:
           ref: gh-pages
       - name: Delete preview and history + push changes
diff --git a/.github/workflows/Enzyme.yml b/.github/workflows/Enzyme.yml
@@ -20,7 +20,7 @@ jobs:
   enzyme:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
 
       - uses: julia-actions/setup-julia@v2
         with:
diff --git a/.github/workflows/IntegrationTest.yml b/.github/workflows/IntegrationTest.yml
@@ -21,13 +21,13 @@ jobs:
           - {user: TuringLang, repo: Turing.jl}
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
       - uses: julia-actions/setup-julia@v2
         with:
           version: 1
       - uses: julia-actions/julia-buildpkg@v1
       - name: Clone Downstream
-        uses: actions/checkout@v4
+        uses: actions/checkout@v5
         with:
           repository: ${{ matrix.package.user }}/${{ matrix.package.repo }}
           path: downstream
diff --git a/.github/workflows/JuliaPre.yml b/.github/workflows/JuliaPre.yml
@@ -23,7 +23,7 @@ jobs:
           - Group2
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
       - uses: julia-actions/setup-julia@v2
         with:
           version: 'pre' # pre-release
diff --git a/HISTORY.md b/HISTORY.md
@@ -5,7 +5,11 @@
 The `varname_leaves` and `varname_and_value_leaves` functions have been moved to AbstractPPL.jl.
 Their behaviour is otherwise identical.
 
-[...]
+## 0.37.2
+
+Make the `resume_from` keyword work for multiple-chain (parallel) sampling as well.
+Prior to this version, it was silently ignored.
+Note that to get the correct behaviour you also need to have a recent version of MCMCChains (v7.2.1).
 
 ## 0.37.1
 
diff --git a/src/sampler.jl b/src/sampler.jl
@@ -95,6 +95,23 @@ function AbstractMCMC.sample(
     )
 end
 
+function AbstractMCMC.sample(
+    rng::Random.AbstractRNG,
+    model::Model,
+    sampler::Sampler,
+    parallel::AbstractMCMC.AbstractMCMCEnsemble,
+    N::Integer,
+    nchains::Integer;
+    chain_type=default_chain_type(sampler),
+    resume_from=nothing,
+    initial_state=loadstate(resume_from),
+    kwargs...,
+)
+    return AbstractMCMC.mcmcsample(
+        rng, model, sampler, parallel, N, nchains; chain_type, initial_state, kwargs...
+    )
+end
+
 # initial step: general interface for resuming and
 function AbstractMCMC.step(
     rng::Random.AbstractRNG, model::Model, spl::Sampler; initial_params=nothing, kwargs...
diff --git a/test/Project.toml b/test/Project.toml
@@ -43,7 +43,7 @@ EnzymeCore = "0.6 - 0.8"
 ForwardDiff = "0.10.12, 1"
 JET = "0.9, 0.10"
 LogDensityProblems = "2"
-MCMCChains = "6.0.4, 7"
+MCMCChains = "7.2.1"
 MacroTools = "0.5.6"
 OrderedCollections = "1"
 ReverseDiff = "1"
diff --git a/test/sampler.jl b/test/sampler.jl
@@ -1,4 +1,105 @@
 @testset "sampler.jl" begin
+    @testset "initial_state and resume_from kwargs" begin
+        # Model is unused, but has to be a DynamicPPL.Model otherwise we won't hit our
+        # overloaded method.
+        @model f() = x ~ Normal()
+        model = f()
+        # This sampler just returns the state it was given as its 'sample'
+        struct S <: AbstractMCMC.AbstractSampler end
+        function AbstractMCMC.step(
+            rng::Random.AbstractRNG,
+            model::Model,
+            sampler::Sampler{<:S},
+            state=nothing;
+            kwargs...,
+        )
+            if state === nothing
+                s = rand()
+                return s, s
+            else
+                return state, state
+            end
+        end
+        spl = Sampler(S())
+
+        function AbstractMCMC.bundle_samples(
+            samples::Vector{Float64},
+            model::Model,
+            sampler::Sampler{<:S},
+            state,
+            chain_type::Type{MCMCChains.Chains};
+            kwargs...,
+        )
+            return MCMCChains.Chains(samples, [:x]; info=(samplerstate=state,))
+        end
+
+        N_iters, N_chains = 10, 3
+
+        @testset "single-chain sampling" begin
+            chn = sample(model, spl, N_iters; progress=false, chain_type=MCMCChains.Chains)
+            initial_value = chn[:x][1]
+            @test all(chn[:x] .== initial_value) # sanity check
+            # using `initial_state`
+            chn2 = sample(
+                model,
+                spl,
+                N_iters;
+                progress=false,
+                initial_state=chn.info.samplerstate,
+                chain_type=MCMCChains.Chains,
+            )
+            @test all(chn2[:x] .== initial_value)
+            # using `resume_from`
+            chn3 = sample(
+                model,
+                spl,
+                N_iters;
+                progress=false,
+                resume_from=chn,
+                chain_type=MCMCChains.Chains,
+            )
+            @test all(chn3[:x] .== initial_value)
+        end
+
+        @testset "multiple-chain sampling" begin
+            chn = sample(
+                model,
+                spl,
+                MCMCThreads(),
+                N_iters,
+                N_chains;
+                progress=false,
+                chain_type=MCMCChains.Chains,
+            )
+            initial_value = chn[:x][1, :]
+            @test all(i -> chn[:x][i, :] == initial_value, 1:N_iters) # sanity check
+            # using `initial_state`
+            chn2 = sample(
+                model,
+                spl,
+                MCMCThreads(),
+                N_iters,
+                N_chains;
+                progress=false,
+                initial_state=chn.info.samplerstate,
+                chain_type=MCMCChains.Chains,
+            )
+            @test all(i -> chn2[:x][i, :] == initial_value, 1:N_iters)
+            # using `resume_from`
+            chn3 = sample(
+                model,
+                spl,
+                MCMCThreads(),
+                N_iters,
+                N_chains;
+                progress=false,
+                resume_from=chn,
+                chain_type=MCMCChains.Chains,
+            )
+            @test all(i -> chn3[:x][i, :] == initial_value, 1:N_iters)
+        end
+    end
+
     @testset "SampleFromPrior and SampleUniform" begin
         @model function gdemo(x, y)
             s ~ InverseGamma(2, 3)
