Design of VarInfo

VarInfo is a fairly simple structure.

struct VarInfo{Tmeta, Tlogp} <: AbstractVarInfo
    metadata::Tmeta
    logp::Base.RefValue{Tlogp}
    num_produce::Base.RefValue{Int}
end

It contains

• a logp field for accumulation of the log-density evaluation, and
• a metadata field for storing information about the realizations of the different variables.

Representing logp is fairly straight-forward: we'll just use a Real or an array of Real, depending on the context.

Representing metadata is a bit trickier. This is supposed to contain all the necessary information for each VarName to enable the different executions of the model + extraction of different properties of interest after execution, e.g. the realization / value corresponding to a variable @varname(x).

To ensure that VarInfo is simple and intuitive to work with, we want VarInfo, and hence the underlying metadata, to replicate the following functionality of Dict:

• keys(::Dict): return all the VarNames present in metadata.
• haskey(::Dict): check if a particular VarName is present in metadata.
• getindex(::Dict, ::VarName): return the realization corresponding to a particular VarName.
• setindex!(::Dict, val, ::VarName): set the realization corresponding to a particular VarName.
• push!(::Dict, ::Pair): add a new key-value pair to the container.
• delete!(::Dict, ::VarName): delete the realization corresponding to a particular VarName.
• empty!(::Dict): delete all realizations in metadata.
• merge(::Dict, ::Dict): merge two metadata structures according to similar rules as Dict.

But for general-purpose samplers, we often want to work with a simple flattened structure, typically a Vector{<:Real}. One can access a vectorised version of a variable's value with the following vector-like functions:

• getindex_internal(::VarInfo, ::VarName): get the flattened value of a single variable.
• getindex_internal(::VarInfo, ::Colon): get the flattened values of all variables.
• getindex_internal(::VarInfo, i::Int): get ith value of the flattened vector of all values
• setindex_internal!(::VarInfo, ::AbstractVector, ::VarName): set the flattened value of a variable.
• setindex_internal!(::VarInfo, val, i::Int): set the ith value of the flattened vector of all values
• length_internal(::VarInfo): return the length of the flat representation of metadata.

The functions have _internal in their name because internally VarInfo always stores values as vectorised.

Moreover, a link transformation can be applied to a VarInfo with link!! (and reversed with invlink!!), which applies a reversible transformation to the internal storage format of a variable that makes the range of the random variable cover all of Euclidean space. getindex_internal and setindex_internal! give direct access to the vectorised value after such a transformation, which is what samplers often need to be able sample in unconstrained space. One can also manually set a transformation by giving setindex_internal! a fourth, optional argument, that is a function that maps internally stored value to the actual value of the variable.

Finally, we want want the underlying representation used in metadata to have a few performance-related properties:

1. Type-stable when possible, but functional when not.
2. Efficient storage and iteration when possible, but functional when not.

The "but functional when not" is important as we want to support arbitrary models, which means that we can't always have these performance properties.

In the following sections, we'll outline how we achieve this in VarInfo.

Type-stability

Ensuring type-stability is somewhat non-trivial to address since we want this to be the case even when models mix continuous (typically Float64) and discrete (typically Int) variables.

Suppose we have an implementation of metadata which implements the functionality outlined in the previous section. The way we approach this in VarInfo is to use a NamedTuple with a separate metadata for each distinct Symbol used. For example, if we have a model of the form

using DynamicPPL, Distributions, FillArrays

@model function demo()
    x ~ product_distribution(Fill(Bernoulli(0.5), 2))
    y ~ Normal(0, 1)
    return nothing
end

demo (generic function with 2 methods)

then we construct a type-stable representation by using a NamedTuple{(:x, :y), Tuple{Vx, Vy}} where

• Vx is a container with eltype Bool, and
• Vy is a container with eltype Float64.

Since VarName contains the Symbol used in its type, something like getindex(varinfo, @varname(x)) can be resolved to getindex(varinfo.metadata.x, @varname(x)) at compile-time.

For example, with the model above we have

# Type-unstable `VarInfo`
varinfo_untyped = DynamicPPL.untyped_varinfo(
    demo(), SampleFromPrior(), DefaultContext(), DynamicPPL.Metadata()
)
typeof(varinfo_untyped.metadata)

DynamicPPL.Metadata{Dict{VarName, Int64}, Vector{Distribution}, Vector{VarName}, Vector{Real}, Vector{Set{DynamicPPL.Selector}}}

# Type-stable `VarInfo`
varinfo_typed = DynamicPPL.typed_varinfo(demo())
typeof(varinfo_typed.metadata)

@NamedTuple{x::DynamicPPL.Metadata{Dict{VarName{:x, typeof(identity)}, Int64}, Vector{Product{Discrete, Bernoulli{Float64}, FillArrays.Fill{Bernoulli{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{VarName{:x, typeof(identity)}}, BitVector, Vector{Set{DynamicPPL.Selector}}}, y::DynamicPPL.Metadata{Dict{VarName{:y, typeof(identity)}, Int64}, Vector{Normal{Float64}}, Vector{VarName{:y, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}

They both work as expected but one results in concrete typing and the other does not:

varinfo_untyped[@varname(x)], varinfo_untyped[@varname(y)]

(Real[true, true], 0.06631845415545398)

varinfo_typed[@varname(x)], varinfo_typed[@varname(y)]

(Bool[1, 1], 0.4541584437699302)

Notice that the untyped VarInfo uses Vector{Real} to store the boolean entries while the typed uses Vector{Bool}. This is because the untyped version needs the underlying container to be able to handle both the Bool for x and the Float64 for y, while the typed version can use a Vector{Bool} for x and a Vector{Float64} for y due to its usage of NamedTuple.

x[1] ~ Bernoulli(0.5)
x[2] ~ Normal(0, 1)

Hence we obtain a "type-stable when possible"-representation by wrapping it in a NamedTuple and partially resolving the getindex, setindex!, etc. methods at compile-time. When type-stability is not desired, we can simply use a single metadata for all VarNames instead of a NamedTuple wrapping a collection of metadatas.

Efficient storage and iteration

Efficient storage and iteration we achieve through implementation of the metadata. In particular, we do so with DynamicPPL.VarNamedVector:

VarNamedVector

julia> using DynamicPPL: ReshapeTransform, VarNamedVector, @varname, setindex!, update!, getindex_internal

julia> vnv = VarNamedVector();

julia> setindex!(vnv, [0.0, 0.0, 0.0, 0.0], @varname(x));

julia> setindex!(vnv, reshape(1:6, (2,3)), @varname(y));

julia> vnv.vals
10-element Vector{Real}:
 0.0
 0.0
 0.0
 0.0
 1
 2
 3
 4
 5
 6

julia> vnv.transforms[vnv.varname_to_index[@varname(y)]] == DynamicPPL.ReshapeTransform((6,), (2,3))
true

julia> update!(vnv, [46.0, 48.0], @varname(x))

julia> vnv.vals
10-element Vector{Real}:
 46.0
 48.0
  0.0
  0.0
  1
  2
  3
  4
  5
  6

julia> println(vnv.num_inactive);
OrderedDict(1 => 2)

julia> vnv[@varname(x)]
2-element Vector{Float64}:
 46.0
 48.0

julia> getindex_internal(vnv, :)
8-element Vector{Real}:
 46.0
 48.0
  1
  2
  3
  4
  5
  6

In a DynamicPPL.VarNamedVector{<:VarName,T}, we achieve the desiderata by storing the values for different VarNames contiguously in a Vector{T} and keeping track of which ranges correspond to which VarNames.

This does require a bit of book-keeping, in particular when it comes to insertions and deletions. Internally, this is handled by assigning each VarName a unique Int index in the varname_to_index field, which is then used to index into the following fields:

• varnames::Vector{<:VarName}: the VarNames in the order they appear in the Vector{T}.
• ranges::Vector{UnitRange{Int}}: the ranges of indices in the Vector{T} that correspond to each VarName.
• transforms::Vector: the transforms associated with each VarName.

Mutating functions, e.g. setindex_internal!(vnv::VarNamedVector, val, vn::VarName), are then treated according to the following rules:

1. If vn is not already present: add it to the end of vnv.varnames, add the val to the underlying vnv.vals, etc.

2. If vn is already present in vnv:

   1. If val has the same length as the existing value for vn: replace existing value.
   2. If val has a smaller length than the existing value for vn: replace existing value and mark the remaining indices as "inactive" by increasing the entry in vnv.num_inactive field.
   3. If val has a larger length than the existing value for vn: expand the underlying vnv.vals to accommodate the new value, update all VarNames occuring after vn, and update the vnv.ranges to point to the new range for vn.

This means that VarNamedVector is allowed to grow as needed, while "shrinking" (i.e. insertion of smaller elements) is handled by simply marking the redundant indices as "inactive". This turns out to be efficient for use-cases that we are generally interested in.

For example, we want to optimize code-paths which effectively boil down to inner-loop in the following example:

# Construct a `VarInfo` with types inferred from `model`.
varinfo = VarInfo(model)

# Repeatedly sample from `model`.
for _ in 1:num_samples
    rand!(rng, model, varinfo)

    # Do something with `varinfo`.
    # ...
end

There are typically a few scenarios where we encounter changing representation sizes of a random variable x:

1. We're working with a transformed version x which is represented in a lower-dimensional space, e.g. transforming a x ~ LKJ(2, 1) to unconstrained y = f(x) takes us from 2-by-2 Matrix{Float64} to a 1-length Vector{Float64}.
2. x has a random size, e.g. in a mixture model with a prior on the number of components. Here the size of x can vary widly between every realization of the Model.

In scenario (1), we're usually shrinking the representation of x, and so we end up not making any allocations for the underlying Vector{T} but instead just marking the redundant part as "inactive".

In scenario (2), we end up increasing the allocated memory for the randomly sized x, eventually leading to a vector that is large enough to hold realizations without needing to reallocate. But this can still lead to unnecessary memory usage, which might be undesirable. Hence one has to make a decision regarding the trade-off between memory usage and performance for the use-case at hand.

To help with this, we have the following functions:

has_inactive(vnv::VarNamedVector)

num_inactive(vnv::VarNamedVector)

num_inactive(vnv::VarNamedVector, vn::VarName)

num_allocated(vnv::VarNamedVector)
num_allocated(vnv::VarNamedVector[, vn::VarName])
num_allocated(vnv::VarNamedVector[, idx::Int])

is_contiguous(vnv::VarNamedVector)

contiguify!(vnv::VarNamedVector)

julia> using DynamicPPL: VarNamedVector, @varname, contiguify!, update!, has_inactive

julia> vnv = VarNamedVector(@varname(x) => [1.0, 2.0, 3.0], @varname(y) => [3.0]);

julia> update!(vnv, [23.0, 24.0], @varname(x));

julia> has_inactive(vnv)
true

julia> length(vnv.vals)
4

julia> contiguify!(vnv);

julia> has_inactive(vnv)
false

julia> length(vnv.vals)
3

julia> vnv[@varname(x)]  # All the values are still there.
2-element Vector{Float64}:
 23.0
 24.0

For example, one might encounter the following scenario:

vnv = DynamicPPL.VarNamedVector(@varname(x) => [true])
println("Before insertion: number of allocated entries  $(DynamicPPL.num_allocated(vnv))")

for i in 1:5
    x = fill(true, rand(1:100))
    DynamicPPL.update!(vnv, x, @varname(x))
    println(
        "After insertion #$(i) of length $(length(x)): number of allocated entries  $(DynamicPPL.num_allocated(vnv))",
    )
end

Before insertion: number of allocated entries  1
After insertion #1 of length 37: number of allocated entries  37
After insertion #2 of length 16: number of allocated entries  37
After insertion #3 of length 86: number of allocated entries  86
After insertion #4 of length 21: number of allocated entries  86
After insertion #5 of length 8: number of allocated entries  86

We can then insert a call to DynamicPPL.contiguify! after every insertion whenever the allocation grows too large to reduce overall memory usage:

vnv = DynamicPPL.VarNamedVector(@varname(x) => [true])
println("Before insertion: number of allocated entries  $(DynamicPPL.num_allocated(vnv))")

for i in 1:5
    x = fill(true, rand(1:100))
    DynamicPPL.update!(vnv, x, @varname(x))
    if DynamicPPL.num_allocated(vnv) > 10
        DynamicPPL.contiguify!(vnv)
    end
    println(
        "After insertion #$(i) of length $(length(x)): number of allocated entries  $(DynamicPPL.num_allocated(vnv))",
    )
end

Before insertion: number of allocated entries  1
After insertion #1 of length 48: number of allocated entries  48
After insertion #2 of length 86: number of allocated entries  86
After insertion #3 of length 22: number of allocated entries  22
After insertion #4 of length 40: number of allocated entries  40
After insertion #5 of length 91: number of allocated entries  91

This does incur a runtime cost as it requires re-allocation of the ranges in addition to a resize! of the underlying Vector{T}. However, this also ensures that the the underlying Vector{T} is contiguous, which is important for performance. Hence, if we're about to do a lot of work with the VarNamedVector without insertions, etc., it can be worth it to do a sweep to ensure that the underlying Vector{T} is contiguous.

Continuing from the example from the previous section, we can use a VarInfo with a VarNamedVector as the metadata field:

# Type-unstable
varinfo_untyped_vnv = DynamicPPL.VectorVarInfo(varinfo_untyped)
varinfo_untyped_vnv[@varname(x)], varinfo_untyped_vnv[@varname(y)]

(Real[true, true], 0.06631845415545398)

# Type-stable
varinfo_typed_vnv = DynamicPPL.VectorVarInfo(varinfo_typed)
varinfo_typed_vnv[@varname(x)], varinfo_typed_vnv[@varname(y)]

(Bool[1, 1], 0.4541584437699302)

If we now try to delete! @varname(x)

haskey(varinfo_untyped_vnv, @varname(x))

true

DynamicPPL.has_inactive(varinfo_untyped_vnv.metadata)

false

# `delete!`
DynamicPPL.delete!(varinfo_untyped_vnv.metadata, @varname(x))
DynamicPPL.has_inactive(varinfo_untyped_vnv.metadata)

false

haskey(varinfo_untyped_vnv, @varname(x))

false

Or insert a differently-sized value for @varname(x)

DynamicPPL.insert!(varinfo_untyped_vnv.metadata, fill(true, 1), @varname(x))
varinfo_untyped_vnv[@varname(x)]

1-element Vector{Real}:
 true

DynamicPPL.num_allocated(varinfo_untyped_vnv.metadata, @varname(x))

1

DynamicPPL.update!(varinfo_untyped_vnv.metadata, fill(true, 4), @varname(x))
varinfo_untyped_vnv[@varname(x)]

4-element Vector{Real}:
 true
 true
 true
 true

DynamicPPL.num_allocated(varinfo_untyped_vnv.metadata, @varname(x))

4

Performance summary

In the end, we have the following "rough" performance characteristics for VarNamedVector:

Other methods

replace_raw_storage(vnv::VarNamedVector, vals::AbstractVector)

julia> using DynamicPPL: VarNamedVector, replace_raw_storage

julia> vnv = VarNamedVector(@varname(x) => [1.0]);

julia> replace_raw_storage(vnv, [2.0])[@varname(x)] == [2.0]
true

julia> using ForwardDiff: ForwardDiff

julia> f(x) = sum(abs2, replace_raw_storage(vnv, x)[@varname(x)])
f (generic function with 1 method)

julia> ForwardDiff.gradient(f, [1.0])
1-element Vector{Float64}:
 2.0

values_as(vnv::VarNamedVector[, T])

julia> using DynamicPPL: VarNamedVector

julia> vnv = VarNamedVector(@varname(x) => 1, @varname(y) => [2.0]);

julia> values_as(vnv) == [1.0, 2.0]
true

julia> values_as(vnv, Vector{Float32}) == Vector{Float32}([1.0, 2.0])
true

julia> values_as(vnv, OrderedDict) == OrderedDict(@varname(x) => 1.0, @varname(y) => [2.0])
true

julia> values_as(vnv, NamedTuple) == (x = 1.0, y = [2.0])
true