AbstractVarInfo: the representation

[phipsgabler]

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

1. Its 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.
2. 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

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

[mohdibntarek]

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.

[mohdibntarek]

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.

[mohdibntarek]

@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.

[mohdibntarek]

You can find my WIP experiment with trie implementations for DynamicPPL and hopefully Soss and Gen here https://github.com/mohamed82008/TrieLab.jl. I am trying to be as generic as possible allowing for any key => value mapping data structure to be used as the underlying data structure for the trie, or for multiple of them to be mixed and matched. Such data structures include Dict, StaticDict, NamedTuple, Tuple, Array, StaticArray etc. The underlying data structure also determines whether the trie is mutable or not and whether it is static or not. It is still a WIP and I have a list of missing features at the bottom of every file so any help is appreciated.

[mohdibntarek]

The idea behind allowing arbitrary mappings is because I want to allow keys at every level to be as arbitrary as possible: symbols, numbers, ranges, integer vectors, strings, structs, etc. Different data structures allow different key types but not others. So it can make sense to use a NamedTuple at the top but a dictionary or array somewhere lower to use numeric, string or struct keys, or vice versa. Also for some (sub-)models, static immutable mappings can make more sense while for others, dynamic mutable ones can be more convenient and being able to nest them seamlessly is going to facilitate nesting Turing, Gen and Soss models if such a trie is adopted everywhere. That's the vision anyways.

[phipsgabler]

❤️ Mohamed, you must have read my mind.

[devmotion]

Probably it would be good to use it to generalize the implementation in DataStructures, as soon as it's reasonably stable. Ref: JuliaCollections/DataStructures.jl#220, JuliaCollections/Tries.jl#2

[cscherrer]

@mohamed82008 it's interesting that you're allowing for more general keys. To this point I've focused on Symbol keys, considering a nested NamedTuple to represent nested namespaces. So for me, once you descend into something other than a Symbol, you're "inside a value", so it's no longer the responsibility of the namespace.

You can see my progress with this at NestedTuples.jl, which will be registered soon (maybe today? I pulled the trigger earlier this week). Everything is static, and I have some tricks for speeding things up that might also be useful for you as well.

Also, Setfield can be useful for this kind of thing, and its successor is in the works, so there's some opportunity to contribute, or at least be sure they're aware of any PPL use cases (most of which will probably come up in other contexts anyway):
https://github.com/JuliaObjects/Accessors.jl

[mohdibntarek]

Well the nice thing about allowing arbitrary mappings is that nested tuples from NestedTuples.jl are also allowed so there is no conflict between our works. I want to use Setfield/Accessors but I need an API that works fine for mutable and immutable structs alike. So something like a combination of BangBang.jl and Accessors.jl would probably be required.

[cscherrer]

I really like the idea of making abstracting things like this. Your static string code is a nice trick, too. If you need to lift more values to the type level, you might check out some of the helper functions in GeneralizedGenerated.

Accessors is still in the early stages, so I'd think it could be possible to have the mutability option there directly. Do you see a reason to have it separated?

I've also wondered about making a wrapper around named tuples to allow implementation of broadcasting, and maybe some other operations that would otherwise be type piracy. No rush on this at this point, but it seems interfacing with your API ought to be straightforward.

[phipsgabler]

FYI: 95 % of this or so was about the representation design (tries etc.), all of which is a bit fluid for now. So I moved the original issue into a discussion were we can keep on, and have split off the part about the sampler-related interface design (previously the "VarInfo proper" subsection) into a separate discussion.

[cscherrer]

@phipsgabler It's been a while since I've thought about this, but I was recently talking to Oliver Shulz about his ValueShapes package:
https://github.com/oschulz/ValueShapes.jl

Just passing it along here in case it can help with these issues

[torfjelde]

Actually, lets split the discussion into separate topics, e.g. the trie representation, the interface for a trace, etc.

Trie representation

[phipsgabler]

❤️ You are really writing out things in the fashion I have been thinking about. Very nice, thanks!

[phipsgabler]

Ad "distribution-equivalent but differently specified models": very good point, I have been thinking for some time about how to nicely refactor the current subsumption mechanism which I'm not quite happy with. This should IMHO be implemented on the level of VarName, though.

Maybe there is a way to formalize the subsumption relations in some mathematical structure (semilattice, preorder, ...?). Also related: first paragraphs from here.

[torfjelde]

Is there a reason why we can't just replace inds in VarName with a Lens from Setfield.jl? And then we use Setfield.jl to update?

If we then also implement subsumes for Lens, aren't we done? This would give us support for arrays and structs (that support Setfield.jl, but AFAIK there's no way to do this completely automatically).

[phipsgabler]

I had the same thought, but it seems suspiciously easy. E.g., wouldn't we want to constrain VarName a bit more? What about DynamicIndexLens and FunctionLens? Aren't there any extra requirements resulting from the trie implentation? Etc.

So I'm not at all thinking this is a bad idea, but on the other hand, lenses are really small to implement.

[torfjelde]

I had the same thought, but it seems suspiciously easy. E.g., wouldn't we want to constrain VarName a bit more? What about DynamicIndexLens and FunctionLens? Aren't there any extra requirements resulting from the trie implentation? Etc

The current subsume also doesn't support DynamicIndexLens and FunctionLens. And FunctionLens seems unnecessary IMO (could be nice syntax, but eh). I actually now have an implementation of VarName using lenses from Setfield.jl locally. The only thing that's left over to make it work is to make vinds return a lens instead of indices, and use @set! for the tilde-statements:)

And after thinking about it a bit, VarName should probably just go back to DPPL. It's too specific given what we're trying to achieve with APPL, and the above suggestions are more targetted towards improving DPPL rather than something we should have in APPL.

[torfjelde]

Interface

Copy-pasted from start of discussion:

• 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

[torfjelde]

IMO, this should be separate from the trie-representation. Nothing should stop us from defining a very simple StaticVarInfo which is only a wrapper around something like a ComponentArray with a logdensity accumulator.

[phipsgabler]

Since BangBang.jl-style updaters are already in use somwhere else in the infrastructure of Turing (AbstractMCMC, IIRC?), perhaps specifying the "mutating functions" in bangbang-style would be a good idea? This then wouldn't preclude immutable/persistent implementations.

[torfjelde]

I like that!

[cscherrer]

Just a general comment, the concreteness of some of the discussion here is pretty confusing. When we talk about "using Setfield" or "using Tries", or even "how VarNames are set up", are we really saying that the abstractions introduced in this library should support those as special cases, or that this abstract library will make some very concrete design decisions?

[phipsgabler]

Yeah, I'm also aways on the edge here between "we in DPPL" and "we specifying the general interface"... for VarNames though, they are already here in DPPL APPL. And IMHO can stay here.

[cscherrer]

they are already here in DPPL. And IMHO can stay here.

Now I'm really confused

[phipsgabler]

Whops, that was a typo. They have been moved from DPPL to APPL. To many acronymes here :P

Thinking about it once more: I assumed that 1. every library that supports the APPL interface needs to have some concept of "named variables" anyway, and 2. in order to define the interface, there needs to be some specification for them. Now, there is much less choices for varnames than for, say, VarInfo/traces, so Hong and I decided to move the VarName implementation here. They are, in a way, "downwards compatible" to plain symbols/nested named tuples, but add functionality like indexing which cannot be nicely represented by only symbols which exists nowhere else.

So far for my defense. Do you think it is a bad idea to have their implementation here? I guess one thing that could be done is, as soon as we have implemented that lense idea, to only provide the "wrapper type" and leave the choice of concrete lenses to implementors.

[cscherrer]

Thanks @phipsgabler . So it's something like, APPL provides an interface that's met by some external packages like Setfield and Accessors, and also has a simple reference implementation?

[phipsgabler]

I haven't really thought yet about that aspect, honestly. Defining VarNames only via an interface seems like an overkill to me, but perhaps that's untrue. I'd go for something like a "base implementation" in APPL (at least at the level of the current VarNames with names and multi-level indexing), which is extensible by other implementations (e.g. with hierarchical namespaces, properties, etc.).

[Moelf]

I'm late to the party, I wonder which page are people on now. I'm currently facing a problem where I need to generate a Turing "model" based on data and a user config file. I already figured out what to do if I can use @model directly, but I'm wondering how can I:

1. do it programmatically
2. manage variable names such that user can query / freeze model parameter by name instead of position of argument

[phipsgabler]

Maybe have a look at TuringGLM.jl, they do a similar thing. Of course it gets difficult if your model config has "dynamic structure", then I think you can only work via metaprogramming.

As for the arguments, they have a similar problem: TuringLang/TuringGLM.jl#36.

AbstractPPL is not really the right place to discuss this, though -- it's only about the interface above DynamicPPL.