[Merged by Bors] - Linearization/flattening of SimpleVarInfo

src/DynamicPPL.jl

@@ -57,7 +57,9 @@ export AbstractVarInfo,
     setorder!,
     istrans,
     link!,
+    link!!,
     invlink!,
+    invlink!!,
     tonamedtuple,
     # VarName (reexport from AbstractPPL)
     VarName,
@@ -150,5 +152,6 @@ include("compat/ad.jl")
 include("loglikelihoods.jl")
 include("submodel_macro.jl")
 include("test_utils.jl")
+include("transforming.jl")
 
 end # module

src/sampler.jl

@@ -67,6 +67,21 @@ function AbstractMCMC.step(
     return vi, nothing
 end
 
+function make_default_varinfo(
+    rng::Random.AbstractRNG, model::Model, sampler::AbstractSampler
+)
+    return make_default_varinfo(rng, model, sampler, DefaultContext())
+end
+function make_default_varinfo(
+    rng::Random.AbstractRNG,
+    model::Model,
+    sampler::AbstractSampler,
+    context::AbstractContext,
+)
+    init_sampler = initialsampler(sampler)
+    return VarInfo(rng, model, init_sampler, context)
+end
+
 # initial step: general interface for resuming and
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
@@ -82,23 +97,17 @@ function AbstractMCMC.step(
     end
 
     # Sample initial values.
-    _spl = initialsampler(spl)
-    vi = VarInfo(rng, model, _spl)
+    vi = default_varinfo(rng, model, spl)
 
     # Update the parameters if provided.
     if init_params !== nothing
-        vi = initialize_parameters!!(vi, init_params, spl)
+        vi = initialize_parameters!!(vi, init_params, spl, model)
 
         # Update joint log probability.
-        # TODO: fix properly by using sampler and evaluation contexts
         # This is a quick fix for https://github.com/TuringLang/Turing.jl/issues/1588
         # and https://github.com/TuringLang/Turing.jl/issues/1563
         # to avoid that existing variables are resampled
-        if _spl isa SampleFromUniform
-            model(rng, vi, SampleFromPrior())
-        else
-            model(rng, vi, _spl)
-        end
+        vi = last(evaluate!!(model, vi, DefaultContext()))
     end
 
     return initialstep(rng, model, spl, vi; init_params=init_params, kwargs...)
@@ -121,7 +130,9 @@ By default, it returns an instance of [`SampleFromPrior`](@ref).
 """
 initialsampler(spl::Sampler) = SampleFromPrior()
 
-function initialize_parameters!!(vi::AbstractVarInfo, init_params, spl::Sampler)
+function initialize_parameters!!(
+    vi::AbstractVarInfo, init_params, spl::Sampler, model::Model
+)
     @debug "Using passed-in initial variable values" init_params
 
     # Flatten parameters.
@@ -132,8 +143,7 @@ function initialize_parameters!!(vi::AbstractVarInfo, init_params, spl::Sampler)
     # Get all values.
     linked = islinked(vi, spl)
     if linked
-        # TODO: Make work with immutable `vi`.
-        invlink!(vi, spl)
+        vi = invlink!!(vi, spl, model)
     end
     theta = vi[spl]
     length(theta) == length(init_theta) ||
@@ -150,8 +160,7 @@ function initialize_parameters!!(vi::AbstractVarInfo, init_params, spl::Sampler)
     # Update in `vi`.
     vi = setindex!!(vi, theta, spl)
     if linked
-        # TODO: Make work with immutable `vi`.
-        link!(vi, spl)
+        vi = link!!(vi, spl, model)
     end
 
     return vi

src/simple_varinfo.jl

@@ -1,8 +1,3 @@
-abstract type AbstractTransformation end
-
-struct NoTransformation <: AbstractTransformation end
-struct DefaultTransformation <: AbstractTransformation end
-
 """
     $(TYPEDEF)
 
@@ -197,6 +192,8 @@ struct SimpleVarInfo{NT,T,C<:AbstractTransformation} <: AbstractVarInfo
     transformation::C
 end
 
+transformation(vi::SimpleVarInfo) = vi.transformation
+
 SimpleVarInfo(values, logp) = SimpleVarInfo(values, logp, NoTransformation())
 
 function SimpleVarInfo{T}(θ) where {T<:Real}
@@ -227,9 +224,17 @@ function SimpleVarInfo{T}(
     return SimpleVarInfo(values, convert(T, getlogp(vi)))
 end
 
+SimpleVarInfo(svi::SimpleVarInfo, spl, x::AbstractVector) = unflatten(svi, x)
+
+unflatten(svi::SimpleVarInfo, spl, x::AbstractVector) = unflatten(svi, x)
+function unflatten(svi::SimpleVarInfo, x::AbstractVector)
+    return Setfield.@set svi.values = unflatten(svi.values, x)
+end
+
 function BangBang.empty!!(vi::SimpleVarInfo)
-    Setfield.@set resetlogp!!(vi).values = empty!!(vi.values)
+    return resetlogp!!(Setfield.@set vi.values = empty!!(vi.values))
 end
+Base.isempty(vi::SimpleVarInfo) = isempty(vi.values)
 
 getlogp(vi::SimpleVarInfo) = vi.logp
 setlogp!!(vi::SimpleVarInfo, logp) = Setfield.@set vi.logp = logp
@@ -308,11 +313,8 @@ end
 # HACK: Needed to disambiguiate.
 Base.getindex(vi::SimpleVarInfo, vns::Vector{<:VarName}) = map(Base.Fix1(getindex, vi), vns)
 
-Base.getindex(vi::SimpleVarInfo, spl::SampleFromPrior) = vi.values
-Base.getindex(vi::SimpleVarInfo, spl::SampleFromUniform) = vi.values
-
-# TODO: Should we do better?
-Base.getindex(vi::SimpleVarInfo, spl::Sampler) = vi.values
+Base.getindex(svi::SimpleVarInfo, ::Colon) = values_as(svi, Vector)
+Base.getindex(svi::SimpleVarInfo, ::Sampler) = svi[:]
 
 # Since we don't perform any transformations in `getindex` for `SimpleVarInfo`
 # we simply call `getindex` in `getindex_raw`.
@@ -365,6 +367,10 @@ function BangBang.setindex!!(vi::SimpleVarInfo, val, vn::VarName)
     return Setfield.@set vi.values = set!!(vi.values, vn, val)
 end
 
+function BangBang.setindex!!(vi::SimpleVarInfo, val, spl::AbstractSampler)
+    return unflatten(vi, spl, val)
+end
+
 # TODO: Specialize to handle certain cases, e.g. a collection of `VarName` with
 # same symbol and same type of, say, `IndexLens`, for improved `.~` performance.
 function BangBang.setindex!!(vi::SimpleVarInfo, vals, vns::AbstractVector{<:VarName})
@@ -509,6 +515,45 @@ end
 
 # HACK: Allows us to re-use the implementation of `dot_tilde`, etc. for literals.
 increment_num_produce!(::SimpleOrThreadSafeSimple) = nothing
+setgid!(vi::SimpleOrThreadSafeSimple, gid::Selector, vn::VarName) = nothing
+
+# We need these to be compatible with how chains are constructed from `AbstractVarInfo` in Turing.jl.
+# TODO: Move away from using these `tonamedtuple` methods.
+function tonamedtuple(vi::SimpleOrThreadSafeSimple{<:NamedTuple{names}}) where {names}
+    nt_vals = map(keys(vi)) do vn
+        val = vi[vn]
+        vns = collect(DynamicPPL.TestUtils.varname_leaves(vn, val))
+        vals = map(Base.Fix1(getindex, vi), vns)
+        (vals, map(string, vns))
+    end
+
+    return NamedTuple{names}(nt_vals)
+end
+
+function tonamedtuple(vi::SimpleOrThreadSafeSimple{<:Dict})
+    syms_to_result = Dict{Symbol,Tuple{Vector{Real},Vector{String}}}()
+    for vn in keys(vi)
+        # Extract the leaf varnames and values.
+        val = vi[vn]
+        vns = collect(DynamicPPL.TestUtils.varname_leaves(vn, val))
+        vals = map(Base.Fix1(getindex, vi), vns)
+
+        # Determine the corresponding symbol.
+        sym = only(unique(map(getsym, vns)))
+
+        # Initialize entry if not yet initialized.
+        if !haskey(syms_to_result, sym)
+            syms_to_result[sym] = (Real[], String[])
+        end
+
+        # Combine with old result.
+        old_vals, old_string_vns = syms_to_result[sym]
+        syms_to_result[sym] = (vcat(old_vals, vals), vcat(old_string_vns, map(string, vns)))
+    end
+
+    # Construct `NamedTuple`.
+    return NamedTuple(pairs(syms_to_result))
+end
 
 # NOTE: We don't implement `settrans!!(vi, trans, vn)`.
 function settrans!!(vi::SimpleVarInfo, trans)
@@ -525,6 +570,8 @@ istrans(vi::SimpleVarInfo) = !(vi.transformation isa NoTransformation)
 istrans(vi::SimpleVarInfo, vn::VarName) = istrans(vi)
 istrans(vi::ThreadSafeVarInfo{<:SimpleVarInfo}, vn::VarName) = istrans(vi.varinfo, vn)
 
+islinked(vi::SimpleVarInfo, ::Union{Sampler,SampleFromPrior}) = istrans(vi)
+
 """
     values_as(varinfo[, Type])
 
@@ -536,6 +583,10 @@ values_as(vi::SimpleVarInfo) = vi.values
 values_as(vi::SimpleVarInfo, ::Type{Dict}) = Dict(pairs(vi.values))
 values_as(vi::SimpleVarInfo, ::Type{NamedTuple}) = NamedTuple(pairs(vi.values))
 values_as(vi::SimpleVarInfo{<:NamedTuple}, ::Type{NamedTuple}) = vi.values
+function values_as(vi::SimpleVarInfo{<:Any,T}, ::Type{Vector}) where {T}
+    isempty(vi.values) && return T[]
+    return mapreduce(v -> vec([v;]), vcat, values(vi.values))
+end
 
 """
     logjoint(model::Model, θ)
@@ -632,3 +683,35 @@ julia> # Truth.
 ```
 """
 Distributions.loglikelihood(model::Model, θ) = loglikelihood(model, SimpleVarInfo(θ))
+
+# Allow usage of `NamedBijector` too.
+function link!!(
+    t::BijectorTransformation{<:Bijectors.NamedBijector},
+    vi::SimpleVarInfo{<:NamedTuple},
+    spl::AbstractSampler,
+    model::Model,
+)
+    # TODO: Make sure that `spl` is respected.
+    b = t.bijector
+    x = vi.values
+    y, logjac = with_logabsdet_jacobian(b, x)
+    lp_new = getlogp(vi) - logjac
+    vi_new = setlogp!!(Setfield.@set(vi.values = y), lp_new)
+    return settrans!!(vi_new, t)
+end
+
+function invlink!!(
+    t::BijectorTransformation{<:Bijectors.NamedBijector},
+    vi::SimpleVarInfo{<:NamedTuple},
+    spl::AbstractSampler,
+    model::Model,
+)
+    # TODO: Make sure that `spl` is respected.
+    b = t.bijector
+    ib = inverse(b)
+    y = vi.values
+    x, logjac = with_logabsdet_jacobian(ib, y)
+    lp_new = getlogp(vi) - logjac
+    vi_new = setlogp!!(Setfield.@set(vi.values = x), lp_new)
+    return settrans!!(vi_new, NoTransformation())
+end

src/transforming.jl

@@ -0,0 +1,98 @@
+function Bijectors.Stacked(
+    model::Model,
+    ::Val{sym2ranges}=Val(false);
+    varinfo::VarInfo=VarInfo(model),
+) where {sym2ranges}
+    dists = vcat([varinfo.metadata[sym].dists for sym in keys(varinfo.metadata)]...)
+
+    num_ranges = sum([
+        length(varinfo.metadata[sym].ranges) for sym in keys(varinfo.metadata)
+    ])
+    ranges = Vector{UnitRange{Int}}(undef, num_ranges)
+    idx = 0
+    range_idx = 1
+
+    # ranges might be discontinuous => values are vectors of ranges rather than just ranges
+    sym_lookup = Dict{Symbol,Vector{UnitRange{Int}}}()
+    for sym in keys(varinfo.metadata)
+        sym_lookup[sym] = Vector{UnitRange{Int}}()
+        for r in varinfo.metadata[sym].ranges
+            ranges[range_idx] = idx .+ r
+            push!(sym_lookup[sym], ranges[range_idx])
+            range_idx += 1
+        end
+
+        idx += varinfo.metadata[sym].ranges[end][end]
+    end
+
+    b = Bijectors.Stacked(map(Bijectors.bijector, dists), ranges)
+    return sym2ranges ? (b, Dict(zip(keys(sym_lookup), values(sym_lookup)))) : b
+end
+
+link!!(vi::AbstractVarInfo, model::Model) = link!!(vi, SampleFromPrior(), model)
+function link!!(t::AbstractTransformation, vi::AbstractVarInfo, model::Model)
+    return link!!(t, vi, SampleFromPrior(), model)
+end
+function link!!(vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
+    # Use `default_transformation` to decide which transformation to use if none is specified.
+    return link!!(default_transformation(model, vi), vi, spl, model)
+end
+function link!!(
+    t::DefaultTransformation, vi::AbstractVarInfo, spl::AbstractSampler, model::Model
+)
+    # TODO: Implement this properly, e.g. using a context or something.
+    # Fall back to `Bijectors.Stacked` but then we act like we're using
+    # the `DefaultTransformation` by setting the transformation accordingly.
+    return settrans!!(
+        link!!(BijectorTransformation(Bijectors.Stacked(model)), vi, spl, model), t
+    )
+end
+function link!!(t::DefaultTransformation, vi::VarInfo, spl::AbstractSampler, model::Model)
+    # TODO: Implement this properly, e.g. using a context or something.
+    DynamicPPL.link!(vi, spl)
+    # TODO: Add `logabsdet_jacobian` correction to `logp`!
+    return vi
+end
+function link!!(
+    t::BijectorTransformation, vi::AbstractVarInfo, spl::AbstractSampler, model::Model
+)
+    b = t.bijector
+    x = vi[spl]
+    y, logjac = with_logabsdet_jacobian(b, x)
+
+    lp_new = getlogp(vi) - logjac
+    vi_new = setlogp!!(unflatten(vi, spl, y), lp_new)
+    return settrans!!(vi_new, t)
+end
+
+invlink!!(vi::AbstractVarInfo, model::Model) = invlink!!(vi, SampleFromPrior(), model)
+function invlink!!(t::AbstractTransformation, vi::AbstractVarInfo, model::Model)
+    return invlink!!(t, vi, SampleFromPrior(), model)
+end
+function invlink!!(vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
+    # Here we extract the `transformation` from `vi` rather than using the default one.
+    return invlink!!(transformation(vi), vi, spl, model)
+end
+function invlink!!(
+    ::DefaultTransformation, vi::AbstractVarInfo, spl::AbstractSampler, model::Model
+)
+    # TODO: Implement this properly, e.g. using a context or something.
+    return invlink!!(BijectorTransformation(Bijectors.Stacked(model)), vi, spl, model)
+end
+function invlink!!(::DefaultTransformation, vi::VarInfo, spl::AbstractSampler, model::Model)
+    # TODO: Implement this properly, e.g. using a context or something.
+    DynamicPPL.invlink!(vi, spl)
+    return vi
+end
+function invlink!!(
+    t::BijectorTransformation, vi::AbstractVarInfo, spl::AbstractSampler, model::Model
+)
+    b = t.bijector
+    ib = inverse(b)
+    y = vi[spl]
+    x, logjac = with_logabsdet_jacobian(ib, y)
+    # TODO: Do we need this?
+    lp_new = getlogp(vi) - logjac
+    vi_new = setlogp!!(unflatten(vi, spl, x), lp_new)
+    return settrans!!(vi_new, NoTransformation())
+end