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

Project.toml

@@ -52,7 +52,7 @@ AbstractMCMC = "5.5"
 Accessors = "0.1"
 AdvancedHMC = "0.3.0, 0.4.0, 0.5.2, 0.6, 0.7, 0.8"
 AdvancedMH = "0.8"
-AdvancedPS = "0.6.0"
+AdvancedPS = "0.7"
 AdvancedVI = "0.4"
 BangBang = "0.4.2"
 Bijectors = "0.14, 0.15"
@@ -65,7 +65,7 @@ DynamicHMC = "3.4"
 DynamicPPL = "0.36.3"
 EllipticalSliceSampling = "0.5, 1, 2"
 ForwardDiff = "0.10.3, 1"
-Libtask = "0.8.8"
+Libtask = "0.9.1"
 LinearAlgebra = "1"
 LogDensityProblems = "2"
 MCMCChains = "5, 6, 7"

src/mcmc/particle_mcmc.jl

@@ -25,9 +25,8 @@ function TracedModel(
             "Sampling with `$(sampler.alg)` does not support models with keyword arguments. See issue #2007 for more details.",
         )
     end
-    return TracedModel{AbstractSampler,AbstractVarInfo,Model,Tuple}(
-        model, sampler, varinfo, (model.f, args...)
-    )
+    evaluator = (model.f, args...)
+    return TracedModel(model, sampler, varinfo, evaluator)
 end
 
 function AdvancedPS.advance!(
@@ -71,8 +70,10 @@ function AdvancedPS.update_rng!(
     return trace
 end
 
-function Libtask.TapedTask(model::TracedModel, ::Random.AbstractRNG, args...; kwargs...) # RNG ?
-    return Libtask.TapedTask(model.evaluator[1], model.evaluator[2:end]...; kwargs...)
+function Libtask.TapedTask(taped_globals, model::TracedModel, args...; kwargs...) # RNG ?
+    return Libtask.TapedTask(
+        taped_globals, model.evaluator[1], model.evaluator[2:end]...; kwargs...
+    )
 end
 
 abstract type ParticleInference <: InferenceAlgorithm end
@@ -402,11 +403,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
+        return (trace === nothing ? varinfo : trace.model.f.varinfo)::AbstractVarInfo
     catch e
         # NOTE: this heuristic allows Libtask evaluating a model outside a `Trace`.
-        if e == KeyError(:__trace) || current_task().storage isa Nothing
+        if e == KeyError(:task_variable)
             return varinfo
         else
             rethrow(e)
@@ -416,11 +417,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
     catch e
         # NOTE: this heuristic allows Libtask evaluating a model outside a `Trace`.
-        if e == KeyError(:__trace) || current_task().storage isa Nothing
+        if e == KeyError(:task_variable)
             return rng
         else
             rethrow(e)
@@ -485,6 +485,6 @@ function AdvancedPS.Trace(
 
     tmodel = TracedModel(model, sampler, newvarinfo, rng)
     newtrace = AdvancedPS.Trace(tmodel, rng)
-    AdvancedPS.addreference!(newtrace.model.ctask.task, newtrace)
+    AdvancedPS.addreference!(newtrace.model.ctask, newtrace)
     return newtrace
 end

test/Project.toml

@@ -43,7 +43,7 @@ ADTypes = "1"
 AbstractMCMC = "5"
 AbstractPPL = "0.9, 0.10, 0.11"
 AdvancedMH = "0.6, 0.7, 0.8"
-AdvancedPS = "=0.6.0"
+AdvancedPS = "0.7"
 AdvancedVI = "0.4"
 Aqua = "0.8"
 BangBang = "0.4"