[Merged by Bors] - Perform invlinking in assume rather than implicitly in getindex

Project.toml

@@ -8,7 +8,9 @@ AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -22,7 +24,9 @@ AbstractPPL = "0.3, 0.4"
 BangBang = "0.3"
 Bijectors = "0.5.2, 0.6, 0.7, 0.8, 0.9, 0.10"
 ChainRulesCore = "0.9.7, 0.10, 1"
+ConstructionBase = "1"
 Distributions = "0.23.8, 0.24, 0.25"
+DocStringExtensions = "0.8"
 MacroTools = "0.5.6"
 Setfield = "0.7.1, 0.8"
 ZygoteRules = "0.2"

src/DynamicPPL.jl

@@ -15,6 +15,8 @@ using Setfield: Setfield
 using Setfield: Setfield
 using BangBang: BangBang
 
+using DocStringExtensions
+
 using Random: Random
 
 import Base:
@@ -176,4 +178,8 @@ include("test_utils.jl")
 @deprecate acclogp!(vi, logp) acclogp!!(vi, logp)
 @deprecate resetlogp!(vi) resetlogp!!(vi)
 
+@deprecate settrans!(vi::AbstractVarInfo, trans::Bool, vn::VarName) settrans!!(
+    vi, trans, vn
+) false
+
 end # module

src/context_implementations.jl

@@ -55,7 +55,7 @@ end
 function tilde_assume(context::PriorContext{<:NamedTuple}, right, vn, vi)
     if haskey(context.vars, getsym(vn))
         vi = setindex!!(vi, vectorize(right, get(context.vars, vn)), vn)
-        settrans!(vi, false, vn)
+        settrans!!(vi, false, vn)
     end
     return tilde_assume(PriorContext(), right, vn, vi)
 end
@@ -64,15 +64,15 @@ function tilde_assume(
 )
     if haskey(context.vars, getsym(vn))
         vi = setindex!!(vi, vectorize(right, get(context.vars, vn)), vn)
-        settrans!(vi, false, vn)
+        settrans!!(vi, false, vn)
     end
     return tilde_assume(rng, PriorContext(), sampler, right, vn, vi)
 end
 
 function tilde_assume(context::LikelihoodContext{<:NamedTuple}, right, vn, vi)
     if haskey(context.vars, getsym(vn))
         vi = setindex!!(vi, vectorize(right, get(context.vars, vn)), vn)
-        settrans!(vi, false, vn)
+        settrans!!(vi, false, vn)
     end
     return tilde_assume(LikelihoodContext(), right, vn, vi)
 end
@@ -86,7 +86,7 @@ function tilde_assume(
 )
     if haskey(context.vars, getsym(vn))
         vi = setindex!!(vi, vectorize(right, get(context.vars, vn)), vn)
-        settrans!(vi, false, vn)
+        settrans!!(vi, false, vn)
     end
     return tilde_assume(rng, LikelihoodContext(), sampler, right, vn, vi)
 end
@@ -194,7 +194,9 @@ end
 
 # fallback without sampler
 function assume(dist::Distribution, vn::VarName, vi)
-    r = vi[vn]
+    # x = vi[vn]
+    r_raw = getindex_raw(vi, vn)
+    r = maybe_invlink(vi, vn, dist, r_raw)
     return r, Bijectors.logpdf_with_trans(dist, r, istrans(vi, vn)), vi
 end
 
@@ -211,16 +213,23 @@ function assume(
         if sampler isa SampleFromUniform || is_flagged(vi, vn, "del")
             unset_flag!(vi, vn, "del")
             r = init(rng, dist, sampler)
-            vi[vn] = vectorize(dist, r)
-            settrans!(vi, false, vn)
+            vi[vn] = vectorize(dist, maybe_link(vi, vn, dist, r))
             setorder!(vi, vn, get_num_produce(vi))
         else
-            r = vi[vn]
+            # Otherwise we just extract it.
+            # r = vi[vn]
+            r_raw = getindex_raw(vi, vn)
+            r = maybe_invlink(vi, vn, dist, r_raw)
         end
     else
         r = init(rng, dist, sampler)
-        push!!(vi, vn, r, dist, sampler)
-        settrans!(vi, false, vn)
+        if istrans(vi)
+            push!!(vi, vn, link(dist, r), dist, sampler)
+            # By default `push!!` sets the transformed flag to `false`.
+            settrans!!(vi, true, vn)
+        else
+            push!!(vi, vn, r, dist, sampler)
+        end
     end
 
     return r, Bijectors.logpdf_with_trans(dist, r, istrans(vi, vn)), vi
@@ -286,7 +295,7 @@ function dot_tilde_assume(context::LikelihoodContext{<:NamedTuple}, right, left,
         var = get(context.vars, vn)
         _right, _left, _vns = unwrap_right_left_vns(right, var, vn)
         set_val!(vi, _vns, _right, _left)
-        settrans!.(Ref(vi), false, _vns)
+        settrans!!.(Ref(vi), false, _vns)
         dot_tilde_assume(LikelihoodContext(), _right, _left, _vns, vi)
     else
         dot_tilde_assume(LikelihoodContext(), right, left, vn, vi)
@@ -305,7 +314,7 @@ function dot_tilde_assume(
         var = get(context.vars, vn)
         _right, _left, _vns = unwrap_right_left_vns(right, var, vn)
         set_val!(vi, _vns, _right, _left)
-        settrans!.(Ref(vi), false, _vns)
+        settrans!!.(Ref(vi), false, _vns)
         dot_tilde_assume(rng, LikelihoodContext(), sampler, _right, _left, _vns, vi)
     else
         dot_tilde_assume(rng, LikelihoodContext(), sampler, right, left, vn, vi)
@@ -326,7 +335,7 @@ function dot_tilde_assume(context::PriorContext{<:NamedTuple}, right, left, vn,
         var = get(context.vars, vn)
         _right, _left, _vns = unwrap_right_left_vns(right, var, vn)
         set_val!(vi, _vns, _right, _left)
-        settrans!.(Ref(vi), false, _vns)
+        settrans!!.(Ref(vi), false, _vns)
         dot_tilde_assume(PriorContext(), _right, _left, _vns, vi)
     else
         dot_tilde_assume(PriorContext(), right, left, vn, vi)
@@ -345,7 +354,7 @@ function dot_tilde_assume(
         var = get(context.vars, vn)
         _right, _left, _vns = unwrap_right_left_vns(right, var, vn)
         set_val!(vi, _vns, _right, _left)
-        settrans!.(Ref(vi), false, _vns)
+        settrans!!.(Ref(vi), false, _vns)
         dot_tilde_assume(rng, PriorContext(), sampler, _right, _left, _vns, vi)
     else
         dot_tilde_assume(rng, PriorContext(), sampler, right, left, vn, vi)
@@ -390,7 +399,9 @@ function dot_assume(
     #     m .~ Normal()
     #
     # in which case `var` will have `undef` elements, even if `m` is present in `vi`.
-    r = vi[vns]
+    # r = vi[vns]
+    r_raw = getindex_raw(vi, vns)
+    r = maybe_invlink(vi, vn, dist, r_raw)
     lp = sum(zip(vns, eachcol(r))) do (vn, ri)
         return Bijectors.logpdf_with_trans(dist, ri, istrans(vi, vn))
     end
@@ -423,7 +434,8 @@ function dot_assume(
     #     m .~ Normal()
     #
     # in which case `var` will have `undef` elements, even if `m` is present in `vi`.
-    r = reshape(vi[vec(vns)], size(vns))
+    r_raw = getindex_raw(vi, vec(vns))
+    r = reshape(maybe_invlink.(Ref(vi), vns, dists, r_raw), size(vns))
     lp = sum(Bijectors.logpdf_with_trans.(dists, r, istrans(vi, vns[1])))
     return r, lp, vi
 end
@@ -462,19 +474,24 @@ function get_and_set_val!(
             r = init(rng, dist, spl, n)
             for i in 1:n
                 vn = vns[i]
-                vi[vn] = vectorize(dist, r[:, i])
-                settrans!(vi, false, vn)
+                vi[vn] = vectorize(dist, maybe_link(vi, vn, dist, r[:, i]))
                 setorder!(vi, vn, get_num_produce(vi))
             end
         else
-            r = vi[vns]
+            r_raw = getindex_raw(vi, vns)
+            r = maybe_invlink(vi, vns, dist, r_raw)
         end
     else
         r = init(rng, dist, spl, n)
         for i in 1:n
             vn = vns[i]
-            push!!(vi, vn, r[:, i], dist, spl)
-            settrans!(vi, false, vn)
+            if istrans(vi)
+                push!!(vi, vn, Bijectors.link(dist, r[:, i]), dist, spl)
+                # `push!!` sets the trans-flag to `false` by default.
+                settrans!!(vi, true, vn)
+            else
+                push!!(vi, vn, r[:, i], dist, spl)
+            end
         end
     end
     return r
@@ -496,12 +513,13 @@ function get_and_set_val!(
             for i in eachindex(vns)
                 vn = vns[i]
                 dist = dists isa AbstractArray ? dists[i] : dists
-                vi[vn] = vectorize(dist, r[i])
-                settrans!(vi, false, vn)
+                vi[vn] = vectorize(dist, maybe_link(vi, vn, dist, r[i]))
                 setorder!(vi, vn, get_num_produce(vi))
             end
         else
-            r = reshape(vi[vec(vns)], size(vns))
+            # r = reshape(vi[vec(vns)], size(vns))
+            r_raw = getindex_raw(vi, vec(vns))
+            r = maybe_invlink.(Ref(vi), vns, dists, reshape(r_raw, size(vns)))
         end
     else
         f = (vn, dist) -> init(rng, dist, spl)
@@ -511,8 +529,13 @@ function get_and_set_val!(
         # 1. Figure out the broadcast size and use a `foreach`.
         # 2. Define an anonymous function which returns `nothing`, which
         #    we then broadcast. This will allocate a vector of `nothing` though.
-        push!!.(Ref(vi), vns, r, dists, Ref(spl))
-        settrans!.(Ref(vi), false, vns)
+        if istrans(vi)
+            push!!.(Ref(vi), vns, link.(Ref(vi), vns, dists, r), dists, Ref(spl))
+            # `push!!` sets the trans-flag to `false` by default.
+            settrans!!.(Ref(vi), true, vns)
+        else
+            push!!.(Ref(vi), vns, r, dists, Ref(spl))
+        end
     end
     return r
 end