Rework of how transformations are handled

[torfjelde]

While working on #555, I ran into the pain that is our current way of handling transformations.

This is partially because we've been slowly adding support for things that wasn't support back in the day, while trying to preserve functionality.

I think it's now time rip out quite a bit of the historical artifacts and rework this.

In DynamicPPL.jl we effectively have a very annoying problem where we can have three different "representation" of a realization:

1. The "internal representation", i.e. how a AbstractVarInfo stores the realization internally.
2. The "linked internal representation", i.e. how a AbstractVarInfo stores a "linked" realization internally (this might differ from (1) in dimensionality, e.g. VarInfo, but also type, e.g. SimpleVarInfo).
3. The "model representation", i.e. how a realization of a variable is expected to be represented when working with it in the model body; this is always decided by the Distribution from which the variable is sampled.

Back in the day when everything was easier, we could define the map from_linked_internal (needed in getindex, assume, etc.) as a simple composition of from_internal (aka reconstruct) and invlink, and similarly for to_linked_internal (needed in setindex!, push!, etc.). But this resulting in multiple bugs once we started working with less convenient distributions, e.g. LKJCholesky where the linked representation is a Vector but the model representation is a Cholesky.

So these days we can't really think about from_linked_internal as a simple composition, but instead has to think about it as an "independent" mapping that can sometimes be represented in the old way not always.

To handle this, we made reconstruct, a function which was meant to take us from "internal representation" to "model representation", also accept the linking transformation, and attempts to construct whatever representation is needed for the pair (invlink_transform(dist), dist). This is okay, but because we have different implementations of AbstractVarInfo, e.g. VarInfo and SimpleVarInfo, each of which have different internal representations, the reconstruct has this behavior where it sometimes does what you expect and sometimes doesn't do anything, and it's difficult to determine exactly when what happens.

This PR removes reconstruct completely, in addition to other related methods, and effectively boils the entire transformation handling down to two mappings:

• from_internal_transform(varinfo, vn, dist): construct a transformation that takes us from the internal representation to the model representation.
• from_internal_transformation(varinfo, vn, dist): construct a transformation that takes us from the internal representation to the model representation.
• from_linked_internal_transform(varinfo, vn, dist): construct a transformation that takes us from the linked internal representation to the model representation. (correction by @yebai)

Notice that these methods construct a transformation, and we want these transformations to also define InverseFunctions.inverse and ChangesOfVariables.with_logabsdet_jacobian so that we can easily define the to_* mappings and obtain the log-abs-det-jacobian corrections for the transformations easily.

But that's really all we need + now everything is very explicit and clear.

There is a longer exposition on the topic in the docs accompanying this PR (see the docs preview under the "internals" tab) where I discuss why we need certain methods that might at a first glance now seem necessary.

Note that there is one drawback with this approach: we might run into type-instabilities since we're constructing a transform. One would hope that union splitting saves us; indeed that seems to work nicely for our test-suite on more recent Julia versions but not on 1.6. I'm not sure how much we should care about this though because once we move away from all the Selector stuff + old Gibbs sampler #573 , there's really no reason why anyone would every link only a subset of a varinfo, at which point we should always use the immutable link method and make whether or not a varinfo is linked available at compile-time, not runtime.

@devmotion @yebai

[torfjelde]

Aight, so it seems we weren't really testing cases such as

product_distribution(fill(Dirichlet(ones(4)), 2, 3))

(which might be because we didn't have support for this in Bijectors.jl until recently), and so I just added a test for this and now things are failiing 🙃

It turns out there's a bug in Bijectors.logpdf_with_trans for these distributions, so we need to fix this in Bijecotrs.jl. Will open an PR.

[torfjelde]

Thanks @mhauru

[torfjelde]

Looks like everything is passing nicely:)

[mhauru]

My only two unresolved comments are very minor, so I'm happy to approve. Likewise @yebai has a few unresolved comments, but he approved as well, so I assume he considers them non-blocking. @torfjelde, good to merge on your side?

[mhauru]

Having read the code more carefully, I think I agree with from_internal_transform being the right thing. Could still have a docstring, but optional.

[shravanngoswamii]

I will open another PR for converting diagrams images to Mermaid!

[mhauru]

The API docs index.html had grown to exceed the threshold size allowed by Documenter by a few kilobytes. I doubled the threshold to 400 KiB. There doesn't seem to be anything wrong with the API page, it's just plain long.

[torfjelde]

Does this count the images? Or is it just the HTML?

[mhauru]

I'm not sure if it would count the images, but they are small, and not on the page that exceeds the threshold.

[mhauru]

Hurray! Great stuff @torfjelde!