fix Turing threads and test with threads

Author: marius311

.github/workflows/tests_and_docs.yml

@@ -13,6 +13,7 @@ jobs:
     strategy:
       matrix:
         julia-version: ['1.7', '1.8', '1.9']
+        threads: ['1', '2']
       fail-fast: false
     steps:
       - uses: actions/checkout@v2
@@ -31,5 +32,6 @@ jobs:
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # if authenticating with GitHub Actions token
           DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }} # if authenticating with SSH deploy key
-          BUILD_DOCS: ${{ matrix.julia-version == '1.7' }} # only build/deploy docs from one version
+          BUILD_DOCS: ${{ matrix.julia-version == '1.7' && matrix.threads == '1'}} # only build/deploy docs from one version
+          JULIA_NUM_THREADS: ${{ matrix.threads }}
         timeout-minutes: 30

src/turing.jl

@@ -187,13 +187,13 @@ standardizeθ(prob::TuringMuseProblem, θ::Number) =
 
 function logLike(prob::TuringMuseProblem, x, z, θ, θ_space)
     trans = is_transformed(θ_space) ? prob.trans_z′_θ′ : prob.trans_z′_θ
-    vi = DynPPL.SimpleVarInfo((;x..., z..., θ...), 0, trans)
+    vi = DynPPL.SimpleVarInfo((;x..., z..., θ...), trans)
     DynPPL.logjoint(prob.model, vi)
 end
 
 function logPriorθ(prob::TuringMuseProblem, θ, θ_space)
     trans = is_transformed(θ_space) ? prob.trans_z′_θ′ : prob.trans_z′_θ
-    vi = DynPPL.SimpleVarInfo((;θ...), 0, trans)
+    vi = DynPPL.SimpleVarInfo((;θ...), trans)
     DynPPL.logprior(prob.model_for_prior, vi)
 end
 
@@ -210,13 +210,23 @@ end
 
 function sample_x_z(prob::TuringMuseProblem, rng::AbstractRNG, θ)
     model = DynPPL.condition(prob.model, θ)
-    vi = DynPPL.SimpleVarInfo((;), 0, prob.trans_z′_θ)
+    vi = DynPPL.SimpleVarInfo((;θ...), prob.trans_z′_θ)
     vars = DynPPL.values_as(last(DynPPL.evaluate!!(model, rng, vi)), NamedTuple)
     (x = ComponentVector(select(vars, prob.observed_vars)), z = ComponentVector(select(vars, prob.latent_vars)))
 end
 
 
 
+# benevolent type-piracy:
+function DynPPL.SimpleVarInfo(nt::NamedTuple, trans::DynPPL.AbstractTransformation)
+    if isempty(nt)
+        T = DynPPL.SIMPLEVARINFO_DEFAULT_ELTYPE
+    else
+        T = DynPPL.float_type_with_fallback(DynPPL.infer_nested_eltype(typeof(nt)))
+    end
+    DynPPL.SimpleVarInfo(nt, zero(T), trans)
+end
+
 function _namedtuple(vi::DynPPL.VarInfo)
     # `values_as` seems to return Real arays so narrow eltype
     map(x -> identity.(x), DynPPL.values_as(vi, NamedTuple))