Storing vectorised vs. raw values

In the Existing accumulators page, we mentioned that there are two accumulators that store values, namely RawValueAccumulator and VectorValueAccumulator. This page attempts to explain the rationale behind this, when each of these needs to be used, and how we could potentially remove the need for VectorValueAccumulator in the future.

Before we go into this, we need to start by talking about VarInfo, which we have not covered at all yet.

The role of VarInfo

As described in the model evaluation documentation page, each tilde-statement is split up into three parts:

1. Initialisation;
2. Transformation; and
3. Accumulation.

Unfortunately, not everything in DynamicPPL follows this clean structure yet. In particular, there is a struct, called VarInfo:

struct VarInfo{Tfm<:AbstractTransformStrategy,V<:VarNamedTuple,A<:AccumulatorTuple}
    transform_strategy::Tfm
    values::V
    accs::A
end

The values field stores either LinkedVectorValues or VectorValues. The transform_strategy field stores an AbstractTransformStrategy which is (as far as possible) consistent with the type of values stored in values.

Here is an example:

using DynamicPPL, Distributions

@model function dirichlet()
    x = zeros(3)
    return x[1:3] ~ Dirichlet(ones(3))
end
dirichlet_model = dirichlet()
vi = VarInfo(dirichlet_model)
vi

VarInfo
 ├─ transform_strategy: UnlinkAll()
 ├─ values
 │  VarNamedTuple
 │  └─ x => PartialArray size=(3,) data::Vector{DynamicPPL.VarNamedTuples.ArrayLikeBlock{VectorValue{Vector{Float64}, typeof(identity), Tuple{Int64}}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}}
 │          ├─ (1,) => DynamicPPL.VarNamedTuples.ArrayLikeBlock{VectorValue{Vector{Float64}, typeof(identity), Tuple{Int64}}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}(VectorValue{Vector{Float64}, typeof(identity), Tuple{Int64}}([0.14920895846702054, 0.0517852481449895, 0.7990057933879898], identity, (3,)), (1:3,), NamedTuple(), (3,))
 │          ├─ (2,) => DynamicPPL.VarNamedTuples.ArrayLikeBlock{VectorValue{Vector{Float64}, typeof(identity), Tuple{Int64}}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}(VectorValue{Vector{Float64}, typeof(identity), Tuple{Int64}}([0.14920895846702054, 0.0517852481449895, 0.7990057933879898], identity, (3,)), (1:3,), NamedTuple(), (3,))
 │          └─ (3,) => DynamicPPL.VarNamedTuples.ArrayLikeBlock{VectorValue{Vector{Float64}, typeof(identity), Tuple{Int64}}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}(VectorValue{Vector{Float64}, typeof(identity), Tuple{Int64}}([0.14920895846702054, 0.0517852481449895, 0.7990057933879898], identity, (3,)), (1:3,), NamedTuple(), (3,))
 └─ accs
    AccumulatorTuple with 3 accumulators
    ├─ LogPrior => LogPriorAccumulator(0.6931471805599453)
    ├─ LogJacobian => LogJacobianAccumulator(0.0)
    └─ LogLikelihood => LogLikelihoodAccumulator(0.0)

In VarInfo, it is mandatory to store LinkedVectorValues or VectorValues as ArrayLikeBlocks (see the Array-like blocks documentation for information on this). The reason is because, if the value is linked, it may have a different size than the number of indices in the VarName. This means that when retrieving the keys, we obtain each block as a single key:

keys(vi.values)

1-element Vector{VarName}:
 x[1:3]

Towards a new framework

In a VarInfo, the accs field is responsible for the accumulation step, just like an ordinary AccumulatorTuple.

However, values serves three purposes in one:

• it is sometimes used for initialisation (when the model's leaf context is DefaultContext, the AbstractTransformedValue to be used in the transformation step is read from it)
• it also determines whether the log-Jacobian term should be included or not (if the value is a LinkedVectorValue, the log-Jacobian is included)
• it is sometimes also used for accumulation (when evaluating a model with a VarInfo, we will potentially store a new AbstractTransformedValue in it!).

The path to removing VarInfo is essentially to separate these three roles:

1. The initialisation role of varinfo.values can be taken over by an initialisation strategy that wraps it. Recall that the only role of an initialisation strategy is to provide an AbstractTransformedValue via DynamicPPL.init. This can be trivially done by indexing into the VarNamedTuple stored in the strategy.

2. Whether the log-Jacobian term should be included or not can be determined by a transform strategy. Much like how we can have an initialisation strategy that takes values from a VarInfo, we can also have a transform strategy that is defined by the existing status of a VarInfo. This is implemented in the DynamicPPL.get_link_strategy(::AbstractVarInfo) function.

3. The accumulation role of varinfo.values can be taken over by a new accumulator, which we call VectorValueAccumulator. This name is chosen because it does not store generic AbstractTransformedValues, but only two subtypes of it, LinkedVectorValue and VectorValue. VectorValueAccumulator is implemented inside src/accs/vector_value.jl.

RawValueAccumulator

Earlier we said that VectorValueAccumulator stores only two subtypes of AbstractTransformedValue: LinkedVectorValue and VectorValue. One might therefore ask about the third subtype, namely, UntransformedValue.

It turns out that it is very often useful to store UntransformedValues. Additionally, since UntransformedValues must always correspond exactly to the indices they are assigned to, we can unwrap them and do not need to store them as array-like blocks!

This is the role of RawValueAccumulator.

oavi = DynamicPPL.OnlyAccsVarInfo(DynamicPPL.RawValueAccumulator(false))
_, oavi = DynamicPPL.init!!(dirichlet_model, oavi, InitFromPrior(), UnlinkAll())
raw_vals = get_raw_values(oavi)

VarNamedTuple
└─ x => PartialArray size=(3,) data::Vector{Float64}
        ├─ (1,) => 0.3127737487043817
        ├─ (2,) => 0.3648323204580915
        └─ (3,) => 0.32239393083752677

Note that when we unwrap UntransformedValues, we also lose the block structure that was present in the model. That means that in RawValueAccumulator, there is no longer any notion that x[1:3] was set together, so the keys correspond to the individual indices.

keys(raw_vals)

3-element Vector{VarName}:
 x[1]
 x[2]
 x[3]

In particular, the outputs of RawValueAccumulator are used for chain construction. This is why indices of keys like x[1:3] ~ dist end up being split up in chains.

Why do we still need to store TransformedValues?

Given that RawValueAccumulator exists, one may wonder why we still need to store the other AbstractTransformedValues at all, i.e. what the purpose of VectorValueAccumulator is.

Currently, the only remaining reason for transformed values is the fact that we may sometimes need to perform DynamicPPL.unflatten!! on a VarInfo, to insert new values into it from a vector.

vi = VarInfo(dirichlet_model)
vi[@varname(x[1:3])]

3-element Vector{Float64}:
 0.1263750296477373
 0.39656291979385505
 0.47706205055840767

vi = DynamicPPL.unflatten!!(vi, [0.2, 0.5, 0.3])
vi[@varname(x[1:3])]

3-element view(::Vector{Float64}, 1:3) with eltype Float64:
 0.2
 0.5
 0.3

If we do not store the vectorised form of the values, we will not know how many values to read from the input vector for each key.

Removing upstream usage of unflatten!! would allow us to completely get rid of TransformedValueAccumulator and only ever use RawValueAccumulator. See this DynamicPPL issue for more information.

One possibility for removing unflatten!! is to turn it into a function that, instead of generating a new VarInfo, instead generates a tuple of new initialisation and link strategies which returns LinkedVectorValues or VectorValues containing views into the input vector. This would be conceptually very similar to how LogDensityFunction currently works.