Expose API for Turing integration (#90)

* Move things around

* Update Project.toml

* Add test, fix comments

---------

Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com>

Project.toml

@@ -1,7 +1,7 @@
 name = "AdvancedPS"
 uuid = "576499cb-2369-40b2-a588-c64705576edc"
 authors = ["TuringLang"]
-version = "0.5.1"
+version = "0.5.2"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

ext/AdvancedPSLibtaskExt.jl

@@ -19,30 +19,29 @@ end
 """
     LibtaskModel{F}
 
-State wrapper to hold `Libtask.CTask` model initiated from `f`
+State wrapper to hold `Libtask.CTask` model initiated from `f`.
 """
-struct LibtaskModel{F1,F2}
-    f::F1
-    ctask::Libtask.TapedTask{F2}
-
-    LibtaskModel(f::F1, ctask::Libtask.TapedTask{F2}) where {F1,F2} = new{F1,F2}(f, ctask)
-end
-
-function LibtaskModel(f, args...)
-    return LibtaskModel(
+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, args...; deepcopy_types=Union{AdvancedPS.TracedRNG,typeof(f)}),
+        Libtask.TapedTask(
+            f, rng, args...; deepcopy_types=Union{AdvancedPS.TracedRNG,typeof(f)}
+        ),
     )
 end
 
-Base.copy(model::LibtaskModel) = LibtaskModel(model.f, copy(model.ctask))
+function Base.copy(model::AdvancedPS.LibtaskModel)
+    return AdvancedPS.LibtaskModel(model.f, copy(model.ctask))
+end
 
-const LibtaskTrace{R} = AdvancedPS.Trace{<:LibtaskModel,R}
+const LibtaskTrace{R} = AdvancedPS.Trace{<:AdvancedPS.LibtaskModel,R}
 
 function AdvancedPS.Trace(
     model::AdvancedPS.AbstractGenericModel, rng::Random.AbstractRNG, args...
 )
-    return AdvancedPS.Trace(LibtaskModel(model, args...), rng)
+    return AdvancedPS.Trace(AdvancedPS.LibtaskModel(model, rng, args...), rng)
 end
 
 # step to the next observe statement and
@@ -56,7 +55,7 @@ function AdvancedPS.advance!(t::LibtaskTrace, isref::Bool=false)
 end
 
 # create a backward reference in task_local_storage
-function addreference!(task::Task, trace::LibtaskTrace)
+function AdvancedPS.addreference!(task::Task, trace::LibtaskTrace)
     if task.storage === nothing
         task.storage = IdDict()
     end
@@ -65,9 +64,7 @@ function addreference!(task::Task, trace::LibtaskTrace)
     return task
 end
 
-current_trace() = current_task().storage[:__trace]
-
-function update_rng!(trace::LibtaskTrace)
+function AdvancedPS.update_rng!(trace::LibtaskTrace)
     rng, = trace.model.ctask.args
     trace.rng = rng
     return trace
@@ -76,12 +73,12 @@ end
 # Task copying version of fork for Trace.
 function AdvancedPS.fork(trace::LibtaskTrace, isref::Bool=false)
     newtrace = copy(trace)
-    update_rng!(newtrace)
+    AdvancedPS.update_rng!(newtrace)
     isref && AdvancedPS.delete_retained!(newtrace.model.f)
     isref && delete_seeds!(newtrace)
 
     # add backward reference
-    addreference!(newtrace.model.ctask.task, newtrace)
+    AdvancedPS.addreference!(newtrace.model.ctask.task, newtrace)
     return newtrace
 end
 
@@ -94,11 +91,11 @@ function AdvancedPS.forkr(trace::LibtaskTrace)
     ctask = Libtask.TapedTask(
         newf, trace.rng; deepcopy_types=Union{AdvancedPS.TracedRNG,typeof(trace.model.f)}
     )
-    new_tapedmodel = LibtaskModel(newf, ctask)
+    new_tapedmodel = AdvancedPS.LibtaskModel(newf, ctask)
 
     # add backward reference
     newtrace = AdvancedPS.Trace(new_tapedmodel, trace.rng)
-    addreference!(ctask.task, newtrace)
+    AdvancedPS.addreference!(ctask.task, newtrace)
     AdvancedPS.gen_refseed!(newtrace)
     return newtrace
 end
@@ -135,9 +132,8 @@ function AbstractMCMC.step(
             AdvancedPS.forkr(copy(state.trajectory))
         else
             trng = AdvancedPS.TracedRNG()
-            gen_model = LibtaskModel(deepcopy(model), trng)
-            trace = AdvancedPS.Trace(LibtaskModel(deepcopy(model), trng), trng)
-            addreference!(gen_model.ctask.task, trace) # Do we need it here ?
+            trace = AdvancedPS.Trace(deepcopy(model), trng)
+            AdvancedPS.addreference!(trace.model.ctask.task, trace) # TODO: Do we need it here ?
             trace
         end
     end
@@ -174,9 +170,8 @@ function AbstractMCMC.sample(
 
     traces = map(1:(sampler.nparticles)) do i
         trng = AdvancedPS.TracedRNG()
-        gen_model = LibtaskModel(deepcopy(model), trng)
-        trace = AdvancedPS.Trace(LibtaskModel(deepcopy(model), trng), trng)
-        addreference!(gen_model.ctask.task, trace) # Do we need it here ?
+        trace = AdvancedPS.Trace(deepcopy(model), trng)
+        AdvancedPS.addreference!(trace.model.ctask.task, trace) # Do we need it here ?
         trace
     end
 
@@ -202,7 +197,9 @@ function AdvancedPS.replay(particle::AdvancedPS.Particle)
     trng = deepcopy(particle.rng)
     Random123.set_counter!(trng.rng, 0)
     trng.count = 1
-    trace = AdvancedPS.Trace(LibtaskModel(deepcopy(particle.model.f), trng), trng)
+    trace = AdvancedPS.Trace(
+        AdvancedPS.LibtaskModel(deepcopy(particle.model.f), trng), trng
+    )
     score = AdvancedPS.advance!(trace, true)
     while !isnothing(score)
         score = AdvancedPS.advance!(trace, true)

src/model.jl

@@ -1,5 +1,5 @@
 """
-    Trace{F,R}
+	Trace{F,R}
 """
 mutable struct Trace{F,R}
     model::F
@@ -10,24 +10,37 @@ const Particle = Trace
 const SSMTrace{R} = Trace{<:AbstractStateSpaceModel,R}
 const GenericTrace{R} = Trace{<:AbstractGenericModel,R}
 
-# reset log probability
 reset_logprob!(::AdvancedPS.Particle) = nothing
-
 reset_model(f) = deepcopy(f)
 delete_retained!(f) = nothing
 
-Base.copy(trace::Trace) = Trace(copy(trace.model), deepcopy(trace.rng))
+"""
+	copy(trace::Trace)
 
-# This is required to make it visible from outside extensions
-function observe end
-function replay end
+Copy a trace. The `TracedRNG` is deep-copied. The inner model is shallow-copied. 
+"""
+Base.copy(trace::Trace) = Trace(copy(trace.model), deepcopy(trace.rng))
 
 """
-    gen_refseed!(particle::Particle)
+	gen_refseed!(particle::Particle)
 
-Generate a new seed for the reference particle
+Generate a new seed for the reference particle.
 """
 function gen_refseed!(particle::Particle)
     seed = split(state(particle.rng.rng), 1)
     return safe_set_refseed!(particle.rng, seed[1])
 end
+
+# A few internal functions used in the Libtask extension. Since it is not possible to access objects defined
+# 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
+
+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
+    ctask::T
+end

src/rng.jl

@@ -117,3 +117,5 @@ 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

test/container.jl

@@ -130,7 +130,7 @@
 
         # Test task copy version of trace
         trng = AdvancedPS.TracedRNG()
-        tr = AdvancedPS.Trace(Model(Ref(0)), trng, trng)
+        tr = AdvancedPS.Trace(Model(Ref(0)), trng)
 
         consume(tr.model.ctask)
         consume(tr.model.ctask)
@@ -143,6 +143,21 @@
         @test consume(a.model.ctask) == 4
     end
 
+    @testset "current trace" begin
+        struct TaskIdModel <: AdvancedPS.AbstractGenericModel end
+
+        function (model::TaskIdModel)(rng::Random.AbstractRNG)
+            # Just print the task it's running in
+            id = objectid(AdvancedPS.current_trace())
+            return Libtask.produce(id)
+        end
+
+        trace = AdvancedPS.Trace(TaskIdModel(), AdvancedPS.TracedRNG())
+        AdvancedPS.addreference!(trace.model.ctask.task, trace)
+
+        @test AdvancedPS.advance!(trace, false) === objectid(trace)
+    end
+
     @testset "seed container" begin
         seed = 1
         n = 3