Array-like blocks

In a number of VNT use cases, it is necessary to associate multiple indices in a VarNamedTuple with an object that is not necessarily the same number of elements.

Consider, for example, this model:

@model function f()
    x = zeros(3)
    return x[1:3] ~ Dirichlet(ones(3))
end

and suppose we want to store the prior distribution associated with each variable in a VarNamedTuple. With a Dict{VarName,Distribution}, we can do this:

d = Dict{VarName,Distribution}(@varname(x[1:3]) => Dirichlet(ones(3)))

but we incur all the costs associated with the use of a Dict, as described before.

With a VarNamedTuple, we cannot store this directly:

vnt.data.x = ... # some array
vnt.data.x[1:3] = Dirichlet(ones(3))  # will error

because Dirichlet is not an array, and setindex! will fail. Nor can we write

vnt.data.x = ... # some array
vnt.data.x[1:3] .= Dirichlet(ones(3))

because although this will not error, it is semantically different: this means that every element x[1], x[2], and x[3] will be assigned the same Dirichlet(ones(3)) object, which is not what we want.

The current solution to this is to use ArrayLikeBlocks, which are thin wrappers around the actual value, but additionally also store the indices used to set the value. The second and third arguments here are the indices (positional and keyword) used to set the value, and the fourth argument is the size of the block.

using DynamicPPL, Distributions
using DynamicPPL.VarNamedTuples: ArrayLikeBlock

alb = ArrayLikeBlock(Dirichlet(ones(3)), 1:3, (;), (3,))

DynamicPPL.VarNamedTuples.ArrayLikeBlock{Dirichlet{Float64, Vector{Float64}, Float64}, UnitRange{Int64}, @NamedTuple{}, Tuple{Int64}}(Dirichlet{Float64, Vector{Float64}, Float64}(alpha=[1.0, 1.0, 1.0]), 1:3, NamedTuple(), (3,))

We then set this ArrayLikeBlock in all the relevant indices. The extra information in the ArrayLikeBlock allows us to forbid partial indexing into it later on. In particular, we want to ensure that users can only retrieve the entire block at once, and not e.g. just x[1] or x[2:3].

Getting and setting: in practice

As a user you should not have to deal with ArrayLikeBlocks directly. Under the hood, templated_setindex!! will automatically wrap values in ArrayLikeBlocks when necessary:

x = zeros(5)
vnt = DynamicPPL.templated_setindex!!(
    VarNamedTuple(), Dirichlet(ones(3)), @varname(x[1:3]), x
)

VarNamedTuple
└─ x => PartialArray size=(5,) data::Vector{DynamicPPL.VarNamedTuples.ArrayLikeBlock{Dirichlet{Float64, Vector{Float64}, Float64}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}}
        ├─ (1,) => DynamicPPL.VarNamedTuples.ArrayLikeBlock{Dirichlet{Float64, Vector{Float64}, Float64}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}(Dirichlet{Float64, Vector{Float64}, Float64}(alpha=[1.0, 1.0, 1.0]), (1:3,), NamedTuple(), (3,))
        ├─ (2,) => DynamicPPL.VarNamedTuples.ArrayLikeBlock{Dirichlet{Float64, Vector{Float64}, Float64}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}(Dirichlet{Float64, Vector{Float64}, Float64}(alpha=[1.0, 1.0, 1.0]), (1:3,), NamedTuple(), (3,))
        └─ (3,) => DynamicPPL.VarNamedTuples.ArrayLikeBlock{Dirichlet{Float64, Vector{Float64}, Float64}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}(Dirichlet{Float64, Vector{Float64}, Float64}(alpha=[1.0, 1.0, 1.0]), (1:3,), NamedTuple(), (3,))

You can access the value again as long as you refer to the full range:

vnt[@varname(x[1:3])]

Dirichlet{Float64, Vector{Float64}, Float64}(alpha=[1.0, 1.0, 1.0])

Because we provided template information, you can access this via any other combination of indexing, as long as it refers to all three indices:

vnt[@varname(x[begin:(end - 2)])]

Dirichlet{Float64, Vector{Float64}, Float64}(alpha=[1.0, 1.0, 1.0])

However, if you try to access only part of the block, you will get an error:

julia> vnt[@varname(x[1])]ERROR: ArgumentError: A non-Array value set with a range of indices must be retrieved with the same
range of indices.

Furthermore, if you set a value into any of the indices covered by the block, the entire block is invalidated and thus removed:

vnt = DynamicPPL.templated_setindex!!(vnt, Normal(), @varname(x[2]), x)

VarNamedTuple
└─ x => PartialArray size=(5,) data::Vector{Normal{Float64}}
        └─ (2,) => Normal{Float64}(μ=0.0, σ=1.0)

Size checks

Currently, when setting any object val as an ArrayLikeBlock, there is a size check: we make sure that the range of indices being set to has the same size as DynamicPPL.VarNamedTuples.vnt_size(val). By default, vnt_size(x) returns Base.size(x).

DynamicPPL.VarNamedTuples.vnt_size(Dirichlet(ones(3)))

(3,)

This is what allows us to set a Dirichlet distribution to three indices. However, trying to set the same distribution to two indices will fail:

julia> vnt = DynamicPPL.templated_setindex!!(
           VarNamedTuple(), Dirichlet(ones(3)), @varname(x[1:2]), zeros(5)
       )ERROR: DimensionMismatch: Assigned value has size (3,), which does not match the size implied by the indices (2,).

Which parts of DynamicPPL use array-like blocks?

Simply put, anywhere where we don't store raw values. Some examples follow. (This list is meant to only be illustrative, not exhaustive!)

VarInfo

In VarInfo, we need to be able to store either linked or unlinked values (in general, AbstractTransformedValues). These are always vectorised values, and the linked and unlinked vectors may have different sizes (this is indeed the case for Dirichlet distributions). This means that we have to collectively assign multiple indices in the VarNamedTuple to a single vector, which may or may not have the same size as the indices.

@model function dirichlet()
    x = zeros(3)
    return x[1:3] ~ Dirichlet(ones(3))
end
dirichlet_model = dirichlet()
vi = VarInfo(dirichlet_model)
vi.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.33574003085380216, 0.004627580598464074, 0.6596323885477337], 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.33574003085380216, 0.004627580598464074, 0.6596323885477337], 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.33574003085380216, 0.004627580598464074, 0.6596323885477337], identity, (3,)), (1:3,), NamedTuple(), (3,))

This means that in the actual VarInfo we do not have a notion of what x[1] is:

julia> vi[@varname(x[1])]ERROR: ArgumentError: A non-Array value set with a range of indices must be retrieved with the same
range of indices.

See the documentation on storing values for more details.

Prior distributions

In extract_priors we use a VNT to store the prior distributions seen at each point in the model. This is exactly the same use case as in the introduction.

DynamicPPL.extract_priors(dirichlet_model)

VarNamedTuple
└─ x => PartialArray size=(3,) data::Vector{DynamicPPL.VarNamedTuples.ArrayLikeBlock{Dirichlet{Float64, Vector{Float64}, Float64}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}}
        ├─ (1,) => DynamicPPL.VarNamedTuples.ArrayLikeBlock{Dirichlet{Float64, Vector{Float64}, Float64}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}(Dirichlet{Float64, Vector{Float64}, Float64}(alpha=[1.0, 1.0, 1.0]), (1:3,), NamedTuple(), (3,))
        ├─ (2,) => DynamicPPL.VarNamedTuples.ArrayLikeBlock{Dirichlet{Float64, Vector{Float64}, Float64}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}(Dirichlet{Float64, Vector{Float64}, Float64}(alpha=[1.0, 1.0, 1.0]), (1:3,), NamedTuple(), (3,))
        └─ (3,) => DynamicPPL.VarNamedTuples.ArrayLikeBlock{Dirichlet{Float64, Vector{Float64}, Float64}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}(Dirichlet{Float64, Vector{Float64}, Float64}(alpha=[1.0, 1.0, 1.0]), (1:3,), NamedTuple(), (3,))

LogDensityFunction

DynamicPPL.LogDensityFunction has to retain information about whether each VarNames is linked or not, and the indices in the vectorised parameters that correspond to each variable.

For example, consider creating a linked LogDensityFunction for (a slightly expanded version of) the Dirichlet model above:

@model function expanded_dirichlet()
    x = zeros(4)
    x[1] ~ Normal()
    return x[2:4] ~ Dirichlet(ones(3))
end
model = expanded_dirichlet()

vi = VarInfo(model)
linked_vi = DynamicPPL.link!!(vi, model)
ldf = LogDensityFunction(model, DynamicPPL.getlogjoint_internal, linked_vi)

When linking a Dirichlet(ones(3)), the resulting vector will have length 2. So, the LogDensityFunction takes in a vector of length 3, for which the first index belongs to x[1], and the next two indices belong to the linked parameters for x[2:4]. Furthermore, it needs to remember that all the variables have been linked. We can see that this is exactly how the LogDensityFunction has stored the information:

ldf._varname_ranges

VarNamedTuple
└─ x => PartialArray size=(4,) data::Vector{Any}
        ├─ (1,) => DynamicPPL.RangeAndLinked{Tuple{}}(1:1, true, ())
        ├─ (2,) => DynamicPPL.VarNamedTuples.ArrayLikeBlock{DynamicPPL.RangeAndLinked{Tuple{Int64}}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}(DynamicPPL.RangeAndLinked{Tuple{Int64}}(2:3, true, (3,)), (2:4,), NamedTuple(), (3,))
        ├─ (3,) => DynamicPPL.VarNamedTuples.ArrayLikeBlock{DynamicPPL.RangeAndLinked{Tuple{Int64}}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}(DynamicPPL.RangeAndLinked{Tuple{Int64}}(2:3, true, (3,)), (2:4,), NamedTuple(), (3,))
        └─ (4,) => DynamicPPL.VarNamedTuples.ArrayLikeBlock{DynamicPPL.RangeAndLinked{Tuple{Int64}}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}(DynamicPPL.RangeAndLinked{Tuple{Int64}}(2:3, true, (3,)), (2:4,), NamedTuple(), (3,))