AbstractVarInfo: the sampler interface

[phipsgabler]

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.

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

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

[phipsgabler]

It follows a summary of arguments and ideas brought forward in several issues about VarInfo; so the statements don't necessarily represent my own views.

The essence:tm:, by Mohamed:

We still need to keep track of distributions, sampler gids and varnames. We still need to get a vectorized form and set a vectorized form of the values for HMC samplers. We still need to cater for variables disappearing and popping at any time. We still need specialize the type of VarInfo to cater for mixed variable types and automatic differentiation in a type stable way.

Linearization/vectorization and indexing

• While it's clearly necessary to be able to linearise the state of the model at some point for the current implementation of HMC, it's not clear to me that it need to be the default mode of operation, and appears to add considerable complexity to the internals.

• Idea: make a new type Variable with information that needs to be associated with a particular random variable on a particular run, and put it into a NamedTuple or similar. Then provide efficient linearization and delinearization function (copyto!?). Perhaps the overhead gets less relevant when AD libraries improve. Optimizations would be possible if (parts of) models were known statically, to allow preallocation.

• Idea: just use a sufficently smart container type.

  I don't think it's necessary to know anything about the order or shape or number of random variables ahead of time. We just shouldn't try to create a container ahead of time but add samples whenever they appear. We just run the model function, and whenever a new sample is needed, we add it to the "NamedTuple on steroids". At this time point we know the type and the value of the sample.

  This would be in line with the "variable trie" representation idea floating around.

• The current idea described by Mohamed: > The way I think about it is that getval and setval! get and set the vectorized version of the value. So getval returns a vector and setval! expects a vector in even if the value was originally a scalar or matrix. getindex and setindex! on the other hand should reconstruct the shape of the original variable, and link/invlink the value based on whether the VarInfo is working in terms of the transformed or original values.

• Some existing complexity around constrained-space parameters and unconstrained-space parameters can be distinguished by introducing additional types UnonstrainedTypedVarInfo, ConstrainedTypedVarInfo. The sampler should be responsible for handling conversion between UnonstrainedTypedVarInfo and ConstrainedTypedVarInfo at the beginning and end of the sampling process.

• For example, calling getindex(:: UnconstrainedTypedVarInfo) should always parameters in the unconstrained space.

• We can also define a specialize_types function that specializes the types of the keys and values, while creating a fallback generic dictionary for taking in previously unseen types.

Interface cleanup/simplification

• Many of our functions should probably never have a do_something(vi, spl) signature -- we should find ways to handle everything on the back end without anyone worrying about how to use VarInfo. A better way would be to have the VarInfo stored somewhere in a shared state or tied to a specific model.

• Idea: make VarInfo only play an internal role, like a state or cache, and let all external functions operate on the model or sampler.

• Alternative idea: make interface dead simple; update!, parameters, logjoint, priors, ...

• Try to remove as much of the metadata as possible. Many of the more obscure fields could be stored at or inferred by the samplers (gid, flags), although getting this right is admittedly difficult.

• Well, we have r = vi[vn], r = vi[spl], vi[vn] = r, vi[spl] = r and push!, each defined for UntypedVarInfo and TypedVarInfo. So that's about 10 different implementations which is an average 40 loc per implementation.

  I.e., split of the indexing code. Also in line with separating representation and sampler interface.

• David proposed the following once:

  • getindex(varinfo, variables...) and setindex!(varinfo, value, variables...) for getting and setting values of variables
  • copyto!(varinfo, vector), copyto!(vector, varinfo), and probably copyto!(varinfo, variables, vector) and copyto!(vector, varinfo, variables) for copying to/from vectors of variables
  • getlogp, setlogp!, resetlogp!, and acclogp! for handling log probabilities
    Which seems like a good starting point.

• Another redundancy that we have comes from the need for r = vi[vn] and vi[vn] = r but also r = vi[spl] and vi[spl] = r. The first 2 are used in a model call, while the other 2 are used in the step! function.

• Get rid of unnecessary information and state: > To me it seems some of the design of AbstractVarInfo originates from having specific samplers in mind. I think ideally AbstractVarInfo should be sampler agnostic. So quite drastically, if gids are needed by the sampler in the next step, it should be part of the state of the sampler.

• Idea: reduce complexity by introducing a type hierarchy (typed/untyped, threaded, constraint/unconstraint, static, ...). Do this in a compositional manner to allow mixing -- variables in the same model might be treated differntly (cf. MixedVarInfo. Perhaps use traits?

Other things

• It would be nice to allow arbitrary types in UntypedVarInfo instead of just Reals. This should allow Turing to be used with distributions that sample letters, words, trees, images, etc. without having to represent them as an array of reals. The conversion to TypedVarInfo should work fine whatever type is in UntypedVarInfo so this may be as simple as changing Real to Any in a bunch of places.

• Getting a vector of support for each variable -- may help for box minimization, and potentially other sampler types.

• A super quick convenience function: getting a vector of variable names with indexing, i.e. names(vi) would return something like ["a[1,1]", "a[1,2]", "sigma"]

[devmotion]

We still need to keep track of [...] sampler gids

With the upcoming Gibbs refactoring in Turing, sampler gids are not used anymore (they are still set though which could be changed in a separate PR; more generally, it seems possible to remove DynamicPPL.Selector completely). @yebai mentioned that the idea was to make the extraction (and assignment?) of variables for a Gibbs component more efficient. It would be good to evaluate if there are any possible performance gains since the variable space of each sampler is known statically.

[torfjelde]

First off, you're an awesome human being for bringing all this together into a readable format ❤️

I'll respond to different parts of the issue in different comments so that we can have separate threads discussing different parts of the design. First: linking.

Linking

I think I'm of the opinion that this should always be done explicitly.

Bijectors.jl can handle both linearized (using Stacked) and NamedTuple-like inputs (using NamedBijector). And there's a lot of efficiency to be gained if you act on the full VarInfo rather than calling the corresponding bijector on a per-variable basis.

And "partial" transformations can be covered by just providing the Identity bijector for the subset of variables which we don't want to transform.

Though making things explicit probably does help a lot with knowing whether it's unconstrained or not, we should still probably have a way of querying the support of the varinfo. We could do this by defining constraints and have bijectors act on the constraints, e.g. Log()([0, ∞)) == (-∞, ∞), and so on. This would also make it much easier to inspect what's going on, improving both debugging and the user-experience.

Maybe https://github.com/invenia/ParameterHandling.jl could be useful?

TL;DR:

• Transformation of AbstractTrace should be explicit, e.g. b(trace), rather than implicit, e.g. trace[idx] # is it transformed or no?.
• Transformation could even be applied to Model, e.g. b(model)::TransformedModel, similar to Bijectors.TransformedDistribution.
• Partial transformations supported by transformations simply not transforming unspecified indices/fields or using Identity transformations.
• Allow querying of support of model/trace, and let transformations act on the supports to create the transformed support, e.g. Log()([0, ∞)) == (-∞, ∞).

[phipsgabler]

Couldn't linearization be treated as a bijection, too? So that you go from trace to transformed_trace to linearized_transformed_trace? Maybe with a variant

linearized_transformed_trace = transform!(buffer, Linearize(), transformed_trace)

?

(But then, I'm a noob in AD/HMC matters.)

[torfjelde]

Yes, it could; just a bit uncertain if it's useful since we don't need an jacobian term adjustment for such a transformation (absolute value of determinant is 1 for reshape and the like).

[phipsgabler]

I guess it's not so useful because of the Jacobian, but because it simplifies implementation, Linearization has always been heavily discussed, so I thought it might be a good idea to treat it in the unified framework of transformations. Perhaps it could be a special kind of bijector whose Jacobian part compiles to a no-op?

[torfjelde]

Yeah for sure, it could be just as Reshape (it's not in Bijectors.jl right now because it doesn't fit nicely with the current approach of encoding dimension in the Bijector; I'll soon rework Bijectors.jl to not do this though, thus allowing more arbitrary usages of input/outputs), though I'm uncertain if this should go into Bijectors.jl or be in a different package (depends on its implementation).