Fix biased PG/CSMC posteriors: make delete_retained! mutate

CLAUDE.md

@@ -98,3 +98,4 @@ User-facing functions accept `initial_params` as a convenience. `_convert_initia
 
   - Non-breaking changes target `main`; breaking changes target the `breaking` branch.
   - Julia ≥ 1.10.8 required (see `[compat]` in `Project.toml`).
+  - `HISTORY.md`: one line for a bugfix or internal change. Only a breaking change or new feature earns more, and only what a user needs to act on it: what broke, and the old → new form. The mechanism and any measurements belong in the commit and the PR.

HISTORY.md

@@ -1,3 +1,7 @@
+# 0.46.1
+
+Fixed a bug, present since v0.41.0, that biased `PG` / `CSMC` posteriors, whether sampled on their own or as a Gibbs component.
+
 # 0.46.0
 
 ## Breaking changes

Project.toml

@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.46.0"
+version = "0.46.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/mcmc/particle_mcmc.jl

@@ -67,8 +67,10 @@ function AdvancedPS.delete_retained!(trace::TracedModel)
     # and choose the reference particle as the trajectory to carry on from.
     # In such a case, we need to ensure that when we continue sampling (i.e.
     # the next time we hit tilde_assume!!), we don't use the values in the 
-    # reference particle but rather sample new values.
-    return TracedModel(trace.model, trace.varinfo, true, trace.fargs, trace.kwargs)
+    # reference particle but rather sample new values. This has to mutate:
+    # `AdvancedPS.fork` calls it for its side effect and discards the return value.
+    trace.resample = true
+    return trace
 end
 
 function AdvancedPS.reset_model(trace::TracedModel)

test/mcmc/particle_mcmc.jl

@@ -3,7 +3,7 @@ module ParticleMCMCTests
 using ..Models: gdemo_default
 using ..SamplerTestUtils: test_chain_logp_metadata
 using AdvancedPS: ResampleWithESSThreshold, resample_systematic, resample_multinomial
-using Distributions: Bernoulli, Beta, Gamma, Normal, sample
+using Distributions: Bernoulli, Beta, Categorical, Gamma, Normal, sample
 using FlexiChains: VNChain
 using Random: Random
 using StableRNGs: StableRNG
@@ -168,6 +168,48 @@ end
         @test length(unique(c[:s])) == 1
     end
 
+    @testset "conditional sweeps target the exact posterior" begin
+        # A reference particle that is not exactly the retained trajectory shows up here. The
+        # chain's stay probability is the other Gibbs component, so every `z[t]` is
+        # re-conditioned when `i` moves, and the observations are sharp enough to make the
+        # weights uneven; either is enough to bias the marginals. Measured over four seeds, a
+        # correct sweep keeps the mean error under 0.005, descendants that copy the reference
+        # rather than branching off it give 0.020 to 0.028, and a reference rebuilt by
+        # replaying random numbers instead of reusing values gives 0.017 to 0.022.
+        #
+        # Enumerating all `2 * 2^8` configurations and weighting them by the model's own log
+        # density keeps the target out of the hands of a reimplementation.
+        means, sd, stay = (-1.0, 1.0), 0.8, (0.35, 0.65)
+        @model function switching(y)
+            i ~ Categorical(2)
+            z = Vector{Int}(undef, length(y))
+            z[1] ~ Categorical([0.5, 0.5])
+            y[1] ~ Normal(means[z[1]], sd)
+            for t in 2:length(y)
+                p = stay[i]
+                z[t] ~ Categorical(z[t - 1] == 1 ? [p, 1 - p] : [1 - p, p])
+                y[t] ~ Normal(means[z[t]], sd)
+            end
+        end
+        y = [-0.9163, -2.4106, -2.1881, 0.3716, 1.3404, -1.2046, -1.8294, -0.3521]
+        model = switching(y)
+
+        confs = [
+            (i, collect(z)) for i in 1:2 for z in Iterators.product(fill(1:2, length(y))...)
+        ]
+        w = exp.([logjoint(model, (; i=i, z=z)) for (i, z) in confs])
+        w ./= sum(w)
+        exact = [
+            sum(w[k] * (confs[k][2][t] - 1) for k in eachindex(w)) for t in eachindex(y)
+        ]
+
+        alg = Gibbs(@varname(i) => MH(), @varname(z) => CSMC(8))
+        chn = sample(StableRNG(468), model, alg, 6_000)
+        zs = stack(collect(z) for z in chn[@varname(z)])
+        marginals = [mean(view(zs, t, :) .== 2) for t in eachindex(y)]
+        @test mean(abs, marginals .- exact) < 0.01
+    end
+
     @testset "addlogprob leads to reweighting" begin
         # Make sure that PG takes @addlogprob! into account. It didn't use to:
         # https://github.com/TuringLang/Turing.jl/issues/1996