From 50d493c7482a12961415cdbeacd73cca83d8b554 Mon Sep 17 00:00:00 2001 From: FredericWantiez Date: Sun, 23 Mar 2025 11:47:13 +0000 Subject: [PATCH 01/19] New libtask interface --- ext/AdvancedPSLibtaskExt.jl | 40 ++++++++++++++++--------------------- src/container.jl | 4 ++-- 2 files changed, 19 insertions(+), 25 deletions(-) diff --git a/ext/AdvancedPSLibtaskExt.jl b/ext/AdvancedPSLibtaskExt.jl index f4029282..e05dc29c 100644 --- a/ext/AdvancedPSLibtaskExt.jl +++ b/ext/AdvancedPSLibtaskExt.jl @@ -24,12 +24,7 @@ State wrapper to hold `Libtask.CTask` model initiated from `f`. function AdvancedPS.LibtaskModel( f::AdvancedPS.AbstractGenericModel, rng::Random.AbstractRNG, args... ) # Changed the API, need to take care of the RNG properly - return AdvancedPS.LibtaskModel( - f, - Libtask.TapedTask( - f, rng, args...; deepcopy_types=Union{AdvancedPS.TracedRNG,typeof(f)} - ), - ) + return AdvancedPS.LibtaskModel(f, Libtask.TapedTask(rng, f, args...)) end """ @@ -51,12 +46,16 @@ end # step to the next observe statement and # return the log probability of the transition (or nothing if done) -function AdvancedPS.advance!(t::LibtaskTrace, isref::Bool=false) - isref ? AdvancedPS.load_state!(t.rng) : AdvancedPS.save_state!(t.rng) - AdvancedPS.inc_counter!(t.rng) +function AdvancedPS.advance!(trace::LibtaskTrace, isref::Bool=false) + # Where is the RNG ? + # isref ? AdvancedPS.load_state!(trace.rng) : AdvancedPS.save_state!(trace.model.ctask.dynamic_scope) # Nasty + isref ? AdvancedPS.load_state!(trace.rng) : AdvancedPS.save_state!(trace.rng) + AdvancedPS.inc_counter!(trace.rng) + + Libtask.set_dynamic_scope!(trace.model.ctask, trace.rng) # Move to next step - return Libtask.consume(t.model.ctask) + return Libtask.consume(trace.model.ctask) end # create a backward reference in task_local_storage @@ -70,8 +69,9 @@ function AdvancedPS.addreference!(task::Task, trace::LibtaskTrace) end function AdvancedPS.update_rng!(trace::LibtaskTrace) - rng, = trace.model.ctask.args - trace.rng = rng + new_rng = deepcopy(trace.rng) + trace.rng = new_rng + Libtask.set_dynamic_scope!(trace.model.ctask, trace.rng) return trace end @@ -81,27 +81,23 @@ function AdvancedPS.fork(trace::LibtaskTrace, isref::Bool=false) AdvancedPS.update_rng!(newtrace) isref && AdvancedPS.delete_retained!(newtrace.model.f) isref && delete_seeds!(newtrace) - - # add backward reference - AdvancedPS.addreference!(newtrace.model.ctask.task, newtrace) return newtrace end # PG requires keeping all randomness for the reference particle # Create new task and copy randomness function AdvancedPS.forkr(trace::LibtaskTrace) - newf = AdvancedPS.reset_model(trace.model.f) + newf = AdvancedPS.reset_model(trace.model.ctask.fargs[1]) Random123.set_counter!(trace.rng, 1) - ctask = Libtask.TapedTask( - newf, trace.rng; deepcopy_types=Union{AdvancedPS.TracedRNG,typeof(trace.model.f)} - ) + ctask = Libtask.TapedTask(trace.rng, newf) new_tapedmodel = AdvancedPS.LibtaskModel(newf, ctask) # add backward reference newtrace = AdvancedPS.Trace(new_tapedmodel, trace.rng) - AdvancedPS.addreference!(ctask.task, newtrace) AdvancedPS.gen_refseed!(newtrace) + + Libtask.set_dynamic_scope!(ctask, trace.rng) # Sync trace and rng return newtrace end @@ -117,7 +113,7 @@ function AdvancedPS.observe(dist::Distributions.Distribution, x) end """ -AbstractMCMC interface. We need libtask to sample from arbitrary callable AbstractModel +AbstractMCMC interface. We need libtask to sample from arbitrary callable AbstractModelext """ function AbstractMCMC.step( @@ -138,7 +134,6 @@ function AbstractMCMC.step( else trng = AdvancedPS.TracedRNG() trace = AdvancedPS.Trace(deepcopy(model), trng) - AdvancedPS.addreference!(trace.model.ctask.task, trace) # TODO: Do we need it here ? trace end end @@ -176,7 +171,6 @@ function AbstractMCMC.sample( traces = map(1:(sampler.nparticles)) do i trng = AdvancedPS.TracedRNG() trace = AdvancedPS.Trace(deepcopy(model), trng) - AdvancedPS.addreference!(trace.model.ctask.task, trace) # Do we need it here ? trace end diff --git a/src/container.jl b/src/container.jl index 3ac64fef..107a2194 100644 --- a/src/container.jl +++ b/src/container.jl @@ -206,12 +206,12 @@ function resample_propagate!( Random.seed!(p.rng, seeds[1]) - children[j += 1] = p + children[j+=1] = p # fork additional children for k in 2:ni part = fork(p, isref) Random.seed!(part.rng, seeds[k]) - children[j += 1] = part + children[j+=1] = part end end end From af4cf444285d5c97a86dfd46fe7501dda101f09e Mon Sep 17 00:00:00 2001 From: FredericWantiez Date: Sun, 23 Mar 2025 12:02:02 +0000 Subject: [PATCH 02/19] Format --- ext/AdvancedPSLibtaskExt.jl | 2 +- src/container.jl | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/ext/AdvancedPSLibtaskExt.jl b/ext/AdvancedPSLibtaskExt.jl index e05dc29c..eb26e4f0 100644 --- a/ext/AdvancedPSLibtaskExt.jl +++ b/ext/AdvancedPSLibtaskExt.jl @@ -49,7 +49,7 @@ end function AdvancedPS.advance!(trace::LibtaskTrace, isref::Bool=false) # Where is the RNG ? # isref ? AdvancedPS.load_state!(trace.rng) : AdvancedPS.save_state!(trace.model.ctask.dynamic_scope) # Nasty - isref ? AdvancedPS.load_state!(trace.rng) : AdvancedPS.save_state!(trace.rng) + isref ? AdvancedPS.load_state!(trace.rng) : AdvancedPS.save_state!(trace.rng) AdvancedPS.inc_counter!(trace.rng) Libtask.set_dynamic_scope!(trace.model.ctask, trace.rng) diff --git a/src/container.jl b/src/container.jl index 107a2194..3ac64fef 100644 --- a/src/container.jl +++ b/src/container.jl @@ -206,12 +206,12 @@ function resample_propagate!( Random.seed!(p.rng, seeds[1]) - children[j+=1] = p + children[j += 1] = p # fork additional children for k in 2:ni part = fork(p, isref) Random.seed!(part.rng, seeds[k]) - children[j+=1] = part + children[j += 1] = part end end end From bf11785131dc404ca415d44cc470b09abacc2450 Mon Sep 17 00:00:00 2001 From: FredericWantiez Date: Sat, 5 Apr 2025 10:06:40 +0100 Subject: [PATCH 03/19] Tests --- test/container.jl | 4 ++-- test/smc.jl | 17 +++++++++++++---- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/test/container.jl b/test/container.jl index ba266499..2ba3927d 100644 --- a/test/container.jl +++ b/test/container.jl @@ -121,7 +121,7 @@ val::Ref{Int} end - function (model::Model)(rng::Random.AbstractRNG) + function (model::Model)() t = [0] while true model.val[] += 1 @@ -149,7 +149,7 @@ @testset "current trace" begin struct TaskIdModel <: AdvancedPS.AbstractGenericModel end - function (model::TaskIdModel)(rng::Random.AbstractRNG) + function (model::TaskIdModel)() # Just print the task it's running in id = objectid(AdvancedPS.current_trace()) return Libtask.produce(id) diff --git a/test/smc.jl b/test/smc.jl index 1fb4392c..6165767f 100644 --- a/test/smc.jl +++ b/test/smc.jl @@ -28,14 +28,16 @@ NormalModel() = new() end - function (m::NormalModel)(rng::Random.AbstractRNG) + function (m::NormalModel)() # First latent variable. + rng = Libtask.get_dynamic_scope() m.a = a = rand(rng, Normal(4, 5)) # First observation. AdvancedPS.observe(Normal(a, 2), 3) # Second latent variable. + rng = Libtask.get_dynamic_scope() m.b = b = rand(rng, Normal(a, 1)) # Second observation. @@ -52,8 +54,11 @@ FailSMCModel() = new() end - function (m::FailSMCModel)(rng::Random.AbstractRNG) + function (m::FailSMCModel)() + rng = Libtask.get_dynamic_scope() m.a = a = rand(rng, Normal(4, 5)) + + rng = Libtask.get_dynamic_scope() m.b = b = rand(rng, Normal(a, 1)) if a >= 4 AdvancedPS.observe(Normal(b, 2), 1.5) @@ -75,8 +80,9 @@ TestModel() = new() end - function (m::TestModel)(rng::Random.AbstractRNG) + function (m::TestModel)() # First hidden variables. + rng = Libtask.get_dynamic_scope() m.a = rand(rng, Normal(0, 1)) m.x = x = rand(rng, Bernoulli(1)) m.b = rand(rng, Gamma(2, 3)) @@ -85,6 +91,7 @@ AdvancedPS.observe(Bernoulli(x / 2), 1) # Second hidden variable. + rng = Libtask.get_dynamic_scope() m.c = rand(rng, Beta()) # Second observation. @@ -159,10 +166,12 @@ DummyModel() = new() end - function (m::DummyModel)(rng) + function (m::DummyModel)() + rng = Libtask.get_dynamic_scope() m.a = rand(rng, Normal()) AdvancedPS.observe(Normal(), m.a) + rng = Libtask.get_dynamic_scope() m.b = rand(rng, Normal()) return AdvancedPS.observe(Normal(), m.b) end From 8518a582ea15cdb84aa9d2b7bd3344dbf22e3ee5 Mon Sep 17 00:00:00 2001 From: FredericWantiez Date: Sat, 5 Apr 2025 19:11:06 +0100 Subject: [PATCH 04/19] Remove internal call --- ext/AdvancedPSLibtaskExt.jl | 7 ++++--- src/container.jl | 4 ++-- test/smc.jl | 14 +++++++------- 3 files changed, 13 insertions(+), 12 deletions(-) diff --git a/ext/AdvancedPSLibtaskExt.jl b/ext/AdvancedPSLibtaskExt.jl index eb26e4f0..60320622 100644 --- a/ext/AdvancedPSLibtaskExt.jl +++ b/ext/AdvancedPSLibtaskExt.jl @@ -87,7 +87,7 @@ end # PG requires keeping all randomness for the reference particle # Create new task and copy randomness function AdvancedPS.forkr(trace::LibtaskTrace) - newf = AdvancedPS.reset_model(trace.model.ctask.fargs[1]) + newf = AdvancedPS.reset_model(trace.model.f) Random123.set_counter!(trace.rng, 1) ctask = Libtask.TapedTask(trace.rng, newf) @@ -109,7 +109,8 @@ AdvancedPS.update_ref!(::LibtaskTrace) = nothing Observe sample `x` from distribution `dist` and yield its log-likelihood value. """ function AdvancedPS.observe(dist::Distributions.Distribution, x) - return Libtask.produce(Distributions.loglikelihood(dist, x)) + Libtask.produce(Distributions.loglikelihood(dist, x)) + return nothing end """ @@ -149,7 +150,7 @@ function AbstractMCMC.step( replayed = AdvancedPS.replay(newtrajectory) return AdvancedPS.PGSample(replayed.model.f, logevidence), - AdvancedPS.PGState(newtrajectory) + AdvancedPS.PGState(replayed) end function AbstractMCMC.sample( diff --git a/src/container.jl b/src/container.jl index 3ac64fef..107a2194 100644 --- a/src/container.jl +++ b/src/container.jl @@ -206,12 +206,12 @@ function resample_propagate!( Random.seed!(p.rng, seeds[1]) - children[j += 1] = p + children[j+=1] = p # fork additional children for k in 2:ni part = fork(p, isref) Random.seed!(part.rng, seeds[k]) - children[j += 1] = part + children[j+=1] = part end end end diff --git a/test/smc.jl b/test/smc.jl index 6165767f..1e365ea7 100644 --- a/test/smc.jl +++ b/test/smc.jl @@ -7,12 +7,12 @@ sampler = AdvancedPS.SMC(15, 0.6) @test sampler.nparticles == 15 @test sampler.resampler === - AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_systematic, 0.6) + AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_systematic, 0.6) sampler = AdvancedPS.SMC(20, AdvancedPS.resample_multinomial, 0.6) @test sampler.nparticles == 20 @test sampler.resampler === - AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_multinomial, 0.6) + AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_multinomial, 0.6) sampler = AdvancedPS.SMC(25, AdvancedPS.resample_systematic) @test sampler.nparticles == 25 @@ -98,7 +98,7 @@ return AdvancedPS.observe(Bernoulli(x / 2), 0) end - chains_smc = sample(TestModel(), AdvancedPS.SMC(100)) + chains_smc = sample(TestModel(), AdvancedPS.SMC(100); progress=false) @test all(isone(particle.x) for particle in chains_smc.trajectories) @test chains_smc.logevidence ≈ -2 * log(2) @@ -112,12 +112,12 @@ sampler = AdvancedPS.PG(60, 0.6) @test sampler.nparticles == 60 @test sampler.resampler === - AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_systematic, 0.6) + AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_systematic, 0.6) sampler = AdvancedPS.PG(80, AdvancedPS.resample_multinomial, 0.6) @test sampler.nparticles == 80 @test sampler.resampler === - AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_multinomial, 0.6) + AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_multinomial, 0.6) sampler = AdvancedPS.PG(100, AdvancedPS.resample_systematic) @test sampler.nparticles == 100 @@ -152,7 +152,7 @@ return AdvancedPS.observe(Bernoulli(x / 2), 0) end - chains_pg = sample(TestModel(), AdvancedPS.PG(10), 100) + chains_pg = sample(TestModel(), AdvancedPS.PG(10), 100; progress=false) @test all(isone(p.trajectory.x) for p in chains_pg) @test mean(x.logevidence for x in chains_pg) ≈ -2 * log(2) atol = 0.01 @@ -177,7 +177,7 @@ end pg = AdvancedPS.PG(1) - first, second = sample(DummyModel(), pg, 2) + first, second = sample(DummyModel(), pg, 2; progress=false) first_model = first.trajectory second_model = second.trajectory From 3eb727693ce44f794e54825d9b5bfd9d15eea13e Mon Sep 17 00:00:00 2001 From: FredericWantiez Date: Sat, 5 Apr 2025 19:16:40 +0100 Subject: [PATCH 05/19] Format --- ext/AdvancedPSLibtaskExt.jl | 3 +-- src/container.jl | 4 ++-- test/smc.jl | 8 ++++---- 3 files changed, 7 insertions(+), 8 deletions(-) diff --git a/ext/AdvancedPSLibtaskExt.jl b/ext/AdvancedPSLibtaskExt.jl index 60320622..e1c9d04f 100644 --- a/ext/AdvancedPSLibtaskExt.jl +++ b/ext/AdvancedPSLibtaskExt.jl @@ -149,8 +149,7 @@ function AbstractMCMC.step( newtrajectory = rand(rng, particles) replayed = AdvancedPS.replay(newtrajectory) - return AdvancedPS.PGSample(replayed.model.f, logevidence), - AdvancedPS.PGState(replayed) + return AdvancedPS.PGSample(replayed.model.f, logevidence), AdvancedPS.PGState(replayed) end function AbstractMCMC.sample( diff --git a/src/container.jl b/src/container.jl index 107a2194..3ac64fef 100644 --- a/src/container.jl +++ b/src/container.jl @@ -206,12 +206,12 @@ function resample_propagate!( Random.seed!(p.rng, seeds[1]) - children[j+=1] = p + children[j += 1] = p # fork additional children for k in 2:ni part = fork(p, isref) Random.seed!(part.rng, seeds[k]) - children[j+=1] = part + children[j += 1] = part end end end diff --git a/test/smc.jl b/test/smc.jl index 1e365ea7..186fd0c5 100644 --- a/test/smc.jl +++ b/test/smc.jl @@ -7,12 +7,12 @@ sampler = AdvancedPS.SMC(15, 0.6) @test sampler.nparticles == 15 @test sampler.resampler === - AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_systematic, 0.6) + AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_systematic, 0.6) sampler = AdvancedPS.SMC(20, AdvancedPS.resample_multinomial, 0.6) @test sampler.nparticles == 20 @test sampler.resampler === - AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_multinomial, 0.6) + AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_multinomial, 0.6) sampler = AdvancedPS.SMC(25, AdvancedPS.resample_systematic) @test sampler.nparticles == 25 @@ -112,12 +112,12 @@ sampler = AdvancedPS.PG(60, 0.6) @test sampler.nparticles == 60 @test sampler.resampler === - AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_systematic, 0.6) + AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_systematic, 0.6) sampler = AdvancedPS.PG(80, AdvancedPS.resample_multinomial, 0.6) @test sampler.nparticles == 80 @test sampler.resampler === - AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_multinomial, 0.6) + AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_multinomial, 0.6) sampler = AdvancedPS.PG(100, AdvancedPS.resample_systematic) @test sampler.nparticles == 100 From 89ac7d3b984434edf8ff92bc975fa4e15165f313 Mon Sep 17 00:00:00 2001 From: FredericWantiez Date: Sat, 5 Apr 2025 19:36:41 +0100 Subject: [PATCH 06/19] Comment --- ext/AdvancedPSLibtaskExt.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/ext/AdvancedPSLibtaskExt.jl b/ext/AdvancedPSLibtaskExt.jl index e1c9d04f..09116ced 100644 --- a/ext/AdvancedPSLibtaskExt.jl +++ b/ext/AdvancedPSLibtaskExt.jl @@ -48,7 +48,6 @@ end # return the log probability of the transition (or nothing if done) function AdvancedPS.advance!(trace::LibtaskTrace, isref::Bool=false) # Where is the RNG ? - # isref ? AdvancedPS.load_state!(trace.rng) : AdvancedPS.save_state!(trace.model.ctask.dynamic_scope) # Nasty isref ? AdvancedPS.load_state!(trace.rng) : AdvancedPS.save_state!(trace.rng) AdvancedPS.inc_counter!(trace.rng) From a54d271f765d85802c6ed6dead216d36d033800d Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Tue, 22 Apr 2025 18:44:15 +0100 Subject: [PATCH 07/19] Update Project.toml --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 7641cce5..e342cf59 100644 --- a/Project.toml +++ b/Project.toml @@ -21,7 +21,7 @@ AdvancedPSLibtaskExt = "Libtask" [compat] AbstractMCMC = "2, 3, 4, 5" Distributions = "0.23, 0.24, 0.25" -Libtask = "0.8" +Libtask = "0.9" Random = "<0.0.1, 1" Random123 = "1.3" Requires = "1.0" From dc5e594b8840b13c65d0b89e39f2e04f54e44925 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Tue, 22 Apr 2025 18:50:18 +0100 Subject: [PATCH 08/19] Update AdvancedPSLibtaskExt.jl --- ext/AdvancedPSLibtaskExt.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ext/AdvancedPSLibtaskExt.jl b/ext/AdvancedPSLibtaskExt.jl index 09116ced..fa3e55be 100644 --- a/ext/AdvancedPSLibtaskExt.jl +++ b/ext/AdvancedPSLibtaskExt.jl @@ -51,7 +51,7 @@ function AdvancedPS.advance!(trace::LibtaskTrace, isref::Bool=false) isref ? AdvancedPS.load_state!(trace.rng) : AdvancedPS.save_state!(trace.rng) AdvancedPS.inc_counter!(trace.rng) - Libtask.set_dynamic_scope!(trace.model.ctask, trace.rng) + Libtask.set_taped_globals!(trace.model.ctask, trace.rng) # Move to next step return Libtask.consume(trace.model.ctask) @@ -70,7 +70,7 @@ end function AdvancedPS.update_rng!(trace::LibtaskTrace) new_rng = deepcopy(trace.rng) trace.rng = new_rng - Libtask.set_dynamic_scope!(trace.model.ctask, trace.rng) + Libtask.set_taped_globals!(trace.model.ctask, trace.rng) return trace end @@ -96,7 +96,7 @@ function AdvancedPS.forkr(trace::LibtaskTrace) newtrace = AdvancedPS.Trace(new_tapedmodel, trace.rng) AdvancedPS.gen_refseed!(newtrace) - Libtask.set_dynamic_scope!(ctask, trace.rng) # Sync trace and rng + Libtask.set_taped_globals!(ctask, trace.rng) # Sync trace and rng return newtrace end From da248b653fc231f752657151411c397b2ca58937 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Fri, 9 May 2025 20:45:43 +0100 Subject: [PATCH 09/19] Update Project.toml --- test/Project.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/Project.toml b/test/Project.toml index c5edae50..0c8fb23a 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -19,8 +19,8 @@ GaussianDistributions = "0.5" Kalman = "0.1" HypothesisTests = "0.11" DynamicIterators = "0.4" -Libtask = "0.8" +Libtask = "0.9" Random123 = "1.3" SSMProblems = "0.1" StableRNGs = "1" -julia = "1.3" +julia = "1.10.8" From 5f2765bb3b38960e4b757903dcfcdc0420f7be2d Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Fri, 9 May 2025 20:46:10 +0100 Subject: [PATCH 10/19] Update Project.toml --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index e342cf59..60e93a80 100644 --- a/Project.toml +++ b/Project.toml @@ -27,7 +27,7 @@ Random123 = "1.3" Requires = "1.0" StatsFuns = "0.9, 1" SSMProblems = "0.1" -julia = "1.7" +julia = "1.10.8" [extras] Libtask = "6f1fad26-d15e-5dc8-ae53-837a1d7b8c9f" From 2d29914f867b0cfe7eb01b48c987d45333069cd9 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Fri, 9 May 2025 20:47:17 +0100 Subject: [PATCH 11/19] Update Project.toml --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 60e93a80..bf35e5e7 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "AdvancedPS" uuid = "576499cb-2369-40b2-a588-c64705576edc" authors = ["TuringLang"] -version = "0.6.1" +version = "0.7" [deps] AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" From 41125c4dc9851d4cdefc089500ae332645321e95 Mon Sep 17 00:00:00 2001 From: Markus Hauru Date: Mon, 2 Jun 2025 09:17:14 +0100 Subject: [PATCH 12/19] Fix calls to get_dynamic_scope --- test/smc.jl | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/test/smc.jl b/test/smc.jl index 186fd0c5..51c745a2 100644 --- a/test/smc.jl +++ b/test/smc.jl @@ -30,14 +30,14 @@ function (m::NormalModel)() # First latent variable. - rng = Libtask.get_dynamic_scope() + rng = Libtask.get_taped_globals(Any) m.a = a = rand(rng, Normal(4, 5)) # First observation. AdvancedPS.observe(Normal(a, 2), 3) # Second latent variable. - rng = Libtask.get_dynamic_scope() + rng = Libtask.get_taped_globals(Any) m.b = b = rand(rng, Normal(a, 1)) # Second observation. @@ -55,10 +55,10 @@ end function (m::FailSMCModel)() - rng = Libtask.get_dynamic_scope() + rng = Libtask.get_taped_globals(Any) m.a = a = rand(rng, Normal(4, 5)) - rng = Libtask.get_dynamic_scope() + rng = Libtask.get_taped_globals(Any) m.b = b = rand(rng, Normal(a, 1)) if a >= 4 AdvancedPS.observe(Normal(b, 2), 1.5) @@ -82,7 +82,7 @@ function (m::TestModel)() # First hidden variables. - rng = Libtask.get_dynamic_scope() + rng = Libtask.get_taped_globals(Any) m.a = rand(rng, Normal(0, 1)) m.x = x = rand(rng, Bernoulli(1)) m.b = rand(rng, Gamma(2, 3)) @@ -91,7 +91,7 @@ AdvancedPS.observe(Bernoulli(x / 2), 1) # Second hidden variable. - rng = Libtask.get_dynamic_scope() + rng = Libtask.get_taped_globals(Any) m.c = rand(rng, Beta()) # Second observation. @@ -167,11 +167,11 @@ end function (m::DummyModel)() - rng = Libtask.get_dynamic_scope() + rng = Libtask.get_taped_globals(Any) m.a = rand(rng, Normal()) AdvancedPS.observe(Normal(), m.a) - rng = Libtask.get_dynamic_scope() + rng = Libtask.get_taped_globals(Any) m.b = rand(rng, Normal()) return AdvancedPS.observe(Normal(), m.b) end From 6dff5f8b4c9cb769ed8942fa2980450b6bfa0859 Mon Sep 17 00:00:00 2001 From: Markus Hauru Date: Thu, 5 Jun 2025 17:43:18 +0100 Subject: [PATCH 13/19] Fix addreference!, I think --- ext/AdvancedPSLibtaskExt.jl | 48 +++++++++++++++++++++++-------------- src/model.jl | 2 -- test/container.jl | 7 +++--- test/smc.jl | 16 ++++++------- 4 files changed, 42 insertions(+), 31 deletions(-) diff --git a/ext/AdvancedPSLibtaskExt.jl b/ext/AdvancedPSLibtaskExt.jl index fa3e55be..2eb9fa7e 100644 --- a/ext/AdvancedPSLibtaskExt.jl +++ b/ext/AdvancedPSLibtaskExt.jl @@ -16,6 +16,17 @@ else using ..Libtask: Libtask end +# In Libtask.TapedTask.taped_globals, this extension sometimes needs to store an RNG, +# and sometimes both an RNG and other information. In Turing.jl this other information +# is a VarInfo. This struct puts those in a single struct. Note the abstract type of +# the second field. This is okay, because `get_taped_globals` needs a type assertion anyway. +struct TapedGlobals{RngType} + rng::RngType + other::Any +end + +TapedGlobals(rng::Random.AbstractRNG) = TapedGlobals(rng, nothing) + """ LibtaskModel{F} @@ -24,7 +35,7 @@ State wrapper to hold `Libtask.CTask` model initiated from `f`. function AdvancedPS.LibtaskModel( f::AdvancedPS.AbstractGenericModel, rng::Random.AbstractRNG, args... ) # Changed the API, need to take care of the RNG properly - return AdvancedPS.LibtaskModel(f, Libtask.TapedTask(rng, f, args...)) + return AdvancedPS.LibtaskModel(f, Libtask.TapedTask(TapedGlobals(rng), f, args...)) end """ @@ -47,30 +58,30 @@ end # step to the next observe statement and # return the log probability of the transition (or nothing if done) function AdvancedPS.advance!(trace::LibtaskTrace, isref::Bool=false) - # Where is the RNG ? - isref ? AdvancedPS.load_state!(trace.rng) : AdvancedPS.save_state!(trace.rng) - AdvancedPS.inc_counter!(trace.rng) + taped_globals = trace.model.ctask.taped_globals + rng = taped_globals.rng + isref ? AdvancedPS.load_state!(rng) : AdvancedPS.save_state!(rng) + AdvancedPS.inc_counter!(rng) - Libtask.set_taped_globals!(trace.model.ctask, trace.rng) + Libtask.set_taped_globals!(trace.model.ctask, TapedGlobals(rng, taped_globals.other)) + trace.rng = rng # Move to next step return Libtask.consume(trace.model.ctask) end # create a backward reference in task_local_storage -function AdvancedPS.addreference!(task::Task, trace::LibtaskTrace) - if task.storage === nothing - task.storage = IdDict() - end - task.storage[:__trace] = trace - +function AdvancedPS.addreference!(task::Libtask.TapedTask, trace::LibtaskTrace) + rng = task.taped_globals.rng + Libtask.set_taped_globals!(task, TapedGlobals(rng, trace)) return task end function AdvancedPS.update_rng!(trace::LibtaskTrace) - new_rng = deepcopy(trace.rng) + taped_globals = trace.model.ctask.taped_globals + new_rng = deepcopy(taped_globals.rng) trace.rng = new_rng - Libtask.set_taped_globals!(trace.model.ctask, trace.rng) + Libtask.set_taped_globals!(trace.model.ctask, TapedGlobals(new_rng, taped_globals.other)) return trace end @@ -86,17 +97,18 @@ end # PG requires keeping all randomness for the reference particle # Create new task and copy randomness function AdvancedPS.forkr(trace::LibtaskTrace) + taped_globals = trace.model.ctask.taped_globals + rng = taped_globals.rng newf = AdvancedPS.reset_model(trace.model.f) - Random123.set_counter!(trace.rng, 1) + Random123.set_counter!(rng, 1) + trace.rng = rng - ctask = Libtask.TapedTask(trace.rng, newf) + ctask = Libtask.TapedTask(TapedGlobals(rng, taped_globals.other), newf) new_tapedmodel = AdvancedPS.LibtaskModel(newf, ctask) # add backward reference - newtrace = AdvancedPS.Trace(new_tapedmodel, trace.rng) + newtrace = AdvancedPS.Trace(new_tapedmodel, rng) AdvancedPS.gen_refseed!(newtrace) - - Libtask.set_taped_globals!(ctask, trace.rng) # Sync trace and rng return newtrace end diff --git a/src/model.jl b/src/model.jl index c836c74f..94be5a05 100644 --- a/src/model.jl +++ b/src/model.jl @@ -64,8 +64,6 @@ function observe end function replay end function addreference! end -current_trace() = current_task().storage[:__trace] - # We need this one to be visible outside of the extension for dispatching (Turing.jl). struct LibtaskModel{F,T} f::F diff --git a/test/container.jl b/test/container.jl index 7f390571..73188b71 100644 --- a/test/container.jl +++ b/test/container.jl @@ -148,17 +148,18 @@ @test consume(a.model.ctask) == 4 end - @testset "current trace" begin + @testset "Back-reference" begin struct TaskIdModel <: AdvancedPS.AbstractGenericModel end function (model::TaskIdModel)() # Just print the task it's running in - id = objectid(AdvancedPS.current_trace()) + trace = Libtask.get_taped_globals(Any).other + id = objectid(trace) return Libtask.produce(id) end trace = AdvancedPS.Trace(TaskIdModel(), AdvancedPS.TracedRNG()) - AdvancedPS.addreference!(trace.model.ctask.task, trace) + AdvancedPS.addreference!(trace.model.ctask, trace) @test AdvancedPS.advance!(trace, false) === objectid(trace) end diff --git a/test/smc.jl b/test/smc.jl index 51c745a2..e33d3716 100644 --- a/test/smc.jl +++ b/test/smc.jl @@ -30,14 +30,14 @@ function (m::NormalModel)() # First latent variable. - rng = Libtask.get_taped_globals(Any) + rng = Libtask.get_taped_globals(Any).rng m.a = a = rand(rng, Normal(4, 5)) # First observation. AdvancedPS.observe(Normal(a, 2), 3) # Second latent variable. - rng = Libtask.get_taped_globals(Any) + rng = Libtask.get_taped_globals(Any).rng m.b = b = rand(rng, Normal(a, 1)) # Second observation. @@ -55,10 +55,10 @@ end function (m::FailSMCModel)() - rng = Libtask.get_taped_globals(Any) + rng = Libtask.get_taped_globals(Any).rng m.a = a = rand(rng, Normal(4, 5)) - rng = Libtask.get_taped_globals(Any) + rng = Libtask.get_taped_globals(Any).rng m.b = b = rand(rng, Normal(a, 1)) if a >= 4 AdvancedPS.observe(Normal(b, 2), 1.5) @@ -82,7 +82,7 @@ function (m::TestModel)() # First hidden variables. - rng = Libtask.get_taped_globals(Any) + rng = Libtask.get_taped_globals(Any).rng m.a = rand(rng, Normal(0, 1)) m.x = x = rand(rng, Bernoulli(1)) m.b = rand(rng, Gamma(2, 3)) @@ -91,7 +91,7 @@ AdvancedPS.observe(Bernoulli(x / 2), 1) # Second hidden variable. - rng = Libtask.get_taped_globals(Any) + rng = Libtask.get_taped_globals(Any).rng m.c = rand(rng, Beta()) # Second observation. @@ -167,11 +167,11 @@ end function (m::DummyModel)() - rng = Libtask.get_taped_globals(Any) + rng = Libtask.get_taped_globals(Any).rng m.a = rand(rng, Normal()) AdvancedPS.observe(Normal(), m.a) - rng = Libtask.get_taped_globals(Any) + rng = Libtask.get_taped_globals(Any).rng m.b = rand(rng, Normal()) return AdvancedPS.observe(Normal(), m.b) end From 4bf2ac5beb76f41c8b247f25847614be6c1dd50b Mon Sep 17 00:00:00 2001 From: Markus Hauru Date: Thu, 19 Jun 2025 15:24:20 +0100 Subject: [PATCH 14/19] Simplify Libtask extension --- Project.toml | 2 +- ext/AdvancedPSLibtaskExt.jl | 56 ++++++++++++++++++++++++------------- test/container.jl | 2 +- 3 files changed, 39 insertions(+), 21 deletions(-) diff --git a/Project.toml b/Project.toml index e72a1912..2bde35e8 100644 --- a/Project.toml +++ b/Project.toml @@ -21,7 +21,7 @@ AdvancedPSLibtaskExt = "Libtask" [compat] AbstractMCMC = "2, 3, 4, 5" Distributions = "0.23, 0.24, 0.25" -Libtask = "0.9" +Libtask = "0.9.2" Random = "<0.0.1, 1" Random123 = "1.3" Requires = "1.0" diff --git a/ext/AdvancedPSLibtaskExt.jl b/ext/AdvancedPSLibtaskExt.jl index 2eb9fa7e..84135b6f 100644 --- a/ext/AdvancedPSLibtaskExt.jl +++ b/ext/AdvancedPSLibtaskExt.jl @@ -17,7 +17,7 @@ else end # In Libtask.TapedTask.taped_globals, this extension sometimes needs to store an RNG, -# and sometimes both an RNG and other information. In Turing.jl this other information +# and sometimes both an RNG and other information. In Turing.jl the other information # is a VarInfo. This struct puts those in a single struct. Note the abstract type of # the second field. This is okay, because `get_taped_globals` needs a type assertion anyway. struct TapedGlobals{RngType} @@ -49,6 +49,29 @@ end const LibtaskTrace{R} = AdvancedPS.Trace{<:AdvancedPS.LibtaskModel,R} +"""Get the RNG from a `LibtaskTrace`.""" +function get_rng(trace::LibtaskTrace) + return trace.model.ctask.taped_globals.rng +end + +"""Set the RNG for a `LibtaskTrace`.""" +function set_rng!(trace::LibtaskTrace, rng::Random.AbstractRNG) + taped_globals = trace.model.ctask.taped_globals + Libtask.set_taped_globals!(trace.model.ctask, TapedGlobals(rng, taped_globals.other)) + trace.rng = rng + return trace +end + +"""Set the other "taped global" variable of a `LibtaskTrace`, other than the RNG.""" +function set_other_global!(trace::LibtaskTrace, other) + rng = get_rng(trace) + Libtask.set_taped_globals!(trace.model.ctask, TapedGlobals(rng, other)) + return trace +end + +"""Get the other "taped global" variable of a `LibtaskTrace`, other than the RNG.""" +get_other_global(trace::LibtaskTrace) = trace.model.ctask.taped_globals.other + function AdvancedPS.Trace( model::AdvancedPS.AbstractGenericModel, rng::Random.AbstractRNG, args... ) @@ -58,30 +81,25 @@ end # step to the next observe statement and # return the log probability of the transition (or nothing if done) function AdvancedPS.advance!(trace::LibtaskTrace, isref::Bool=false) - taped_globals = trace.model.ctask.taped_globals - rng = taped_globals.rng + rng = get_rng(trace) isref ? AdvancedPS.load_state!(rng) : AdvancedPS.save_state!(rng) AdvancedPS.inc_counter!(rng) - - Libtask.set_taped_globals!(trace.model.ctask, TapedGlobals(rng, taped_globals.other)) - trace.rng = rng - + set_rng!(trace, rng) # Move to next step return Libtask.consume(trace.model.ctask) end -# create a backward reference in task_local_storage -function AdvancedPS.addreference!(task::Libtask.TapedTask, trace::LibtaskTrace) - rng = task.taped_globals.rng - Libtask.set_taped_globals!(task, TapedGlobals(rng, trace)) - return task +""" +Set a backreference so that the TapedTask in `trace` stores the `trace` itself in the +taped globals. +""" +function AdvancedPS.addreference!(trace::LibtaskTrace) + set_other_global!(trace, trace) + return trace end function AdvancedPS.update_rng!(trace::LibtaskTrace) - taped_globals = trace.model.ctask.taped_globals - new_rng = deepcopy(taped_globals.rng) - trace.rng = new_rng - Libtask.set_taped_globals!(trace.model.ctask, TapedGlobals(new_rng, taped_globals.other)) + set_rng!(trace, deepcopy(get_rng(trace))) return trace end @@ -91,19 +109,19 @@ function AdvancedPS.fork(trace::LibtaskTrace, isref::Bool=false) AdvancedPS.update_rng!(newtrace) isref && AdvancedPS.delete_retained!(newtrace.model.f) isref && delete_seeds!(newtrace) + AdvancedPS.addreference!(newtrace) return newtrace end # PG requires keeping all randomness for the reference particle # Create new task and copy randomness function AdvancedPS.forkr(trace::LibtaskTrace) - taped_globals = trace.model.ctask.taped_globals - rng = taped_globals.rng + rng = get_rng(trace) newf = AdvancedPS.reset_model(trace.model.f) Random123.set_counter!(rng, 1) trace.rng = rng - ctask = Libtask.TapedTask(TapedGlobals(rng, taped_globals.other), newf) + ctask = Libtask.TapedTask(TapedGlobals(rng, get_other_global(trace)), newf) new_tapedmodel = AdvancedPS.LibtaskModel(newf, ctask) # add backward reference diff --git a/test/container.jl b/test/container.jl index 73188b71..b54452f4 100644 --- a/test/container.jl +++ b/test/container.jl @@ -159,7 +159,7 @@ end trace = AdvancedPS.Trace(TaskIdModel(), AdvancedPS.TracedRNG()) - AdvancedPS.addreference!(trace.model.ctask, trace) + AdvancedPS.addreference!(trace) @test AdvancedPS.advance!(trace, false) === objectid(trace) end From f863a7d47ed75a6b92f9514ae0357bb17506325a Mon Sep 17 00:00:00 2001 From: Markus Hauru Date: Thu, 19 Jun 2025 16:35:06 +0100 Subject: [PATCH 15/19] Small simplification to Libtask extension --- ext/AdvancedPSLibtaskExt.jl | 11 ++--------- src/rng.jl | 2 -- 2 files changed, 2 insertions(+), 11 deletions(-) diff --git a/ext/AdvancedPSLibtaskExt.jl b/ext/AdvancedPSLibtaskExt.jl index 84135b6f..5e765e0b 100644 --- a/ext/AdvancedPSLibtaskExt.jl +++ b/ext/AdvancedPSLibtaskExt.jl @@ -84,7 +84,6 @@ function AdvancedPS.advance!(trace::LibtaskTrace, isref::Bool=false) rng = get_rng(trace) isref ? AdvancedPS.load_state!(rng) : AdvancedPS.save_state!(rng) AdvancedPS.inc_counter!(rng) - set_rng!(trace, rng) # Move to next step return Libtask.consume(trace.model.ctask) end @@ -98,15 +97,10 @@ function AdvancedPS.addreference!(trace::LibtaskTrace) return trace end -function AdvancedPS.update_rng!(trace::LibtaskTrace) - set_rng!(trace, deepcopy(get_rng(trace))) - return trace -end - # Task copying version of fork for Trace. function AdvancedPS.fork(trace::LibtaskTrace, isref::Bool=false) newtrace = copy(trace) - AdvancedPS.update_rng!(newtrace) + set_rng!(newtrace, deepcopy(get_rng(newtrace))) isref && AdvancedPS.delete_retained!(newtrace.model.f) isref && delete_seeds!(newtrace) AdvancedPS.addreference!(newtrace) @@ -119,7 +113,6 @@ function AdvancedPS.forkr(trace::LibtaskTrace) rng = get_rng(trace) newf = AdvancedPS.reset_model(trace.model.f) Random123.set_counter!(rng, 1) - trace.rng = rng ctask = Libtask.TapedTask(TapedGlobals(rng, get_other_global(trace)), newf) new_tapedmodel = AdvancedPS.LibtaskModel(newf, ctask) @@ -143,7 +136,7 @@ function AdvancedPS.observe(dist::Distributions.Distribution, x) end """ -AbstractMCMC interface. We need libtask to sample from arbitrary callable AbstractModelext +AbstractMCMC interface. We need libtask to sample from arbitrary callable AbstractModel """ function AbstractMCMC.step( diff --git a/src/rng.jl b/src/rng.jl index 86109299..55e05da2 100644 --- a/src/rng.jl +++ b/src/rng.jl @@ -118,5 +118,3 @@ Random123.set_counter!(r::TracedRNG, n::Integer) = r.count = n Increase the model step counter by `n` """ inc_counter!(r::TracedRNG, n::Integer=1) = r.count += n - -function update_rng! end From 97e69ed61f001d1b37e24b266d53259575edad93 Mon Sep 17 00:00:00 2001 From: Markus Hauru Date: Thu, 19 Jun 2025 17:14:24 +0100 Subject: [PATCH 16/19] Remove addreference! --- ext/AdvancedPSLibtaskExt.jl | 22 +++++++++++----------- src/model.jl | 1 - test/container.jl | 1 - 3 files changed, 11 insertions(+), 13 deletions(-) diff --git a/ext/AdvancedPSLibtaskExt.jl b/ext/AdvancedPSLibtaskExt.jl index 5e765e0b..115062ca 100644 --- a/ext/AdvancedPSLibtaskExt.jl +++ b/ext/AdvancedPSLibtaskExt.jl @@ -49,6 +49,12 @@ end const LibtaskTrace{R} = AdvancedPS.Trace{<:AdvancedPS.LibtaskModel,R} +function Base.copy(trace::LibtaskTrace) + newtrace = AdvancedPS.Trace(copy(trace.model), deepcopy(trace.rng)) + set_other_global!(newtrace, newtrace) + return newtrace +end + """Get the RNG from a `LibtaskTrace`.""" function get_rng(trace::LibtaskTrace) return trace.model.ctask.taped_globals.rng @@ -75,7 +81,11 @@ get_other_global(trace::LibtaskTrace) = trace.model.ctask.taped_globals.other function AdvancedPS.Trace( model::AdvancedPS.AbstractGenericModel, rng::Random.AbstractRNG, args... ) - return AdvancedPS.Trace(AdvancedPS.LibtaskModel(model, rng, args...), rng) + trace = AdvancedPS.Trace(AdvancedPS.LibtaskModel(model, rng, args...), rng) + # Set a backreference so that the TapedTask in `trace` stores the `trace` itself in its + # taped globals. + set_other_global!(trace, trace) + return trace end # step to the next observe statement and @@ -88,22 +98,12 @@ function AdvancedPS.advance!(trace::LibtaskTrace, isref::Bool=false) return Libtask.consume(trace.model.ctask) end -""" -Set a backreference so that the TapedTask in `trace` stores the `trace` itself in the -taped globals. -""" -function AdvancedPS.addreference!(trace::LibtaskTrace) - set_other_global!(trace, trace) - return trace -end - # Task copying version of fork for Trace. function AdvancedPS.fork(trace::LibtaskTrace, isref::Bool=false) newtrace = copy(trace) set_rng!(newtrace, deepcopy(get_rng(newtrace))) isref && AdvancedPS.delete_retained!(newtrace.model.f) isref && delete_seeds!(newtrace) - AdvancedPS.addreference!(newtrace) return newtrace end diff --git a/src/model.jl b/src/model.jl index 94be5a05..1b89198d 100644 --- a/src/model.jl +++ b/src/model.jl @@ -62,7 +62,6 @@ end # in an extension, we just define dummy in the main module and implement them in the extension. function observe end function replay end -function addreference! end # We need this one to be visible outside of the extension for dispatching (Turing.jl). struct LibtaskModel{F,T} diff --git a/test/container.jl b/test/container.jl index b54452f4..07308e21 100644 --- a/test/container.jl +++ b/test/container.jl @@ -159,7 +159,6 @@ end trace = AdvancedPS.Trace(TaskIdModel(), AdvancedPS.TracedRNG()) - AdvancedPS.addreference!(trace) @test AdvancedPS.advance!(trace, false) === objectid(trace) end From 423d731096f96e7ca3ca998aa6a83f889b397fc6 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Mon, 23 Jun 2025 20:56:55 +0100 Subject: [PATCH 17/19] Update ext/AdvancedPSLibtaskExt.jl --- ext/AdvancedPSLibtaskExt.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ext/AdvancedPSLibtaskExt.jl b/ext/AdvancedPSLibtaskExt.jl index 115062ca..1206196b 100644 --- a/ext/AdvancedPSLibtaskExt.jl +++ b/ext/AdvancedPSLibtaskExt.jl @@ -62,8 +62,8 @@ end """Set the RNG for a `LibtaskTrace`.""" function set_rng!(trace::LibtaskTrace, rng::Random.AbstractRNG) - taped_globals = trace.model.ctask.taped_globals - Libtask.set_taped_globals!(trace.model.ctask, TapedGlobals(rng, taped_globals.other)) + other = Libtask.get_other_global(trace.model.ctask) + Libtask.set_taped_globals!(trace.model.ctask, TapedGlobals(rng, other)) trace.rng = rng return trace end From 9f3f47c67b9bd33ac7dfdf10383edef6d16a097d Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Mon, 23 Jun 2025 21:01:54 +0100 Subject: [PATCH 18/19] Update ext/AdvancedPSLibtaskExt.jl --- ext/AdvancedPSLibtaskExt.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/AdvancedPSLibtaskExt.jl b/ext/AdvancedPSLibtaskExt.jl index 1206196b..67febfd5 100644 --- a/ext/AdvancedPSLibtaskExt.jl +++ b/ext/AdvancedPSLibtaskExt.jl @@ -62,7 +62,7 @@ end """Set the RNG for a `LibtaskTrace`.""" function set_rng!(trace::LibtaskTrace, rng::Random.AbstractRNG) - other = Libtask.get_other_global(trace.model.ctask) + other = get_other_global(trace.model.ctask) Libtask.set_taped_globals!(trace.model.ctask, TapedGlobals(rng, other)) trace.rng = rng return trace From 0f6435f64a318610f2c754c7e81b8c651b2cfe64 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Mon, 23 Jun 2025 21:08:49 +0100 Subject: [PATCH 19/19] Update ext/AdvancedPSLibtaskExt.jl --- ext/AdvancedPSLibtaskExt.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/AdvancedPSLibtaskExt.jl b/ext/AdvancedPSLibtaskExt.jl index 67febfd5..d3ba9e02 100644 --- a/ext/AdvancedPSLibtaskExt.jl +++ b/ext/AdvancedPSLibtaskExt.jl @@ -62,7 +62,7 @@ end """Set the RNG for a `LibtaskTrace`.""" function set_rng!(trace::LibtaskTrace, rng::Random.AbstractRNG) - other = get_other_global(trace.model.ctask) + other = get_other_global(trace) Libtask.set_taped_globals!(trace.model.ctask, TapedGlobals(rng, other)) trace.rng = rng return trace