Linearization/flattening of SimpleVarInfo

[torfjelde]

This PR introduces a couple of things, though these are related:

1. unflatten(original[, spl], x): converts from a certain input x, usually a Vector, into an instance similar to original.
   • Effectively the same as the current constructor VarInfo(varinfo_old, spl, x).
   • I looked into using ParameterHandling.jl for this but decided against it for a couple of reasons:
     • Seems overkill.
     • unflatten-equivalent is constructed as a closure, which means that we need to keep track of this returned method rather than just using a "template" AbstractVarInfo + construction of unflattening requires construction of the flatten representation + the way one specifies the types is a bit too opinionated (which causes some issues with certain AD-frameworks) + closures can have less desirable performance characteristics.
     • The current Turing.jl-codebase is easily adapted to this unflatten since it's really just a matter of replacing calls VarInfo(varinfo_old, spl, x) with unflatten(varinfo_old, spl, x). A ParameterHandling.jl approach will require more work.
2. link!! and invlink!!, BangBang-versions of link! and invlink!, respectively, with some differences:
   • These take additional arguments which should always be sufficient to determine the transformation. These are:
     • model
     • sampler
     • t::AbstractTransformation. This sets us up for allowing alternative transformations to be used. As of right now, this only has an affect when calling link!! and invlink!!, not when used inside of the tilde-pipeline.
   • Also adds the logabsdet-jacobian term to the logp, so that getlogp(vi) ≠ getlogp(link!!(vi)) holds. This allows us to compute, say, logjoint by first linking vi in a single pass, and then compute logjoint(settrans!(vi, NoTransformation()), θ_constrained). Such a pattern, in particular if the transformation has been specified by the user themselves, will usually have much better performance than the logpdf_with_trans(..., true) within the tilde-callstack.
3. make_default_varinfo(rng, model, sampler) which allows one to overload on a, say, per-model or model-sampler-combination basis to specify which implementation of AbstractVarInfo to use.
   • Not entirely happy with this approach :confused:

EDIT: This should be merged after #420

[torfjelde]

This should be ready for a review @devmotion @yebai

[torfjelde]

bors try

[torfjelde]

bors try

[bors[bot]]

try

Already running a review

[bors[bot]]

try

Build failed:

• Turing.jl

[torfjelde]

bors try

[bors[bot]]

try

Build failed:

• Turing.jl

[torfjelde]

So the tests are breaking because the Emcee sampler in Turing is explictily calling initialize_parameters!!. It's not a function we're exporting to technically this shouldn't be a breaking chance, but given that it is indeed breaking downstream packages, should I make this a breaking release?

[yebai]

Yes, happy to make this a breaking PR to avoid surprises. This PR contains significant new code/features anyway.

[torfjelde]

Done. But as I said, everything passes except this one thing, so updating Turing is trivial.

[yebai]

bors try

[bors[bot]]

try

Build failed:

• test (1, ubuntu-latest, x64, 2)

[torfjelde]

bors try

[bors[bot]]

try

Build failed:

• test (1, ubuntu-latest, x64, 1)

[torfjelde]

bors try

[yebai]

bors cancel

[yebai]

It seems bors is on strike again. Let me know when this is ready for another look; I'll review and merge it manually.

[yebai]

bors try

[bors[bot]]

try

Already running a review

[yebai]

bors cancel

[yebai]

@torfjelde Did you run the tests locally?

[torfjelde]

They succeeded before I did the refactoring, so let me give it a go once more first.

[torfjelde]

bors try

[bors[bot]]

try

Already running a review

[torfjelde]

This is now passing with and without threads locally on Julia 1.8. Do we know what's wrong with bors?

[yebai]

I'll take a careful look later today and try to merge it manually.

Re bors, there seems to be a pattern that bors will strike if we feed it too many changes in one PR.

[torfjelde]

Cool!

Re bors, there seems to be a pattern that bors will strike if we feed it too many changes in one PR.

Regarding bors: but is there a way to reset it?

[yebai]

Re bors, there seems to be a pattern that bors will strike if we feed it too many changes in one PR. Regarding bors: but is there a way to reset it?

I didn't find one, but I didn't try very hard either.

[torfjelde]

Aight!

Btw, let me click the "Merge" button. You just let me know when you've reviewed and you're happy. This is just in case I discover something in the mean time.

[devmotion]

Regarding bors: but is there a way to reset it?

A bit late, but one can check the dashboard here if it is unclear if/what the problem is: https://app.bors.tech/repositories/24589 (one can navigate to it from the bors website or, IIRC, from the bors checks on Github)

One can cancel borg with bors r- or bors try- (depending on whether it was triggered with bors r+ or bors try). One can also run bors retry (and other things: https://bors.tech/documentation/). It seems the error message shows up if one runs bors try but there is already a bors try command running.