`AbstractVarInfo` interface

[phipsgabler]

There are many aspects that make VarInfo a very complex data structure:

1. VarInfo really consists of multiple mappings: Metadata maps variables to their values,
   their distributions, their range in the underlying vector, and additional metadata such as flags,
   GIDs, and orders. Additionally, there are scalars logp and num_produce associated with the
   whole thing.
2. These mappings are not normal dictionaries. First, the order of variables needs to be preserved
   (would an ordered dictionary be enough for that?). Second, keys are themselves hierarchical,
   since they contain index values, and we need to be able to handle "sub-keys": get out, say,
   x[1:3] from the entry x.
3. The mapping between keys and values gets complicated even more, since some of the values need to
   be transformed between the distribution supports and Euclidean space. This is currently
   implemented by special logic in some getters/setters, which automatically call link/invlink.

Correct me if I'm missing something.

Now, it seems to me that a sensible separation of concerns would involve a separate interface of the
dictionary part, taking care of the mappings and the "sub-key" handling. Then the actual VarInfo
can just wrap this special dictionary (or several thereof), and add the additional metadata and
special logic to it. Then different implementation of the dictionary essentially replace the
different kinds of Metadata that separate UntypedVarInfo from TypedVarInfo.

I haven’t really worked with traits yet, but they might also have some advantages. For example, the
recently added ThreadSafeVarInfo is really only adding a thin layer to add an orthogonal concept
to the other variants.

Dictionary Trie Interface

Should describe an ordered mapping from VarName to whatever: actual values (<:Real or
<:AbstractArray{<:Real}?), distributions, flags, etc. – but only one such mapping at a time.

• Pure functions:
  • iterate, yielding pairs of VarName and the stored value
  • IteratorEltype == HasEltype(), IteratorSize = HasLength()
  • keys, values, pairs, length consistent with iterate
  • eltype, keytype, valuetype
  • get, getindex, haskey for indexing by VarName
  • merge to join two VarInfos
• Mutating functions:
  • push!, merge! to add and join elements
  • setindex!
  • empty!, delete!
• Conversions:
  • To array, for vectorization/linearization of values. Could be a method of Array, or copyto!.
  • To (nested?) named tuple
• Trie-specific functions:
  • Search for matching keys
  • Extract subkeys

David has argued that linearized storage destroys shape information. I think that is does not need
to: one could have one big array, with reshaped views into it – scalars would then just be
zero-dimensional views (akin to Refs). On the other hand, a tree-like implementation could have some other advantages, both with respect to ease of implementation and behaviour.

Subsumption, the issue of matching "sub-keys", can behave in different ways depending on the value
type. For sampled values, using sub-keys should always work: when a value of a vector-of-vectors
x is stored, then x[1][1] must be a valid key. However, this does not need to be the case for
distributions: when S comes from a matrix valued distribution, then S[1,1] is not necessarily
available.

Some pain points from the current implementation: vectorized access, shape preservation, association
of names and values.

Related: Gen's ChoiceMap

VarInfo proper

• getlogp, setlogp!, resetlogp!, and acclogp! for handling log probabilities

• gen_num_produce, set_num_produce, increment_num_produce, reset_num_produce; although these
  are specific to particle samplers, and would ideally be handled by them (maybe with a custom
  context in Turing?). A similar argument holds for the order information and certain flags.

• Linking: Since I haven't yet written any HMC samplers or the like, I unfortunately don't know how
  feasible/practical it is to make most of this more explicit. Mohamed, in a recent
  PR, has introduced methods for the
  getters/setters that take the distribution as additional argument to figure out the right thing:

  vi[vn, dist]

  Since that indexing in Julia supports keyword arguments, even somthing like the following would be possible:

  vi[vn, link=dist]

  The other getters and setters could then use the same kind of syntax.

  Then, there would also have to be (I guess?) additional methods for the conversion methods, to
  perform bijections for the array of vectorized values.

  Also, link!/invlink! for the complete mapping – is that necessary?

• Probably some resemblance of many of the old getters and setters: getidx, getdist, getrange,
  getall, getval, getvals, getinds, syms, getgid, flags. Hopefully some of them will
  become redundant.

• There should be one well-defined way to build up a VarInfo object. Perhaps an empty one + muliple
  calls to push!. What about push!!, with the possibility of immutable implementations?
  Alternatively: based on mergeing? Maybe merge can be general enough to cover both
  construction of all complex VarInfos, and producing MixedVarInfos when required.

VarName

This is working pretty well in its current form, but the logic for subsumption could be defined a
bit more explicitely. Some helper methods are certainly needed in order to properly support
subsumption in the dictionary interface, and would become useful elsewhere.

In the long term: what about more complex things, like getproperty-based variables?

Generated quantities

Different PPLs handle these in different ways – one alternative is to include them in the trace
(e.g., Soss.jl now uses a nested tuple). If this is considered useful, it needs to included in the
interface.

Related work/discussions

• Formalise VarInfo/AbstractVarInfo API DynamicPPL.jl#68
• Mohamed’s refactoring: https://github.com/TuringLang/DynamicPPL.jl/pull/188/files
• Collection of gripes and anti-gripes: Gripes and anti-gripes about VarInfo DynamicPPL.jl#119
• Splitting up code, refactoring: Splitting varinfo.jl DynamicPPL.jl#118
• Non-real values: Allow non-Real values in UntypedVarInfo DynamicPPL.jl#107
• Some older gripes: VarInfo refactor DynamicPPL.jl#5
• Older interface proposal: VarInfo Goals DynamicPPL.jl#7
• One more interface discussion: Standardising APIs between compiler and inference methods DynamicPPL.jl#16
• Generated quantities: Support transformed RVs & track expressions DynamicPPL.jl#94

CC @yebai, @cpfiffer, @devmotion, @mohamed82008

[mohamed82008]

The "dictionary" data structure is what I suggested to implement and test in a separate package and just use here. I also think we need hierarchical keys not just a dictionary-like data structure. We should allow keys to be more complex than just indexing so long as there is a way to getindex it from the data structure. This is especially useful when you have hierarchical models where the key of a variable in a lower level model is defined using its local key as well as the address/key of its model in the larger model. I will experiment with implementing this data structure in the Christmas break. Implementing VarInfos is how I like to spend my Christmas 😄 afterall TypedVarInfo was also drafted during Christmas 2 years ago!

Regarding vectorisation or not, I think we kind of agreed in the last meeting to have a vectorised global data structure in case random variables come in as scalars. But where they come as a whole, we should keep their shape. Making this type stable requires using TypeSortedCollections.jl.

[mohamed82008]

The Trie data structure from Gen is interesting https://github.com/probcomp/Gen.jl/blob/master/src/trie.jl

[femtomc]

@phipsgabler @devmotion linearization does not destroy information if you have a reference construction shape.

https://github.com/probcomp/Gen.jl/blob/579badeffa41747237a16e04386dff6414dc7e6e/src/choice_map.jl#L145-L225

This works even when primitive random variables have vector values.

(this is deconstruction to array of a map from addresses to values, and then back again given an existing map from which we can extract the shape)

[femtomc]

@mohamed82008 The trie structure is used because Gen treats model calls as first class - e.g. you internal nodes mark model calls, and leaf nodes mark primitive random choices.

This has the effective of creating a hierarchical address space - e.g. (if we consider tuple addresses) (:x, :y, :z => 1) means "inside the call at :x, then inside the call at :y, at address :z => 1".

You could also use a flat dictionary structure where the keys are tuples - but there are access and filtering penalties which are not optimal for implementing the core interfaces at model call site boundaries.

[devmotion]

linearization does not destroy information if you have a reference construction shape.

For me, the main pain points are not loss of information but that 1) linearization only works with samples that can be linearized (we would like to work with samples of arbitrary type) and 2) for many samplers linearization is not needed and creates additional overhead, also for developers.

Currently reconstruction is still annoying though, e.g., for samples of different element types or for scalar samples - they can be distinguished from multivariate samples of dimension 1 only by dispatching on the distribution it was generated from. Handling these special cases when implementing samplers is pretty annoying but could possibly be handled with better abstractions in VarInfo.

[mohamed82008]

@mohamed82008 The trie structure is used because Gen treats model calls as first class - e.g. you internal nodes mark model calls, and leaf nodes mark primitive random choices.

This is why I want to start from the Trie and adapt it for our needs and for performance. Also makes it easy to use it in Gen in the future.

[femtomc]

@devmotion I don’t quite understand the arbitrary samples but - do you mean arbitrary shaped samples ? If a primitive sample is any subtype of array, you should be able to flatten that and even reconstruct it correctly.

If it’s a sort of structure - I think you’re correct in pointing out that that’s not supported yet. I wonder if a clever usage of Flux’s restructure is the right thing.

[devmotion]

@devmotion I don’t quite understand the arbitrary samples but - do you mean arbitrary shaped samples ? If a primitive sample is any subtype of array, you should be able to flatten that and even reconstruct it correctly.

No, I mean samples that are not arrays but of arbitrary type, e.g., graphs, trees, text etc. Of course, I guess even for many of these other structures there exists some way to encode the information in an array - but probably it is quite unnatural and a simple reshaping operation is definitely not enough.

BTW is the ChoiceMap that you linked to above only used for array-like samples? I.e. not for scalars? Or do you vectorize scalars as well?

[femtomc]

A “choice” can come from any of the primitive distributions defined in the distributions section of Gen repo - so it can be anything (scalar, vector, even a custom structure if you define your own distribution using a DSL provided by the lib)

[femtomc]

@devmotion I’m hopping on a call with @phipsgabler to chat about relationship of VarInfo stuff and Gen at 10:30 AM EST - if you want to join chat send email!

[devmotion]

Thanks! Unfortunately, I won't be able to join (and I also won't make it to the Turing gathering tomorrow in case you plan to attend it), I'm busy with teaching duties.

[cscherrer]

FWIW this is the same reason we're adding nested named tuple support in Soss. Nested names tuple correspond to nested scope and allow models within other models. And we currently have SossGen, but I think through this we can make Gen support much more direct.

[phipsgabler]

I thought proposing a trie would be more controversial ("dictionary" was just me speaking losely), but nice that it seems to be consensus already.

[cscherrer]

Yeah, nesting is pretty natural. Something like Accessors can be very useful here too:
https://github.com/JuliaObjects/Accessors.jl

[devmotion]

It's a bit unfortunate that the trie implementation in DataStructures does not allow general keys.