New libtask interface

Author: FredericWantiez

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
 

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