Skip to content

Native particle MCMC (SMC/PG), revised: dropping AdvancedPS - #2859

Merged
sunxd3 merged 24 commits into
mainfrom
smc-rewrite-2
Aug 27, 2026
Merged

Native particle MCMC (SMC/PG), revised: dropping AdvancedPS#2859
sunxd3 merged 24 commits into
mainfrom
smc-rewrite-2

Conversation

@yebai

@yebai yebai commented Aug 26, 2026

Copy link
Copy Markdown
Member

SMC and PG/CSMC now run directly on Libtask and DynamicPPL. Resamplers are
types, stratified by default; conditional sweeps draw ancestors multinomially.
Philox seeding makes draws reproducible across Julia versions. Chains carry
log_normalizing_constant. Roughly twice as quick as main, allocating half.

Builds on #2848. Supersedes #2853.

Fix #2781
Fix TuringLang/AdvancedPS.jl#110
Fix TuringLang/AdvancedPS.jl#39
Fix TuringLang/AdvancedPS.jl#6
Close #2848. Close #2853.

yebai and others added 2 commits August 26, 2026 12:30
`SMC` and `PG`/`CSMC` are implemented directly on Libtask and DynamicPPL, and the
AdvancedPS dependency is gone. Each `observe` is one filtering step: under
`SMCContext` every likelihood term becomes a `Libtask.produce`, so a particle is
a suspended model execution the sweep advances, weights and resamples.

The conditional-SMC reference reproduces the retained trajectory by reusing its
values through `InitFromParams`, which is what keeps it on that trajectory when
another Gibbs component re-conditions the latents; a fork forgets them and
samples afresh. Particle state lives in the task's taped globals rather than
`task_local_storage`, and `@addlogprob!` reweights through the likelihood
accumulator (#1996).

Resampling schemes become types -- `StratifiedResampler`, `SystematicResampler`,
`MultinomialResampler`, optionally wrapped in `ESSThresholdResampler` -- with
stratified as the default, since systematic is order-dependent and can fail to be
consistent as the particle count grows. Conditional sweeps draw their ancestors
multinomially, because a correct conditional version of stratified or systematic
resampling is scheme-specific rather than "pin one draw and keep the rest".

Internal seeds derive through a counter-based Philox generator, so a fixed user
seed gives the same draws on every Julia version and platform (#2781,
AdvancedPS.jl#110). `SMC` resamples once at the end of a sweep so the returned
particles are an equal-weight sample, and chains carry
`log_normalizing_constant`, with `ess_per_step` for `SMC`.

Fixes #2781
Fixes TuringLang/AdvancedPS.jl#110
Fixes TuringLang/AdvancedPS.jl#39
Fixes TuringLang/AdvancedPS.jl#6

Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com>
Co-Authored-By: Claude Code <noreply@anthropic.com>
…eference

The tilde handlers took a varinfo argument and ignored it, reading
`particle.varinfo` instead, so anything the model body had accumulated since the
previous tilde was discarded. `@addlogprob!` and `:=` reach the varinfo directly
rather than through a handler, so a chain's `loglikelihood` omitted every
`@addlogprob!` term and `:=` variables never appeared at all, where every other
sampler reports both. The handlers now thread the argument through; a trailing
update, which has no later handler to carry it, is mirrored where it happens.
That mirroring keys on the produce-aware accumulator, so `gibbs_update_state!!`
swaps in the plain one for its out-of-task re-evaluation -- outside a task,
`get_taped_globals` returns whichever particle was last consumed rather than
failing.

A reference now carries only the retained trajectory's values. The set of
addresses each particle had assumed went with it, and so did the deep copy it
incurred on every fork. Its stated justification -- that a slice assume `x[1:2]`
cannot be checked against the stored `x[1]`, `x[2]` -- was wrong: `haskey`
resolves it.

`check_MoGtest_default` compares against the exact posterior means rather than
the idealised cluster labels, which sit far enough from the truth to spend most of
the tolerance before any Monte Carlo error.

Draws are unchanged throughout. On a T=200 state-space model, 100 PG(16)
iterations run in 9.1 s against 10.9 s and allocate 9.4 GiB against 11.6 GiB.

Co-Authored-By: Claude Code <noreply@anthropic.com>
@github-actions

Copy link
Copy Markdown
Contributor

Turing.jl documentation for PR #2859 is available at:
https://TuringLang.github.io/Turing.jl/previews/PR2859/

@codecov

codecov Bot commented Aug 26, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 94.26230% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 85.90%. Comparing base (c55358d) to head (7e3f1ff).

Files with missing lines Patch % Lines
src/mcmc/particle_mcmc.jl 94.04% 14 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #2859      +/-   ##
==========================================
+ Coverage   85.10%   85.90%   +0.79%     
==========================================
  Files          23       23              
  Lines        1517     1603      +86     
==========================================
+ Hits         1291     1377      +86     
  Misses        226      226              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

yebai added 7 commits August 26, 2026 14:23
Drop a work-stage reference that means nothing to a reader with only
the repository, and a stray line about nondeterministic evaluation
order that sits above `centred_normal` instead of the model it
describes.

Assisted-by: AI coding agent
Libtask only instruments calls marked with `@might_produce`, and the
`produce` that reweights an `@addlogprob!` term sits below
`acclogp_and_mirror!!`, which was unmarked. Reweighting worked only
because the compiler chose to inline that helper: force `@noinline`
and the term stops reweighting with no error of any kind, leaving PG
at the prior (mean 0.10 against a target of 0.79).

Assisted-by: AI coding agent
`resample_propagate!` ran at the top of the sweep loop, so every sweep
resampled once on the equal weights the particles start with. For a
multinomial draw, which every conditional sweep uses, that duplicates
particles before any data has been seen and costs diversity for
nothing.

One resample is still spent after the last observation: a particle
reveals that it has finished only by producing nothing, one pass
later, so the sweep cannot know which observation was the last.

Assisted-by: AI coding agent
PG's state was the retained `Particle`, which owns a live
`Libtask.TapedTask`. Serialising a chain sampled with
`save_state=true` therefore died inside the taped code, in
`write(::IOBuffer, ::Core.IntrinsicFunction)`. `PGState` now holds
only the retained trajectory's raw values, which is all the next
sweep's reference needs, and `Particle` takes those values directly.

`gibbs_update_state!!` re-derives the trajectory on a
`trajectory_varinfo`, so the produce-aware accumulator is never
installed outside a sweep and the swap that used to guard against its
inert `produce` goes away. Nothing carries a generator over: the
reference consumes no randomness of its own, and every other particle
is seeded from the sampler's `rng`.

Assisted-by: AI coding agent
PG draws its first sweep from the prior and dropped `initial_params`
without a word, where SMC warns. Gibbs passes the keyword to every
component's first step, so the same silence applied there.

Assisted-by: AI coding agent
Every particle scoring `-Inf` used to surface as `DomainError with
[NaN, NaN, ...]: Categorical: vector p is not a probability vector`,
naming neither the model nor the cause. Report it from the sweep
instead, at the observation no trajectory could account for.

Assisted-by: AI coding agent
A sweep whose last act was a resample leaves every weight zero, so the
closing resample drew a second time over equal weights and duplicated
particles for nothing. Return the sweep's own population instead.

Assisted-by: AI coding agent
@yebai
yebai requested a review from sunxd3 August 26, 2026 14:28

@sunxd3 sunxd3 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The core logic looks good to me.

Agent found two minor issues I think worth a second look. They are unlikely to impact user too much though.
Here are the MWEs.

using Turing, Distributions, StableRNGs, Test

@model function constrained_posterior()
    x ~ Normal()
    @addlogprob! (; logprior = x >= 0 ? 0.0 : -Inf)
    0.0 ~ Normal(x, 1)
end

smc = sample(
    StableRNG(1), constrained_posterior(), SMC(), 1_000; progress=false
)
pg = sample(
    StableRNG(1), constrained_posterior(), PG(10), 200; progress=false
)

for chain in (smc, pg)
    xs = vec(collect(chain[:x]))
    logjoints = vec(collect(chain[:logjoint]))

    @show count(x -> x < 0, xs)
    @test all(x -> x >= 0, xs)
    @test all(i -> xs[i] >= 0 || logjoints[i] != -Inf, eachindex(xs))
end
SMC: 507 / 1000 negative draws
PG:   95 / 200 negative draws

logprior didn't affect particle weight (loglikelihood is fine otherwise)


using Turing, Distributions, MCMCChains, StableRNGs, Test

@model function one_observation()
    x ~ Normal()
    0.0 ~ Normal(x, 1)
end

chain = sample(
    StableRNG(1),
    one_observation(),
    SMC(),
    32;
    chain_type=MCMCChains.Chains,
    progress=false,
)

ess = vec(collect(chain[:ess_per_step]))

@show names(chain, :internals)
@show ess
@test any(x -> !ismissing(x), ess)
names(chain, :internals) =
[:log_normalizing_constant, :ess_per_step,
 :logprior, :loglikelihood, :logjoint]

ess =
Union{Missing, Float64}[missing, missing, ..., missing]

agent explained that this is because ess_per_step being vector will be replaced by missing

yebai added 15 commits August 26, 2026 22:22
An explicit `@addlogprob!` prior factor is not part of the prior
proposal. Omitting it from particle weights silently changes the target
distribution.

Assisted-by: Codex <codex@openai.com>
MCMCChains drops vector-valued statistics. Store each filtering step
as a scalar internal so the complete diagnostic survives multi-chain
concatenation.

Assisted-by: Codex <codex@openai.com>
A particle sweep requires at least one particle. Reject non-positive
counts at the public boundary instead of failing later during weight
reduction.

Assisted-by: Codex <codex@openai.com>
`@addlogprob! (; logprior=...)` reaches DynamicPPL's `acclogp!!`, which
already routes the field on to `acclogprior!!` -- overridden here anyway.
Producing there drops a 27-line override that had to restate DynamicPPL's
field-name validation to reach the one line that mattered.

Assisted-by: Claude Code <noreply@anthropic.com>
Every other Turing sampler's `sample` defaults to `verbose=true`, so the
same call reported differently through the ensemble wrapper than when
made directly. Inert today: no `post_sample_hook` is defined for `SMC`.

Assisted-by: Claude Code <noreply@anthropic.com>
`particle_rng`'s default argument has no caller, and its `Random.Sampler`
wrapper draws the same stream as `rand(rng, UInt64)`. The `basic model`
testset asserted nothing on a model four other testsets already sample.

`bundle_smc_samples` gains the rationale a reader needs: specialising
`AbstractMCMC.bundle_samples` on `SMC` instead would recurse when the
extension hands the flattened transitions back.

Assisted-by: Claude Code <noreply@anthropic.com>
`chain[:logevidence]` became `chain[:log_normalizing_constant]`, and
SMC's per-particle `weight` is gone. Both break user code silently and
neither was listed.

Assisted-by: Claude Code <noreply@anthropic.com>
DynamicPPL's `acclogp!!` splits a two-field `@addlogprob!` across
`acclogprior!!` and `accloglikelihood!!`, so one statement produces
twice; the tests covered the two fields only in isolation.

One entry of `ess_per_step` is therefore one likelihood term, not one
observation as the changelog claimed.

Assisted-by: Claude Code <noreply@anthropic.com>
Every particle owns its own `Philox2x`, seeded from the sampler's
generator at creation and again at every fork, so a step that skips
resampling has nothing to decorrelate: the streams simply continue. The
refresh was a vestige of AdvancedPS's `TracedRNG`, which replayed
per-step seeds to rebuild the reference trajectory; this implementation
replays values instead.

Draws therefore change. Seeding `Philox2x(UInt64, seed)` directly,
rather than reseeding a generator the bare constructor had already
seeded from the OS, leaves the state and the stream identical.

Assisted-by: Claude Code <noreply@anthropic.com>
`mirror_onto_particle` had one call site. The comment claiming the RNG
section must come first because `Particle` names the generator type in
its signature was wrong: `Particle{RT<:AbstractRNG,WT<:Real}` does not.
The rest restated its own code, and the changelog carried mechanism that
belongs in the commit log.

Assisted-by: Claude Code <noreply@anthropic.com>
The effective sample size lies in `[1, nparticles]`, so a threshold
outside `[0, 1]` silently meant "always resample" or "never". Worse,
`Bool <: Real`, so `SMC(SystematicResampler(), true)` read `true` as a
threshold.

`PG` now takes any `Integer` particle count; `PG(Int32(4))` was a
`MethodError`.

Assisted-by: Claude Code <noreply@anthropic.com>
`acclogprior!!` produced before checking `ignore_missing_accumulator`, so
a varinfo without a `LogPrior` accumulator would have contributed a
weight it never accumulated. Unreachable today, since
`trajectory_varinfo` always installs one, but it inverts the invariant
the produce rests on.

The `hasacc` check is repeated rather than pushed into
`acclogp_and_mirror!!` behind a keyword: a keyword call lowers through
`Core.kwcall`, which `Libtask.@might_produce` does not name, so the
produce could go uninstrumented.

Assisted-by: Claude Code <noreply@anthropic.com>
`save_state=true` stored `nothing` as the chain's sampler state and
`initial_state` was dropped in `kwargs...` without a word, while every
other keyword SMC cannot honour already warns. Both presuppose an MCMC
loop with state to carry, which one sweep does not have.

`SMC` had an `SMCState` before this rewrite, so this is a user-visible
change and the changelog now says so. Making room for the new cases also
removed three duplicate `sample` calls: `@test_logs` returns the chain.

Assisted-by: Claude Code <noreply@anthropic.com>
`E[Oᵏ] = n·Wᵏ` is what makes a resampling scheme unbiased and nothing
checked it, only that the indices stayed in bounds. For systematic
resampling the count is within one offspring of `n·Wᵏ` on every draw,
which is sharper than the mean; that bound does not hold for stratified,
where a particle spanning strata can fall outside it.

The exact-posterior grids were not exact to the tolerance they are
compared against: `range(0.05, 4.0; length=400)` left weight 2e-5 on its
last point, shifting the reference mean for `q` by 0.0075 and
understating its standard deviation by 6%. Widened until the moments
agree to seven digits with a far wider grid, with the coverage asserted
rather than assumed.

Assisted-by: Claude Code <noreply@anthropic.com>
"both fields of a named addlogprob are weighted" pins the produced value
of each field, so the sampling-level testset covering the same ground
goes, as does a `@test_throws ErrorException` standing beside an
assertion on the same message.

The exact-SSM testsets were 4m12s of the file's 6m41s. Their tolerances
are batch-means standard errors, so 16 particles rather than 32 widens
them to match the worse mixing and the assertions stay as sharp.

Assisted-by: Claude Code <noreply@anthropic.com>
@yebai

yebai commented Aug 27, 2026

Copy link
Copy Markdown
Member Author

Thanks, @sunxd3 -- both fixed.

@sunxd3
sunxd3 added this pull request to the merge queue Aug 27, 2026
Merged via the queue into main with commit 7eebb9b Aug 27, 2026
53 of 54 checks passed
@sunxd3
sunxd3 deleted the smc-rewrite-2 branch August 27, 2026 13:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

2 participants