Match AdvancedPS 0.4 API (#1858)

* Conflict

* Conflict

* Weird typo

* Fix `LogPoisson` typo in docstring (#1888)

Simple fix that I saw.

* Conflict

* test

* Update state, still missing one

* Sneaky RNG

* Fix logevidence

* Small test

* AdvancedPS

* Updating the RNG properly

* Update Inference.jl

* Update TuringCI.yml

* Update DynamicHMC.yml

* Update Numerical.yml

* Conflict

* Conflict

* Weird typo

* Conflict

* test

* Update state, still missing one

* Sneaky RNG

* Fix logevidence

* Small test

* AdvancedPS

* Updating the RNG properly

* Update Inference.jl

* Update TuringCI.yml

* Update DynamicHMC.yml

* Update Numerical.yml

* Scaffolding

* Clean up

* Forward kwargs

* Optional args, TArray

* Fix invalid reference error

* Remove mentions of TArray to avoid confusion

* Remove section on Libtask

* Proper copy of TracedModel

* Updating the VarInfo

* Clean up comments, update reset_model

* Increase sample numbers

Co-authored-by: Frederic Wantiez <[email protected]>
Co-authored-by: Hong Ge <[email protected]>
Co-authored-by: Jose Storopoli <[email protected]>
Co-authored-by: Frederic Wantiez <[email protected]>

Author: 5 people

.github/workflows/DynamicHMC.yml

@@ -12,7 +12,7 @@ jobs:
     strategy:
       matrix:
         version:
-          - '1.6'
+          - '1.7'
           - '1'
         os:
           - ubuntu-latest

.github/workflows/Numerical.yml

@@ -12,7 +12,7 @@ jobs:
     strategy:
       matrix:
         version:
-          - '1.6'
+          - '1.7'
           - '1'
         os:
           - ubuntu-latest

.github/workflows/TuringCI.yml

@@ -13,7 +13,7 @@ jobs:
     strategy:
       matrix:
         version:
-          - '1.6'
+          - '1.7'
           - '1'
         os:
           - ubuntu-latest
@@ -23,15 +23,15 @@ jobs:
           - 1
           - 2
         include:
-          - version: '1.6'
+          - version: '1.7'
             os: ubuntu-latest
             arch: x86
             num_threads: 2
-          - version: '1.6'
+          - version: '1.7'
             os: windows-latest
             arch: x64
             num_threads: 2
-          - version: '1.6'
+          - version: '1.7'
             os: macOS-latest
             arch: x64
             num_threads: 2
@@ -58,13 +58,13 @@ jobs:
         env:
           JULIA_NUM_THREADS: ${{ matrix.num_threads }}
       - uses: julia-actions/julia-processcoverage@v1
-        if: matrix.version == '1.6' && matrix.os == 'ubuntu-latest' && matrix.num_threads == 1
+        if: matrix.version == '1.7' && matrix.os == 'ubuntu-latest' && matrix.num_threads == 1
       - uses: codecov/codecov-action@v1
-        if: matrix.version == '1.6' && matrix.os == 'ubuntu-latest' && matrix.num_threads == 1
+        if: matrix.version == '1.7' && matrix.os == 'ubuntu-latest' && matrix.num_threads == 1
         with:
           file: lcov.info
       - uses: coverallsapp/github-action@master
-        if: matrix.version == '1.6' && matrix.os == 'ubuntu-latest' && matrix.num_threads == 1
+        if: matrix.version == '1.7' && matrix.os == 'ubuntu-latest' && matrix.num_threads == 1
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}
           path-to-lcov: lcov.info

Project.toml

@@ -38,7 +38,7 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 AbstractMCMC = "4"
 AdvancedHMC = "0.3.0"
 AdvancedMH = "0.6.8"
-AdvancedPS = "0.3.4"
+AdvancedPS = "0.4"
 AdvancedVI = "0.1"
 BangBang = "0.3"
 Bijectors = "0.8, 0.9, 0.10"
@@ -49,8 +49,8 @@ DocStringExtensions = "0.8, 0.9"
 DynamicPPL = "0.21"
 EllipticalSliceSampling = "0.5, 1"
 ForwardDiff = "0.10.3"
-Libtask = "0.6.7, 0.7"
 LogDensityProblems = "0.12, 1"
+Libtask = "0.7, 0.8"
 MCMCChains = "5"
 NamedArrays = "0.9"
 Reexport = "0.2, 1"

docs/src/library/api.md

@@ -42,15 +42,3 @@ BinomialLogit
 VecBinomialLogit
 OrderedLogistic
 ```
-
-
-## Data Structures
-```@docs
-TArray
-```
-
-## Utilities
-```@docs
-tzeros
-```
-

docs/src/using-turing/guide.md

@@ -10,7 +10,7 @@ title: Guide
 
 ### Introduction
 
-A probabilistic program is Julia code wrapped in a `@model` macro. It can use arbitrary Julia code, but to ensure correctness of inference it should not have external effects or modify global state. Stack-allocated variables are safe, but mutable heap-allocated objects may lead to subtle bugs when using task copying. To help avoid those we provide a Turing-safe datatype `TArray` that can be used to create mutable arrays in Turing programs.
+A probabilistic program is Julia code wrapped in a `@model` macro. It can use arbitrary Julia code, but to ensure correctness of inference it should not have external effects or modify global state. Stack-allocated variables are safe, but mutable heap-allocated objects may lead to subtle bugs when using task copying. By default Libtask deepcopies `Array` and `Dict` objects when copying task to avoid bugs with data stored in mutable structure in Turing models. 
 
 
 To specify distributions of random variables, Turing programs should use the `~` notation:
@@ -369,7 +369,7 @@ The element type of a vector (or matrix) of random variables should match the `e
     1. `Real` to enable auto-differentiation through the model which uses special number types that are sub-types of `Real`, or
     2. Some type parameter `T` defined in the model header using the type parameter syntax, e.g. `function gdemo(x, ::Type{T} = Float64) where {T}`.
 Similarly, when using a particle sampler, the Julia variable used should either be:
-    1. A `TArray`, or
+    1. An `Array`, or
     2. An instance of some type parameter `T` defined in the model header using the type parameter syntax, e.g. `function gdemo(x, ::Type{T} = Vector{Float64}) where {T}`.
 
 
@@ -596,48 +596,6 @@ plot(chn) # Plots statistics of the samples.
 There are numerous functions in addition to `describe` and `plot` in the `MCMCChains` package, such as those used in convergence diagnostics. For more information on the package, please see the [GitHub repository](https://github.com/TuringLang/MCMCChains.jl).
 
 
-### Working with Libtask.jl
-
-
-The [Libtask.jl](https://github.com/TuringLang/Libtask.jl) library provides write-on-copy data structures that are safe for use in Turing's particle-based samplers. One data structure in particular is often required for use – the [`TArray`](http://turing.ml/docs/library/#Libtask.TArray). The following sampler types require the use of a `TArray` to store distributions:
-
-
-  * `IPMCMC`
-  * `IS`
-  * `PG`
-  * `PMMH`
-  * `SMC`
-
-
-If you do not use a `TArray` to store arrays of distributions when using a particle-based sampler, you may experience errors.
-
-
-Here is an example of how the `TArray` (using a `TArray` constructor function called `tzeros`) can be applied in this way:
-
-
-```julia
-# Turing model definition.
-@model function BayesHmm(y)
-    # Declare a TArray with a length of N.
-    s = tzeros(Int, N)
-    m = Vector{Real}(undef, K)
-    T = Vector{Vector{Real}}(undef, K)
-    for i = 1:K
-        T[i] ~ Dirichlet(ones(K)/K)
-        m[i] ~ Normal(i, 0.01)
-    end
-
-    # Draw from a distribution for each element in s.
-    s[1] ~ Categorical(K)
-    for i = 2:N
-        s[i] ~ Categorical(vec(T[s[i-1]]))
-        y[i] ~ Normal(m[s[i]], 0.1)
-    end
-    return (s, m)
-end;
-```
-
-
 ### Changing Default Settings
 
 

src/essential/container.jl

@@ -9,14 +9,13 @@ function TracedModel(
     model::Model,
     sampler::AbstractSampler,
     varinfo::AbstractVarInfo,
+    rng::Random.AbstractRNG
 ) 
-    # evaluate!!(m.model, varinfo, SamplingContext(Random.AbstractRNG, m.sampler, DefaultContext()))
-    context = SamplingContext(DynamicPPL.Random.GLOBAL_RNG, sampler, DefaultContext())
+    context = SamplingContext(rng, sampler, DefaultContext())
     evaluator = _get_evaluator(model, varinfo, context)
     return TracedModel{AbstractSampler,AbstractVarInfo,Model,Tuple}(model, sampler, varinfo, evaluator)
 end
 
-# Smiliar to `evaluate!!` except that we return the evaluator signature without execution.
 # TODO: maybe move to DynamicPPL
 @generated function _get_evaluator(
     model::Model{_F,argnames}, varinfo, context
@@ -39,42 +38,49 @@ end
     end
 end
 
-function Base.copy(trace::AdvancedPS.Trace{<:TracedModel})
-    f = trace.f
-    newf = TracedModel(f.model, f.sampler, deepcopy(f.varinfo))
-    return AdvancedPS.Trace(newf, copy(trace.task))
+
+function Base.copy(model::AdvancedPS.GenericModel{<:TracedModel})
+    newtask = copy(model.ctask)
+    newmodel = TracedModel{AbstractSampler,AbstractVarInfo,Model,Tuple}(deepcopy(model.f.model), deepcopy(model.f.sampler), deepcopy(model.f.varinfo), deepcopy(model.f.evaluator))
+    gen_model = AdvancedPS.GenericModel(newmodel, newtask)
+    return gen_model
 end
 
-function AdvancedPS.advance!(trace::AdvancedPS.Trace{<:TracedModel})
-    DynamicPPL.increment_num_produce!(trace.f.varinfo)
-    score = consume(trace.task)
+function AdvancedPS.advance!(trace::AdvancedPS.Trace{<:AdvancedPS.GenericModel{<:TracedModel}}, isref::Bool=false)
+    # Make sure we load/reset the rng in the new replaying mechanism
+    DynamicPPL.increment_num_produce!(trace.model.f.varinfo)
+    isref ? AdvancedPS.load_state!(trace.rng) : AdvancedPS.save_state!(trace.rng)
+    score = consume(trace.model.ctask)
     if score === nothing
         return
     else
-        return score + DynamicPPL.getlogp(trace.f.varinfo)
+        return score + DynamicPPL.getlogp(trace.model.f.varinfo)
     end
 end
 
-function AdvancedPS.delete_retained!(f::TracedModel)
-    DynamicPPL.set_retained_vns_del_by_spl!(f.varinfo, f.sampler)
-    return
+function AdvancedPS.delete_retained!(trace::TracedModel)
+    DynamicPPL.set_retained_vns_del_by_spl!(trace.varinfo, trace.sampler)
+    return trace
 end
 
-function AdvancedPS.reset_model(f::TracedModel)
-    newvarinfo = deepcopy(f.varinfo)
-    DynamicPPL.reset_num_produce!(newvarinfo)
-    return TracedModel(f.model, f.sampler, newvarinfo)
+function AdvancedPS.reset_model(trace::TracedModel)
+    DynamicPPL.reset_num_produce!(trace.varinfo)
+    return trace
 end
 
-function AdvancedPS.reset_logprob!(f::TracedModel)
-    DynamicPPL.resetlogp!!(f.varinfo)
-    return
+function AdvancedPS.reset_logprob!(trace::TracedModel)
+    DynamicPPL.resetlogp!!(trace.model.varinfo)
+    return trace
 end
 
-function Libtask.TapedTask(model::TracedModel)
-    return Libtask.TapedTask(model.evaluator[1], model.evaluator[2:end]...)
+function AdvancedPS.update_rng!(trace::AdvancedPS.Trace{AdvancedPS.GenericModel{TracedModel{M,S,V,E}, F}, R}) where {M,S,V,E,F,R} 
+    args = trace.model.ctask.args
+    _, _, container, = args
+    rng = container.rng
+    trace.rng = rng
+    return trace
 end
 
-function Libtask.TapedTask(model::TracedModel, ::Random.AbstractRNG)
-    return Libtask.TapedTask(model)
+function Libtask.TapedTask(model::TracedModel, rng::Random.AbstractRNG; kwargs...)
+    return Libtask.TapedTask(model.evaluator[1], model.evaluator[2:end]...; kwargs...)
 end