Need alternative to NamedTuple for SimpleVarInfo
#528
Comments
|
Dictionaries.jl will unfortunately not give us contiguous memory (at least not by default): julia> using Dictionaries
julia> vnv = Dictionary{VarName}(OrderedDict(@varname(a) => [1.0,2,3], @varname(b) => [4,5,6.]))
2-element Dictionary{VarName, Vector{Float64}}
a │ [1.0, 2.0, 3.0]
b │ [4.0, 5.0, 6.0]
julia> vnv.values
2-element Vector{Vector{Float64}}:
[1.0, 2.0, 3.0]
[4.0, 5.0, 6.0] |
|
How much does contiguous memory actually matter? I would imagine CPU cache fills in a lot of the performance gap. |
|
Well, it depends. Generally speaking, I'm expecting the access patterns here to be views mainly, and so if you then reshape these into larger matrices which is then used for downstream computations, then we'd expect it to make a noticeable difference, no? But no matter what we do, we should benchmark these things at some point. |
|
And rolling our own shouldn't be too difficult. For example, we could use something like """
VarNameDict
A `VarNameDict` is a vector-like collection of values that can be indexed by `VarName`.
This is basically like a `OrderedDict{<:VarName}` but ensures that the underlying values
are stored contiguously in memory.
"""
struct VarNameDict{K,T,V<:AbstractVector{T},D<:AbstractDict{K}} <: AbstractDict{K,T}
values::V
varname_to_ranges::D
end
function VarNameDict(dict::OrderedDict)
offset = 0
ranges = map(values(dict)) do x
r = (offset + 1):(offset + length(x))
offset = r[end]
r
end
vals = mapreduce(DynamicPPL.vectorize, vcat, values(dict))
return VarNameDict(vals, OrderedDict(zip(keys(dict), ranges)))
end
# Dict-like functionality.
Base.keys(vnd::VarNameDict) = keys(vnd.varname_to_ranges)
Base.values(vnd::VarNameDict) = vnd.values
Base.length(vnd::VarNameDict) = length(vnd.values)
Base.getindex(vnd::VarNameDict, i) = getindex(vnd.values, i)
Base.setindex!(vnd::VarNameDict, val, i) = setindex!(vnd.values, val, i)
function nextrange(vnd::VarNameDict, x)
n = length(vnd)
return n + 1:n + length(x)
end
function Base.getindex(vnd::VarNameDict, vn::VarName)
return getindex(vnd.values, vnd.varname_to_ranges[vn])
end
function Base.setindex!(vnd::VarNameDict, val, vn::VarName)
# If we don't have `vn` in the dictionary, then we need to add it.
if !haskey(vnd.varname_to_ranges, vn)
# Set the range for the new variable.
r = nextrange(vnd, val)
vnd.varname_to_ranges[vn] = r
# Resize the underlying vector to accommodate the new values.
resize!(vnd.values, r[end])
else
# Existing keys needs to be handled differently depending on
# whether the size of the value is increasing or decreasing.
r = vnd.varname_to_ranges[vn]
n_val = length(val)
n_r = length(r)
if n_val > n_r
# Remove the old range.
delete!(vnd.varname_to_ranges, vn)
# Add the new range.
r_new = nextrange(vnd, val)
vnd.varname_to_ranges[vn] = r_new
# Resize the underlying vector to accommodate the new values.
resize!(vnd.values, r_new[end])
else n_val < n_r
# Just decrease the current range.
vnd.varname_to_ranges[vn] = r[1]:r[1] + n_val - 1
end
# TODO: Keep track of unused ranges so we can perform sweeps
# every now and then to free up memory and re-contiguize the
# underlying vector.
end
return setindex!(vnd.values, val, vnd.varname_to_ranges[vn])
end
function BangBang.setindex!!(vnd::VarNameDict, val, vn::VarName)
setindex!(vnd, val, vn)
return vnd
end
function Base.iterate(vnd::VarNameDict, state=nothing)
res = state === nothing ? iterate(vnd.varname_to_ranges) : iterate(vnd.varname_to_ranges, state)
res === nothing && return nothing
(vn, range), state_new = res
return vn => vnd.values[range], state_new
endAdding this to DPPL and exporting: julia> using DynamicPPL
julia> model = DynamicPPL.TestUtils.demo_one_variable_multiple_constraints()
Model{typeof(DynamicPPL.TestUtils.demo_one_variable_multiple_constraints), (Symbol("##arg#375"),), (), (), Tuple{DataType}, Tuple{}, DefaultContext}(DynamicPPL.TestUtils.demo_one_variable_multiple_constraints, (var"##arg#375" =
Vector{Float64},), NamedTuple(), DefaultContext())
julia> x = rand(OrderedDict, model)
OrderedDict{Any, Any} with 4 entries:
x[1] => 1.18456
x[2] => 1.83516
x[3] => 0.726668
x[4:5] => [0.191595, 0.808405]
julia> vnd = VarNameDict(x)
VarNameDict{Any, Float64, Vector{Float64}, OrderedDict{Any, UnitRange{Int64}}} with 5 entries:
x[1] => [1.18456]
x[2] => [1.83516]
x[3] => [0.726668]
x[4:5] => [0.191595, 0.808405]
julia> vnd[@varname(x[1])]
1-element Vector{Float64}:
1.1845589710704487
julia> vnd[@varname(x[4:5])]
2-element Vector{Float64}:
0.19159541239321576
0.8084045876067844
julia> # Can create a `SimpleVarInfo` from a `VarNameDict`.
vi = SimpleVarInfo(vnd)
SimpleVarInfo(VarNameDict{Any, Float64, Vector{Float64}, OrderedDict{Any, UnitRange{Int64}}}(x[1] => [1.1845589710704487], x[2] => [1.8351581568345814], x[3] => [0.726667648763101], x[4:5] => [0.19159541239321576, 0.8084045876067844]), 0.0)
julia> # Inherits from `AbstractDict`
vi[@varname(x[1])]
1-element Vector{Float64}:
1.1845589710704487
julia> vi[@varname(x[4:5])]
2-element Vector{Float64}:
0.19159541239321576
0.8084045876067844
julia> vi[@varname(x[4:5][1])]
0.19159541239321576
julia> vi_linked = link(vi, model)
Transformed SimpleVarInfo(VarNameDict{Any, Float64, Vector{Float64}, OrderedDict{Any, UnitRange{Int64}}}(x[1] => [1.1845589710704487], x[2] => [0.6071306668031453], x[3] => [-1.2135885975805698], x[4:5] => [-1.4396767388530345]), -3.354889980209639)
julia> vi_invlinked = invlink(vi_linked, model)
SimpleVarInfo(VarNameDict{Any, Float64, Vector{Float64}, OrderedDict{Any, UnitRange{Int64}}}(x[1] => [1.1845589710704487], x[2] => [1.8351581568345814], x[3] => [0.726667648763101], x[4:5] => [0.19159541239321573, 0.8084045876067842]), -3.5819390425447497)Uncertain if it's really worth it to make it a |
|
I like the EDIT: reserving a field for some metadata, other than |
|
I suggest that we rename |
Out of curiosity, what's the issue preventing us from using |
Indexing using |
I also thought about this but didn't want to add it in the initial version. We could do this, yeah, but it might also just be best to be kept separately in the varinfo? A bit uncertain. |
Keeping it inside |
|
Ah, one quite important thing we need from something like this In |
|
It's not too bad to save this "shape" information. I think it can be done transparently when adding new variables to DynamicPPL.jl/src/simple_varinfo.jl Lines 420 to 430 in ad53ca9 |
Problem
Now that we properly support usage of different sizes in the underlying storage of the varinfo after linking, the current usage of
NamedTuplefor both the "ground truth" in TestUtils, e.g.DynamicPPL.jl/src/test_utils.jl
Lines 33 to 37 in 549d9b1
and in
SimpleVarInfo, makes less sense than it did before.To see why, let's consider the following example:
The issue here can really just be boiled down to the fact that we're trying to use the varname
to index a
NamedTuplewhich is after the transformation represented by a 1-length vector rather than 2-length vector.In contrast,
SimpeVarInfo{<:AbstractDict}will work just fine because here each varname gets its own entry:"Luckily" it has always been the plan that
SimpleVarInfoshould be able to use different underlying representations fairly easily, e.g. I've successfully used it withComponentVectorfrom ComponentArrays.jl many times before. And so we should probably find a more flexible default representation forSimpleVarInfothat can be used in more cases.Solution
Option 1: Use
OrderedDictby defaultThis one is obviously not great becuase of performance reasons, but it will "just work" in all cases and it's very simple to reason about.
Option 2: Dict-like flattened representation
In an ideal world, the underlying representation of the values in a varinfo would have the following properties:
VarName.Something like an
OrderedDictfails in two regards:Current
Metadataused byVarInfoThe
Metadatatype inVarInfois a good example of something that satisfies all three properties (of course, the "when possible" in Property (1) is not concrete, butVarInfouses aNamedTupleofMetadatato achieve this in most common use-cases).As a reminder, here is what
Metadatalooks like:DynamicPPL.jl/src/varinfo.jl
Lines 39 to 72 in e05bb09
Most importantly for a dict-like storage of values, are the following lines:
DynamicPPL.jl/src/varinfo.jl
Lines 46 to 58 in e05bb09
With this, it's fairly easy to implement nice indexing behavior for
VarInfo. Here's a simple sketch of what agetindexcould look like forMetadata:This is effectively the
getvalcurrently implemented:https://github.com/TuringLang/DynamicPPL.jl/blob/e05bb0935a1e1a06027c603cc04c20f23195a6c4/src/varinfo.jl#L318C44-L318C44
This then results in a
Vectorof the flattened representation ofvn.Our current implementation of
Base.getindexforVarInfothen contains more complexity to convert theVectorback into the original form expected by corresponding distribution, and it's usage looks likeSince the
distis also stored in theMetadata, the above in fact works the same if you dovarinfo[varname]ifvarinfo isa VarInfoand not aSimpleVarInfo. But, as have been discussed many times before, this is not great because it doesn't properly handle dynamics constraints, etc.; we want to use thedistat the place of index, not from the construction of thevarinfo.Nonetheless, value-storage part of
Metadataarguably proves quite a nice way to store values in a dict-like way while satisfying the three properties above.So why not just use
Metadata?Well, we probably should be. But if we're doing so, we should probably simplify its structure quite a bit.
For example, should we drop the following fields?
dists: As mentioned, this is often not the correct thing to use.gids: This is used by the Gibbs sampler, and will at some point not be of use anymore since we now have ways of programmatically conditioning and deconditioning models.orders: Only used by particle methods to keep track of the number of observe statements hit. This should probably either be moved somewhere else or at least not be hardcoded into the "main" dict-like object.flags: this might be generally useful, but the flags current used (istransanddelete) are no longer that useful (istransshould be replaced by explicit transformations, as is done inSimpleVarInfo, anddeleteshould also no longer be needed as now have a clear way of indicating whether we're running a model in "sampling mode" or not usingSamplingContext).But the problem of doing this, is that we'll break a lot of code currently dependent on
VarInfofunctioning as is.This is also the main reason why we introduced
SimpleVarInfo: to allow us to create simpler and different representations of varinfos without breaking existing code.So what should we do?
For now, it might be a good idea to just introduce a type very similar to
Metadatabut simpler in its form, i.e. mainly just a value container.We could either implement out own, or we could see if there are existing implementations in the ecosystem that could benefit us, e.g. Dictionaries.jl seems like it might be suitable.
The text was updated successfully, but these errors were encountered: