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

Project.toml

@@ -54,7 +54,7 @@ AbstractPPL = "0.11.0"
 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"
@@ -67,7 +67,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.2"
 LinearAlgebra = "1"
 LogDensityProblems = "2"
 MCMCChains = "5, 6, 7"
@@ -85,7 +85,7 @@ Statistics = "1.6"
 StatsAPI = "1.6"
 StatsBase = "0.32, 0.33, 0.34"
 StatsFuns = "0.8, 0.9, 1"
-julia = "1.10.2"
+julia = "1.10.8"
 
 [extras]
 DynamicHMC = "bbc10e6e-7c05-544b-b16e-64fede858acb"

src/mcmc/particle_mcmc.jl

@@ -25,9 +25,8 @@
             "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!(
@@ -59,20 +58,10 @@
     return trace
 end
 
-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
-
-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::Any, 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 +391,11 @@
 
 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 +405,10 @@
 
 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)
@@ -481,6 +469,25 @@
 
     tmodel = TracedModel(model, sampler, newvarinfo, rng)
     newtrace = AdvancedPS.Trace(tmodel, rng)
-    AdvancedPS.addreference!(newtrace.model.ctask.task, newtrace)
     return newtrace
 end
+
+# We need to tell Libtask which calls may have `produce` calls within them. In practice most
+# of these won't be needed, because of inlining and the fact that `might_produce` is only
+# called on `:invoke` expressions rather than `:call`s, but since those are implementation
+# details of the compiler, we set a bunch of methods as might_produce = true. We start with
+# `acclogp_observe!!` which is what calls `produce` and go up the call stack.
+Libtask.might_produce(::Type{<:Tuple{typeof(DynamicPPL.acclogp_observe!!),Vararg}}) = true
+Libtask.might_produce(::Type{<:Tuple{typeof(DynamicPPL.tilde_observe!!),Vararg}}) = true
+Libtask.might_produce(::Type{<:Tuple{typeof(DynamicPPL.evaluate!!),Vararg}}) = true
+function Libtask.might_produce(
+    ::Type{<:Tuple{typeof(DynamicPPL.evaluate_threadsafe!!),Vararg}}
+)
+    return true
+end
+function Libtask.might_produce(
+    ::Type{<:Tuple{typeof(DynamicPPL.evaluate_threadunsafe!!),Vararg}}
+)
+    return true
+end
+Libtask.might_produce(::Type{<:Tuple{<:DynamicPPL.Model,Vararg}}) = true

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"

test/mcmc/particle_mcmc.jl

@@ -34,6 +34,7 @@ using Turing
 
         tested = sample(normal(), SMC(), 100)
 
+        # TODO(mhauru) This needs an explanation for why it fails.
         # failing test
         @model function fail_smc()
             a ~ Normal(4, 5)