AdvancedPS v0.7 (and thus Libtask v0.9) support #2585

mhauru · 2025-06-06T13:44:39Z

The complement PR of TuringLang/AdvancedPS.jl#114, which adds support for the newly rewritten Libtask.

~~Work in progress, currently blocked by TuringLang/Libtask.jl#186~~

mhauru · 2025-06-19T15:02:33Z

The tests that I had the patience to run locally now pass. Waiting for the AdvancedPS release to be able to run the full test suite on CI.

Some indicators of speed:

julia> module MWE

       using Turing

       @model function gdemo(x, y)
           s ~ InverseGamma(2, 3)
           m ~ Normal(0, sqrt(s))
           x ~ Normal(m, sqrt(s))
           y ~ Normal(m, sqrt(s))
           return s, m
       end

       @time chn = sample(gdemo(2.5, 1.0), PG(10), 10_000)
       describe(chn)

       end

On main:

104.715858 seconds (58.48 M allocations: 13.259 GiB, 1.10% gc time, 1.20% compilation time)

On this branch:

 16.612050 seconds (116.52 M allocations: 16.296 GiB, 8.27% gc time, 5.91% compilation time)

julia> module MWE

       using Turing

       @model function f(dim=20, ::Type{T}=Float64) where T
           s = Vector{Bool}(undef, dim)
           x = Vector{T}(undef, dim)
           for i in 1:dim
               s[i] ~ Bernoulli()
               if s[i]
                   x[i] ~ Normal()
                else
                   x[i] ~ Beta()
                end
                0.0 ~ Normal(x[i])
           end
           return nothing
       end

       alg = Gibbs(
           @varname(s)=>PG(10),
           @varname(x)=>HMC(0.1, 5),
       )
       @time chn = sample(f(), alg, 1_000)

       end

On main:

 49.682945 seconds (65.43 M allocations: 9.463 GiB, 1.79% gc time, 8.23% compilation time)

On this branch:

  9.180071 seconds (61.58 M allocations: 4.028 GiB, 4.04% gc time, 55.49% compilation time)

Obviously the speed gains are all due to @willtebbutt's fantastic work on Libtask, everything else is just wrapping that work.

github-actions · 2025-06-23T20:59:46Z

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

Project.toml

yebai · 2025-06-23T21:28:33Z

src/mcmc/particle_mcmc.jl

@@ -402,11 +391,11 @@ end

 function trace_local_varinfo_maybe(varinfo)
    try
-        trace = AdvancedPS.current_trace()
-        return trace.model.f.varinfo
+        trace = Libtask.get_taped_globals(Any).other


If we change Libtask.get_taped_globals to return nothing if not inside a running TapedTask, the following try .. catch ... end can be removed.

Libtask would then lose a distinction between having nothing as the taped global within a task, and just not being within a task at all. I wonder if that distinction could be useful in some situations.

yebai · 2025-06-23T21:28:57Z

src/mcmc/particle_mcmc.jl

@@ -416,11 +405,10 @@ end

 function trace_local_rng_maybe(rng::Random.AbstractRNG)
    try
-        trace = AdvancedPS.current_trace()
-        return trace.rng
+        return Libtask.get_taped_globals(Any).rng


Same with above.

codecov · 2025-06-23T21:35:52Z

Codecov Report

Attention: Patch coverage is 41.17647% with 10 lines in your changes missing coverage. Please review.

Project coverage is 81.44%. Comparing base (a5408fb) to head (90b1cf7).

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

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #2585       +/-   ##
===========================================
+ Coverage   58.50%   81.44%   +22.94%     
===========================================
  Files          22       22               
  Lines        1458     1466        +8     
===========================================
+ Hits          853     1194      +341     
+ Misses        605      272      -333

☔ View full report in Codecov by Sentry.
📢 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.

coveralls · 2025-06-23T22:10:24Z

Pull Request Test Coverage Report for Build 16350159377

Details

7 of 17 (41.18%) changed or added relevant lines in 1 file are covered.
1 unchanged line in 1 file lost coverage.
Overall coverage decreased (-0.6%) to 84.904%

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
src/mcmc/particle_mcmc.jl	7	17	41.18%

Files with Coverage Reduction	New Missed Lines	%
src/mcmc/particle_mcmc.jl	1	84.34%

Totals
Change from base Build 16349895799:	-0.6%
Covered Lines:	1243
Relevant Lines:	1464

💛 - Coveralls

penelopeysm · 2025-06-24T01:17:50Z

Is this reviewable? The tests are failing, there's a method ambiguity that Aqua complains about, there's a Gibbs failure on 1.12 which should be disabled with @test_broken, and the sampling in mcmc/Inference is returning numerically inaccurate values:

beta binomial: Test Failed at /home/runner/work/Turing.jl/Turing.jl/test/test_utils/numerical_tests.jl:55
  Expression: ≈(E, val, atol = atol, rtol = rtol)
   Evaluated: 0.5430087089005213 ≈ 0.7142857142857143 (atol=0.05, rtol=0.0)

beta binomial: Test Failed at /home/runner/work/Turing.jl/Turing.jl/test/test_utils/numerical_tests.jl:55
  Expression: ≈(E, val, atol = atol, rtol = rtol)
   Evaluated: 0.4 ≈ 0.7142857142857143 (atol=0.1, rtol=0.0)

I don't want to speak for @mhauru in his absence but last time we spoke about this PR, it was clear that there were still a few gaps to bridge. If I were to review it at this stage, my sole comment would be to fix the tests.

mhauru · 2025-07-08T16:25:32Z

The remaining test failures on v1.10 and v1.11 should be fixed once TuringLang/Libtask.jl#192 and TuringLang/Libtask.jl#191 are merged and released (I see the tests passing locally). Will then have a look at what's going on with v1.12, and check code quality, and hopefully this could then be done.

Project.toml

Co-authored-by: Hong Ge <[email protected]>

mhauru · 2025-07-11T13:32:01Z

Seems I spoke too soon. All tests now pass on some version of Julia. But the beta binomial test still fails on v1.11, whereas the ESS test fails on v1.10. I'll try to reproduce locally, although no immediate success.

mhauru · 2025-07-14T11:12:58Z

src/mcmc/particle_mcmc.jl

-function AdvancedPS.update_rng!(
-    trace::AdvancedPS.Trace{<:AdvancedPS.LibtaskModel{<:TracedModel}}
-)
-    # Extract the `args`.
-    args = trace.model.ctask.args
-    # From `args`, extract the `SamplingContext`, which contains the RNG.
-    sampling_context = args[3]
-    rng = sampling_context.rng
-    trace.rng = rng
-    return trace
-end


This was used to keep the internal state of trace consistent, and should now be taken care of in AdvancedPS.

For my edification: AdvancedPS uses its own RNG, right? Is that distinct from the RNG in the SamplingContext (that was used here) / how are they related?

Don't know, but I don't think they are always the same. AdvancedPS does all sorts of seed splitting stuff to make sure each trace has its own source of randomness. For this purpose it also uses a particular RNG type called Random123. Frankly not sure how that will relate to the one in the context.

mhauru · 2025-07-14T11:18:42Z

src/mcmc/particle_mcmc.jl

@@ -481,6 +469,25 @@ function AdvancedPS.Trace(

    tmodel = TracedModel(model, sampler, newvarinfo, rng)
    newtrace = AdvancedPS.Trace(tmodel, rng)
-    AdvancedPS.addreference!(newtrace.model.ctask.task, newtrace)


This was about keeping internal state of newtrace consistent, and is now dealt with in the AdvancedPS.Trace constructor.

test/mcmc/Inference.jl

test/mcmc/ess.jl

mhauru · 2025-07-14T11:23:11Z

@penelopeysm, this is ready for attention now. The only test failures are Libtask v1.12 issues.

penelopeysm

Regarding correctness: I have tested quite extensively using StableRNGs on both 1.10.10 and 1.11.6 and this PR makes no difference to the sampled values. I can't figure out why the tests needed increased sample counts, but I at least feel reasonably confident that it's not because of a correctness problem.

Happy to merge whenever you like, just had one question about the rng which is more for my understanding, and it shouldn't block this PR.

mhauru · 2025-07-17T14:05:34Z

Thanks for testing that, I find that reassuring. I'll try to see if Libtask on v1.12 is an easy fix, and then either get that working or merge this as-is.

Waaaait, what, one of the v1.10 CI runs just segfaulted. Ewww.

mhauru · 2025-07-17T16:06:02Z

Passed after rerunning CI. Do we have an indeterministic test, or did we get unlucky and GHA bugged out?

penelopeysm · 2025-07-17T16:08:32Z

I would be happy to ignore it, unless and until it shows up again.

AdvancedPS v0.7 support, work in progress

f687db0

github-actions bot assigned mhauru Jun 6, 2025

Fixing particle_mcmc.jl

2366bfa

mhauru and others added 3 commits June 20, 2025 11:03

Remove use of AdvancedPS.addreference!

b4823d9

Improve a comment

d34dd3d

Merge branch 'main' into mhauru/advancedps-0.7

276e56e

Update Project.toml (#2598)

7cf8ee0

yebai marked this pull request as ready for review June 23, 2025 21:23

yebai requested a review from penelopeysm June 23, 2025 21:23

yebai reviewed Jun 23, 2025

View reviewed changes

Project.toml Show resolved Hide resolved

yebai reviewed Jun 23, 2025

View reviewed changes

mhauru and others added 2 commits July 3, 2025 14:37

Merge branch 'main' into mhauru/advancedps-0.7

8784b8e

Fix a bug and a test

1c6fad9

mhauru mentioned this pull request Jul 8, 2025

access to undefined reference TuringLang/Libtask.jl#190

Closed

Merge branch 'main' into mhauru/advancedps-0.7

bbbcebc

yebai reviewed Jul 9, 2025

View reviewed changes

Project.toml Outdated Show resolved Hide resolved

Bump Libtask to 0.9.3

df31909

Co-authored-by: Hong Ge <[email protected]>

mhauru added 2 commits July 11, 2025 15:41

Fix seed setting, increase iterations

dd84c33

Increate a test iteration count

c0fb941

mhauru commented Jul 14, 2025

View reviewed changes

test/mcmc/Inference.jl Show resolved Hide resolved

mhauru commented Jul 14, 2025

View reviewed changes

test/mcmc/Inference.jl Show resolved Hide resolved

mhauru commented Jul 14, 2025

View reviewed changes

test/mcmc/ess.jl Show resolved Hide resolved

Merge branch 'main' into mhauru/advancedps-0.7

32423b6

penelopeysm approved these changes Jul 17, 2025

View reviewed changes

Merge branch 'main' into mhauru/advancedps-0.7

90b1cf7

AdvancedPS v0.7 (and thus Libtask v0.9) support #2585

Are you sure you want to change the base?

AdvancedPS v0.7 (and thus Libtask v0.9) support #2585

Uh oh!

Conversation

mhauru commented Jun 6, 2025 • edited by yebai Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mhauru commented Jun 19, 2025

Uh oh!

github-actions bot commented Jun 23, 2025

Uh oh!

Uh oh!

yebai Jun 23, 2025

Choose a reason for hiding this comment

Uh oh!

mhauru Jul 14, 2025

Choose a reason for hiding this comment

Uh oh!

yebai Jun 23, 2025

Choose a reason for hiding this comment

Uh oh!

codecov bot commented Jun 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

coveralls commented Jun 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request Test Coverage Report for Build 16350159377

Details

💛 - Coveralls

Uh oh!

penelopeysm commented Jun 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mhauru commented Jul 8, 2025

Uh oh!

Uh oh!

mhauru commented Jul 11, 2025

Uh oh!

mhauru Jul 14, 2025

Choose a reason for hiding this comment

Uh oh!

penelopeysm Jul 17, 2025

Choose a reason for hiding this comment

Uh oh!

mhauru Jul 17, 2025

Choose a reason for hiding this comment

Uh oh!

mhauru Jul 14, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mhauru commented Jul 14, 2025

Uh oh!

penelopeysm left a comment

Choose a reason for hiding this comment

Uh oh!

mhauru commented Jul 17, 2025

Uh oh!

mhauru commented Jul 17, 2025

Uh oh!

penelopeysm commented Jul 17, 2025

Uh oh!

Uh oh!

mhauru commented Jun 6, 2025 •

edited by yebai

Loading

codecov bot commented Jun 23, 2025 •

edited

Loading

coveralls commented Jun 23, 2025 •

edited

Loading

penelopeysm commented Jun 24, 2025 •

edited

Loading