Reimplement particle MCMC (SMC and PG) natively, removing AdvancedPS

[yebai]

Opened for discussion.

Reimplements SMC and PG/CSMC directly on Libtask + DynamicPPL and drops the AdvancedPS
dependency. A rewrite based on #2848 (thanks @charlesknipp; commits co-authored).

Each observe is one filtering step: under ParticleMCMCContext, every likelihood term
becomes a Libtask.produce, so a particle is a suspended model execution we advance, weight,
and resample. The conditional-SMC reference is reproduced by replaying a per-particle
TracedRNG (no stored trajectory); a fork is reseeded to branch off. Particle state lives in
the task's taped globals (no task_local_storage), and @addlogprob! reweights via the
likelihood accumulator.

Posteriors, log-evidence, and exact reference reproduction match the old implementation;
gibbs.jl, Aqua, and log-density-metadata tests pass; performance is on par.

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

[github-actions]

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

[codecov]

Codecov Report

❌ Patch coverage is 97.10145% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 86.14%. Comparing base (11f081b) to head (8663adf).

Files with missing lines	Patch %	Lines
src/mcmc/particle_mcmc.jl	97.10%	6 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2853      +/-   ##
==========================================
+ Coverage   85.09%   86.14%   +1.05%     
==========================================
  Files          23       23              
  Lines        1516     1566      +50     
==========================================
+ Hits         1290     1349      +59     
+ Misses        226      217       -9     

☔ 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.

[charlesknipp]

Overall, I think this is a very nice revision of #2848. Although, there are a few minor choices I intentionally made when drafting that should be reconsidered here.

1. Passing AbstractMCMCEnsemble through to reweight! (only for SMC) for consistency
2. Producing the log score upon tilde_observe (or accloglikelihood!! for @addlogprob!) to ensure log score is up to date with log likelihood accumulation.

[yebai]

Will merge this in 72 hours.

[charlesknipp]

This largely resembles my initial redesign, so I don't hate it. As I point out, there are also some design changes which better motivate some of my earlier suggestions.

Reference Trajectories

My main criticism here is with respect to handling reference trajectories. This approach induces a ton of type unions and abstractions which are only necessary to accommodate a single reference trajectory.

In my PR, I introduce ReferencedContainer which assumes reference trajectories are the same type as other particles, so forking them during resampling ensures my particles are all type stable. To be fair, this is only possible with TracedRNG.

The PG inconsistency you pointed out was due to the missing DynamicPPL.InitFromParams in the tilde_assume pipeline, not the RNG tracing. Otherwise, the particle samplers in main are also incorrect.

AbstractMCMC Interface

My comment on AbstractMCMC.sample says it all. I think this is absolutely the right call, but it should be pushed further to adhere to the interface.

Passing PG(; multithreading=true) in conjunction with MCMCThreads() is not only nonsensical, but potentially dangerous. Playing devils advocate, I wonder if we could see some nice gains when pairing multithreaded sweep and parallel chains with MCMCDistributed().

Hooking into Libtask

While not mentioned in my comments, I have some concerns about the libtask usage here.

My original concern from the review was initially addressed, then subsequently replaced with the desynchronized accumulated varinfo. I understand this choice, but it still doesn't sit right with me.

I also don't like the self referential aspect of storing the entire particle in the taped globals. I think a much cleaner approach is outlined in my PR where we only need two things: the RNG and the varinfo; everything else is unnecessary storage. At the very least, I wonder if you could abuse Ref to avoid redundant copies.

Final Remarks

Compared to my implementation, this rewrite drops much of the elegance and modularity without provable gains to speed or efficiency. While I understand much of the design choices, the bloated containers and heavy indentation in functions like reweight and resample_propagate just leave a lot to be desired.

With all of that being said, if this is sufficient to pass the unit tests and guarantee correctness, it would be a waste not to merge.

[charlesknipp]

I left an earlier comment which questioned the exclusion of InitFromParams which is necessary for PG, and hence why you mentioned task based RNG wouldn't handle this alone.

[yebai]

It sounds like that would require combining conditioning with InitFromParams for reference particles. I haven't explored that, but would be happy if there is a clean mechanism.

[charlesknipp]

Is this necessary on a per particle basis? I think this should be handled in the sampler code, like in my version.

[charlesknipp]

Instead of the named tuple, why not use a raw value accumulator?

[charlesknipp]

If I understand correctly, this is necessary for the reference varinfo type to match the rest of the particles. My original draft of the PR (which is now closed) needed this in order to do subsequent CSMC steps, but ultimately led me to reimplement the traced RNG approach.

[charlesknipp]

I like how you took the approach I suggested, dropping the AbstractMCMC.step paradigm. I think you should take this a step further.

I would argue that each particle trajectory is indeed Markov, which implies that a set of particle trajectories is the same structure as a set of parallel Markov chains. Therefore, I think including AbstractMCMCEnsemble here and dispatching it to sweep! would be the correct way to parallelize this algorithm that is still consistent with the AbstractMCMC interface.

This further strengthens my earlier suggestion to replace SMC(; multithreaded=true) by passing MCMCThreads() to a dispatched reweight. See my branch to understand how it is dispatched, since it works for both multiprocessing and multithreading.

[charlesknipp]

It seems like @code_warntype is going to have a field day here.

While I understand it is not entirely possible to have a fully parametric and type stable SMC sampler with DPPL, my version at least guarantees that particles are parametrically typed and that each element shares the SAME type. The only abstract type I had was the varinfo stored in the task globals. Everything else is parametric and consistent between elements.

When the rest of the cloud (besides the reference trajectory) is the same type, it seems wasteful to make the compiler work for it.

My suggestion is to run a profiler and check to see how much bloat is from abstract types. Maybe I'm off base here, but my prior is that it will have a whole lot of type inference.

[charlesknipp]

I'm a bit surprised that this passes given that we drop the traced RNG component; as long as forking a reference changes the per particle RNG I think this is golden.

[yebai]

We will do another iteration of the design, likely as a new PR.