Attempt at implementation of VarNameVector (Metadata alternative)

[torfjelde]

This is an attempt at a implementation similar to what was discussed in #528 (comment).

The idea is to introduce a simpler alternative to Metadata, which contains a fair bit of not-so-useful-anymore functionality and isn't particularly accessible.

With the current state of this PR, it's possible to evaluate a model with a VarNameVector:

julia> using DynamicPPL

julia> model = last(DynamicPPL.TestUtils.DEMO_MODELS);

julia> x = rand(OrderedDict, model)
OrderedDict{Any, Any} with 2 entries:
  s => [0.91846 1.26951]
  m => [0.379781, -2.30755]

julia> vnv = VarNameVector(x)
[s = [0.9184600730345488 1.2695082629237986], m = [0.3797814300783162, -2.307552998945772]

julia> logjoint(model, VarInfo(vnv))
-15.691863927983814

But it's not possible to sample.

What are we trying to achieve here?

The idea behind VarNameVector is to enable both an efficient and convenient representation of a varinfo, i.e. an efficient representation of a realization from a @model with additional information necessary for both a sampler developer and end-user to extract what they need.

In that vein, it seems like we need something that sort of acts like a Vector, and sort of acts like an OrderedDict, but is neither (at least, I don't think it fits in either).

1. We need a Vector because we need:
   1. A straight-forward way to work with gradient-based samplers (and more generally, most sampler implementations assumes a simple Vector representation).
   2. Efficient representation of the realizations, e.g. having variables be in contiguous blocks of memory is good (e.g. in the case of for loop over variables x[i] in a model, we want the values for x to be stored in a contiguous chunk of memory).
2. We need an OrderedDict because we need:
   1. A simple way for the end-user to get information related to a particular VarName, e.g. trace[vn] should result in the realization in it's "constrained" space (i.e. in the space that the distribution of which it came from commonly works in).

Moreover, the above should be achieved while ensuring that we allow both mutable and immutable implementations, in addition to maintaining type-stability whenever possible and falling back to type unstable when not.

Current implementation: VarInfo with Metadata

The current implementation of VarInfo with Metadata as it's underlying storage achieves some of these properties through a few different means:

1. Type stability for varname-specific operations, e.g. getindex(varinfo, varname), is achieved by effectively grouping the varnames contained in a VarInfo by their symbol (which is part of the type), and putting each group in a separate Metadata in a NamedTuple. Basically, instead of representing pairs like (@varname(x) => 1.0, @varname(y) => 2.0) in a flattened way, we represent them as (x = (@varname(x) => 1.0,), y = (@varname(y) => 2.0,)), such that in the scenario where they different in types, e.g. y isa Int, we can dispatch on the sym in VarName{sym} to determine which entry of the NamedTuple to extract.
   1. This, IMO, is a good approach, and one we should continue as this grouping by symbol is generally what the end-user does anyways, e.g. x is continuous while y is discrete. It also provides very straight-forward guidelines on how to speed up models, i.e. "group variables into a single higher-variate random variable, e.g. x[1] ~ Normal(); x[2] ~ Normal() into x ~ fill(Normal(), 2)".
2. Run once with type-unstable varinfo to get type-information, and then use this for the subsequent runs to obtain type-stability (and thus much improved performance).
   1. IMO, we should also keep this idea, as it works very well in practice.
   2. But the current impl of type-stable VarInfo is somewhat lacking, e.g. it does not support changing support even if the types stay the same,
3. Interaction for the end-user is mainly done through varinfo[vn], but is currently a) limited, and b) very confusing.
   1. Currently, VarInfo implements a somewhat confusing subset of both AbstractDict and AbstractVector interfaces, but neither of which are "complete" in any sense.
   2. There is little coherency between the operations one can perform on VarInfo and Metadata, even though VarInfo is effectively just a wrapper around Metadata / NamedTuple containing Metadata, which adds further confusion.
   3. Metadata comes with a lot of additional fields and information that we have slowly been moving away from using, as the resulting codebase ends up being overly complicated and difficult to work with. This also significantly reduces the utility of VarInfo for the end-user.

Replacing Metadata with VarNameVector

VarNameVector is an attempt at replacing Metadata while preserving some of the good ideas from VarInfo with Metadata (and in the process, establish a bare-minimum of functionality required to implement a underlying storage for varinfo):

1. VarNameVector (should) implement (almost) all operations for AbstractDict + all non-linear-algebra operations for a AbstractVector{<:Real}, to the point where a user will find it easy to use VarNameVector (and thus a VarInfo wrapping a VarNameVector).
   1. AbstractDict interface is implemented as if keys are of type VarName and values are of type corresponding to varinfo[varname].
   2. AbstractVector interface is implemented wrt. underlying "raw" storage, e.g. if we're working with a unconstrained representation, then the vector will be in unconstrained space (unlike varinfo[varname] which will be in constrained space).
2. VarNameVector uses a contiguous chunk of memory to store the values in a flattened manner, both in the type stable and unstable scenarios, which overall should lead to better performance in both scenarios.
3. VarNameVector reduces the number of fields + replaces some now-less-useful fields such as dist with fields requiring less information such as transformss (holding the transformations used to map from "unconstrained" to "constrained" representation).

Updates

2023-11-13

Okay, so I've made some new additions:

1. Much improved testing of the core functionalities for VarNameVector.
2. It's now possible to push! and update! a VarNameVector where:
   • push! means what it usually means, but the varname / key must be unique.
   • update! works with either existing keys or new keys.

push! & update!

The idea behind the implementation is as follows:

1. When we see a new varname, push! (and update!, which is equivalent in this scenario), is straight-forward: we just call push! and setindex! on all the necessary underlying containers, e.g. ranges.
2. When wee see an already existing varname, things become a bit more complicated for a few fields, namely ranges and values (everthing else is just call to setindex!):
   • If the new value requires less memory than the previous value, we can make use of the part of values already allocated to varname.
   • If the new value requires more memory than the previous value, things become a bit more difficult (once again). One approach would be to simply allocate a larger chunk of values to varname and then shift all the values occuring after this part; doing this on every update! will be expensive! Instead, we just move the new value to the end of values and mark the old location as "inactive". This leads to much more efficient update!, and then we can just perform a "sweep" to re-contiguify the underlying storage every now and then.

To make things a bit more concrete, consider the following example:

julia> using DynamicPPL

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

julia> vnv.varname_to_index
OrderedDict{VarName{sym, Setfield.IdentityLens} where sym, Int64} with 2 entries:
  x => 1
  y => 2

julia> vnv.ranges
2-element Vector{UnitRange{Int64}}:
 1:1
 2:2

julia> OrderedDict(vnv)
OrderedDict{VarName{sym, Setfield.IdentityLens} where sym, Any} with 2 entries:
  x => 1.0
  y => [2.0]

julia> vnv[:]
2-element Vector{Float64}:
 1.0
 2.0

Then we update the entry for @varname(x) to a differently sized value:

julia> DynamicPPL.update!(vnv, @varname(x), [3.0, 4.0, 5.0]);


julia> vnv.varname_to_index
OrderedDict{VarName{sym, Setfield.IdentityLens} where sym, Int64} with 2 entries:
  x => 1
  y => 2

julia> vnv.ranges
2-element Vector{UnitRange{Int64}}:
 3:5
 2:2

julia> OrderedDict(vnv)
OrderedDict{VarName{sym, Setfield.IdentityLens} where sym, Vector{Float64}} with 2 entries:
  x => [3.0, 4.0, 5.0]
  y => [2.0]

julia> vnv[:]
4-element Vector{Float64}:
 3.0
 4.0
 5.0
 2.0

Notice that the order is still preserved, even though the underlying ranges is no longer ordered.

But, if we inspect the underlying values, this contains now-inactive entries:

julia> vnv.vals
5-element Vector{Float64}:
 1.0
 2.0
 3.0
 4.0
 5.0

But in the scenario where we care about performance, we can easily fix this:

julia> DynamicPPL.inactive_ranges_sweep!(vnv);

julia> vnv.varname_to_index
OrderedDict{VarName{sym, Setfield.IdentityLens} where sym, Int64} with 2 entries:
  x => 1
  y => 2

julia> vnv.ranges
2-element Vector{UnitRange{Int64}}:
 1:3
 4:4

julia> OrderedDict(vnv)
OrderedDict{VarName{sym, Setfield.IdentityLens} where sym, Vector{Float64}} with 2 entries:
  x => [3.0, 4.0, 5.0]
  y => [2.0]

julia> vnv[:]
4-element Vector{Float64}:
 3.0
 4.0
 5.0
 2.0

Type-stable sampling for dynamic suppport

The result of this is that we can even perform type-stable sampling for models with changing support:

julia> using Distributions

julia> @model function demo_random_num_variables(::Type{TV}=Vector{Float64}) where {TV}
           α ~ Dirichlet(ones(10))
           d ~ Categorical(α)
           x = TV(undef, d)
           for i = 1:d
               x[i] ~ Normal()
           end

           return (; α, d, x)
       end
demo_random_num_variables (generic function with 4 methods)

julia> model = demo_random_num_variables();

julia> x = rand(OrderedDict, model);

julia> vi = VarInfo(VarNameVector(x));

julia> first(DynamicPPL.evaluate!!(model, empty!!(vi), SamplingContext()))
(α = [0.38811824437465864, 0.04808086905976418, 0.2038431779706373, 0.07337088670210494, 0.039133882408812014, 0.012910572919725475, 0.014187317531812365, 0.05987651255436473, 0.12002476163571087, 0.040453774842409536], d = 1, x = [-0.26582545522147744])

julia> OrderedDict(last(DynamicPPL.evaluate!!(model, empty!!(vi), SamplingContext())).metadata)
OrderedDict{VarName, Any} with 8 entries:
  α    => [0.0206437, 0.048932, 0.106592, 0.016048, 0.00882465, 0.383303, 0.173889, 0.06695…  d    => 6.0
  x[1] => 0.755316
  x[2] => -1.22984
  x[3] => 1.0898
  x[4] => 0.678128
  x[5] => -1.17279
  x[6] => 0.654284

julia> OrderedDict(last(DynamicPPL.evaluate!!(model, empty!!(vi), SamplingContext())).metadata)
OrderedDict{VarName, Any} with 9 entries:
  α    => [0.115859, 0.120242, 0.213961, 0.069815, 0.0541034, 0.207712, 0.0977965, 0.062019…  d    => 7.0
  x[1] => -1.78224
  x[2] => 0.357474
  x[3] => -0.39225
  x[4] => -0.154527
  x[5] => -1.00411
  x[6] => 1.01872
  x[7] => 0.262885

The really nice thing here is that, unlike with TypedVarInfo, we don't need to mess around with boolean flags to indicate whether something should be resampled, etc. Instead we just call similar on the VarNameVector and push! onto this.

We can also make this work nicely with TypedVarInfo:)

2024-01-26T15:00:38

Okay, so now all tests should be passing and we should have, effectively, feature parity with VarInfo using Metadata.

But this required quite a lot of code, and there are a few annoyances that are worth pointing out / discussing:

Transformations are (still) a pain

My original idea was that we really did not want to attach the entire Distribution from which the random variable came from to the metadata for a couple of reasons:

1. This can technically change between evaluations of a model (or, maybe a more realistic scenario, we can't use a VarInfo from ModelA in ModelB, even though they only differ by the RHS of a single ~ statement).
2. We're really only using the Distribution to determine the transformation from the vectorized / flattened representation to the original one we want and linking.

In fact, the above is only partially true: getindex inside a model evaluation actually uses getindex(vi, vn, dist) and uses the "tilde-local" dist for the reshaping and linking, not the dist present in vi.

Hence it seemed to me that one immediate improvement of Metadata is to remove the dist completely, and instead just store the transformation from the vectorized / flattened representation to whatever desired form we want (be that linked or invlinked).

And this is what we do currently with VarNameVector.

This then "resolves" (1) since now all we need is a map f that takes a vector and outputs something we can work with.

However, (2) is still an issue: we still need the "tilde-local" dist to determine the necessary transformation for the particular realization we're interested in, while simultaenously wanting getindex(vi, vn) to also function as intended outside of a model, i.e. I should still be able to do

@model demo() = x ~ LKJCholesky(2, 1.0)
vi = DynamicPPL.link(VarInfo(model), model)
vi[@varname(x)] == vi[@varname(x), LKJCholesky(2, 1.0)] # => true

Sooo I don't think we can get around having to keep transformations in the metadata object that might not actually get used within a model evaluation if we want allow the indexing behavior above 😕

2024-01-31: Update on transformations

See #575

[coveralls]

Pull Request Test Coverage Report for Build 10924238467

Warning: This coverage report may be inaccurate.

This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.

• For more information on this, see Tracking coverage changes with pull request builds.
• To avoid this issue with future PRs, see these Recommended CI Configurations.
• For a quick fix, rebase this PR at GitHub. Your next report should be accurate.

Details

• 542 of 626 (86.58%) changed or added relevant lines in 11 files are covered.
• 40 unchanged lines in 4 files lost coverage.
• Overall coverage increased (+0.08%) to 77.586%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
src/abstract_varinfo.jl	0	1	0.0%
src/utils.jl	19	21	90.48%
src/varnamedvector.jl	337	375	89.87%
src/varinfo.jl	137	180	76.11%

Files with Coverage Reduction	New Missed Lines	%
ext/DynamicPPLEnzymeCoreExt.jl	1	50.0%
src/utils.jl	2	82.4%
src/abstract_varinfo.jl	6	78.62%
src/varinfo.jl	31	80.84%

Totals
Change from base Build 10701838580:	0.08%
Covered Lines:	3150
Relevant Lines:	4060

---

💛 - Coveralls

[codecov]

Codecov Report

Attention: 200 lines in your changes are missing coverage. Please review.

Comparison is base (c33eeae) 84.32% compared to head (ff68206) 79.25%.

Files	Patch %	Lines
src/varinfo.jl	48.83%	110 Missing ⚠️
src/varnamevector.jl	87.03%	35 Missing ⚠️
src/abstract_varinfo.jl	62.74%	19 Missing ⚠️
src/context_implementations.jl	9.52%	19 Missing ⚠️
src/utils.jl	81.08%	7 Missing ⚠️
src/extract_priors.jl	20.00%	4 Missing ⚠️
src/threadsafe.jl	66.66%	4 Missing ⚠️
src/simple_varinfo.jl	86.66%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #555      +/-   ##
==========================================
- Coverage   84.32%   79.25%   -5.07%     
==========================================
  Files          26       27       +1     
  Lines        3183     3621     +438     
==========================================
+ Hits         2684     2870     +186     
- Misses        499      751     +252     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[mhauru]

This isn't quite finished, but it's close enough that I'll open it for review, since I'm off work tomorrow. Things that remain to be done:

• Sort out the only test failure, which is some strange Zygote error that appeared when I changed something seemingly unrelated. Might be a Zygote bug.
• Decide on a migration path from Metadata + old Gibbs sampler to VarNamedVector + new Gibbs sampler. Currently this PR makes VarNamedVector the default metadata type for VarInfo to see that tests for it pass, but that needs to be reverted, assuming we won't bother making VarNamedVector work with the old Gibbs sampler. Doing that would require putting in structures for gids, which seems like a waste of effort.
• There's one substantial difference between Metadata and VarNamedVector from a user's point of view: The former doesn't resample values if you call model(vi) again unless you set the "del" flags, while the latter does. There are a couple of places where I gloss over this, in assume and in setval_and_resample!. I think that's all okay, but I need to understand this better to really convince myself.

[mhauru]

This method claimed to work for AbstactChains, but in fact only worked for MCMCChains. Hence I moved it to the MCMCChains extension as I refactored this function to work with VarNamedVectors.

[mhauru]

These tests were checking that calling a model with a VarInfo doesn't by default overwrite existing values with newly sampled ones. This is, intentionally, no longer the case with VarNamedVector.

[mhauru]

The splitting of FromVec to UnwrapSingletonTransform and ReshapeTransform isn't really core to this PR, it's something that come up when I was changing how VarNamedVector does transformations. It fixes a problem where from_vec_transform(x)(tovec(x)) was not a no-op for 0-dimensional Arrays. If you'd like to review it separately let me know and I can do some cherry-picking to separate it to a distinct PR.

[yebai]

It's okay to keep it in the PR, but perhaps add it to this PR's summary so we can easily see it in the future.

[torfjelde]

Agreed; not a huge thing, so might as well just keep it here 👍

[mhauru]

I would request a review from @torfjelde, but I can't because this is his PR :)

[yebai]

@sunxd3, can you also provide your review, given that you have worked with VarInfo extensively?

[sunxd3]

Absolutely, will take a look.

[yebai]

Dict actually supports this push! API, so we could move it to the list of L24-L30 above. For example,

julia> d = Dict()
Dict{Any, Any}()

julia> push!(d, :s=>nothing)
Dict{Any, Any} with 1 entry:
  :s => nothing

[yebai]

Maybe I am missing something. What is the difference between update! and setindex!?

[mhauru]

update! can take the transform argument, and can also deal with variables changing dimension. setindex! is more low level and just directly writes to the storage vector.

The more I think about it, the less the distinction seems necessary to me. @torfjelde, did you have a reason to keep these separate, or would you mind essentially renaming update! to setindex!, and making it be synonymous push! be synonymous with it too? This would also be in line with how push! works for Dicts, where it allows both adding and changing values.

[yebai]

It's okay to keep it in the PR, but perhaps add it to this PR's summary so we can easily see it in the future.

[mhauru]

A Zygote issue for the failing test here: FluxML/Zygote.jl#1523 Will have to think about workarounds.

[torfjelde]

Yeah, I'd say no; for now, let's be explicit about it and test it properly 👍

[torfjelde]

Hmm, is this worth it? As in, would it not be better to just pass an instantiated metadata? Not a huuuuge fan of passing types like this.

[mhauru]

Docs build now, and @willtebbutt helped me put in place a temporary workaround for the Zygote issue.

[mhauru]

@yebai asked me to make a slide or two on VarNamedVector. Here's a sketch (needs figures/prettying), may be helpful for reviewers:

VarNamedVector

• A data structure for holding the values of random variables.
• When indexed with VarNames, works like a dictionary/mapping. E.g. vnv[@varname(x)] = rand(MvNormal([1.0, 1.0]))
• When indexed with integers or Colon, works like a vector. E.g. all_values_in_a_vector = vnv[:].
• Values for all variables stored in one contiguous chunk of memory.
• Can handle variables changing dimension, usually without allocating new memory.
• Stores everything as a vector, together with a transformation that maps the vector back into the full value (e.g. reshapes matrix variate values).
  • Indexing with a VarName returns the full value in the "original" space.
  • Indixing with an integer returns the value as stored internally.
• Also stores a flag for whether a variable is known to have been transformed so that the domain of the internal storage is all of Euclidean space.
  • Useful for linking in sampling: If the flag is true, then one can modify the internal storage directly with setindex_raw!(vnv, y, @varname(x)) where y can be any real number of real vector of the correct dimension, and vnv[@varname(x)] should be a valid value for the variable @varname(x).
• Indexing into a VarNamedVector with complex VarNames I think still requires some work. E.g. if vnv[@varname(x.lu[1])] = 0.1 and vnv[@varname(x.lu.a)] = randn(2,2), then what should vnv[@varname(x)] be? Currently there are several different "get" and "set" methods and they handle these things differently. I would leave this for a future PR where we rework VarInfo/SimpleVarInfo.

[torfjelde]

What is this ignorable thingy?

[torfjelde]

Or rather, why is it needed?

[mhauru]

unset_flag! errors if you try to set a flag which can not be set for VNV, on a "rather error than silently do the wrong thing" basis. assume calls unset_flag!(vi, vn, "del"), which would thus error. It seems to be fine to ignore the fact that the call actually does nothing for VNV, so I added an argument to be able to do so, and gets tests to pass for now. I haven't understood fully why assume wants to unset said flag, but still seems to work fine even if it can't do so. I expect to remove ignorable once I've figured that out.

[mhauru]

We just had a long chat about this PR with @sunxd3, @willtebbutt, and @torfjelde. We agreed that given the lack of a clear design document, and the fact that much of this is still working around various aspects of the old Gibbs sampler, @willtebbutt and @sunxd3 will only review the new files: varnamedvector.jl, its tests, and its docs. They are relatively self-contained, although the motivation for some of their features will remain opaque without the surrounding context. That's fine. We'll do another round of reviewing the whole structure around VarInfo once the old Gibbs is out and we've cleaned up after it.

Some other points that came up:

• Should check some literature, see if contiguify should be called coalesce: https://en.wikipedia.org/wiki/Coalescing_(computer_science)
• Add some more @inference tests to check for type stability.