Fix copy mechanism (#71)

* Fix copy mechanism

* Update Project.toml

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

Author: FredericWantiez

Project.toml

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

examples/gaussian-process/script.jl

@@ -54,8 +54,6 @@ function gp_update(model::GPSSM, state, step)
     return Normal(μ[1], σ[1])
 end
 
-Libtask.tape_copy(model::GPSSM) = deepcopy(model)
-
 AdvancedPS.initialization(::GPSSM) = h(model)
 AdvancedPS.transition(model::GPSSM, state, step) = gp_update(model, state, step)
 AdvancedPS.observation(model::GPSSM, state, step) = logpdf(g(model, state, step), y[step])

examples/particle-gibbs/script.jl

@@ -33,9 +33,9 @@ Parameters = @NamedTuple begin
 end
 
 mutable struct NonLinearTimeSeries <: AbstractMCMC.AbstractModel
-    X::TArray
+    X::Array
     θ::Parameters
-    NonLinearTimeSeries(θ::Parameters) = new(TArray(Float64, θ.T), θ)
+    NonLinearTimeSeries(θ::Parameters) = new(zeros(Float64, θ.T), θ)
 end
 
 f(model::NonLinearTimeSeries, state, t) = Normal(model.θ.a * state, model.θ.q)
@@ -87,10 +87,6 @@ function (model::NonLinearTimeSeries)(rng::Random.AbstractRNG)
     end
 end
 
-# `AdvancedPS` relies on `Libtask` to copy models during their execution but we need to make sure the 
-# internal data of each model is properly copied over as well.
-Libtask.tape_copy(model::NonLinearTimeSeries) = deepcopy(model)
-
 # Here we use the particle gibbs kernel without adaptive resampling.
 model = NonLinearTimeSeries(θ₀)
 pgas = AdvancedPS.PG(Nₚ, 1.0)

src/model.jl

@@ -10,7 +10,11 @@ struct GenericModel{F1,F2} <: AbstractMCMC.AbstractModel
     GenericModel(f::F1, ctask::Libtask.TapedTask{F2}) where {F1,F2} = new{F1,F2}(f, ctask)
 end
 
-GenericModel(f, args...) = GenericModel(f, Libtask.TapedTask(f, args...))
+function GenericModel(f, args...)
+    return GenericModel(
+        f, Libtask.TapedTask(f, args...; deepcopy_types=Union{TracedRNG,typeof(f)})
+    )
+end
 Base.copy(model::GenericModel) = GenericModel(model.f, copy(model.ctask))
 
 """
@@ -74,7 +78,10 @@ function forkr(trace::GenericTrace)
     newf = reset_model(trace.model.f)
     Random123.set_counter!(trace.rng, 1)
 
-    ctask = Libtask.TapedTask(newf, trace.rng)
+    ctask = Libtask.TapedTask(
+        newf, trace.rng; deepcopy_types=Union{TracedRNG,typeof(trace.model.f)}
+    )
+    #ctask = Libtask.TapedTask(newf, trace.rng)
     new_tapedmodel = GenericModel(newf, ctask)
 
     # add backward reference

test/container.jl

@@ -190,6 +190,6 @@
 
         AdvancedPS.advance!(ref)
         child = AdvancedPS.fork(ref, true)
-        @test length(new_tr.rng.keys) == 1
+        @test length(child.rng.keys) == 1
     end
 end

test/pgas.jl

@@ -79,7 +79,6 @@
 
     @testset "rng stability" begin
         model = BaseModel(Params(0.9, 0.32, 1))
-
         seed = 10
         rng = Random.MersenneTwister(seed)