Vector bijectors

[penelopeysm]

This copies the interface over from https://github.com/penelopeysm/VectorBijectors.jl, as discussed in previous meetings, and extends it to support a wider range of distributions than Bijectors itself does.

Closes #237
Closes #258
Closes #398

Progress

• Univariates
  • LocationScale
  • Truncated
  • Censored
• Multivariates
• ReshapedDistribution
• Matrix
  • MatrixNormal, MatrixTDist
  • InverseWishart,Wishart generate positive (semi)definite matrices, and there is a PDBijector.
  • LKJ uses VecCorrBijector (although this one is not perfect, see recent issues)
  • Note that there are quite a few other matrix distributions (https://juliastats.org/Distributions.jl/stable/matrix/#Distributions), but we could conceivably ignore them for now — Bijectors' original implementation does not cover them, so skipping them for now would not be a regression.
• LKJCholesky
• Product distributions
  • tuples
  • arrays
  • namedtuples
• Mixture models. Note that we can probably only handle univariate mixtures sensibly and I think those should already subtype UnivariateDistribution so should end up using the truncated bijector --- if my instincts are correct, then that should already work out of the box and so the only thing we need to do is to add a test case.
• Order statistics and joint order statistics (should be trivial to implement by forwarding to the base distribution) It was quite tricky...

[github-actions]

Bijectors.jl documentation for PR #429 is available at:
https://TuringLang.github.io/Bijectors.jl/previews/PR429/

[penelopeysm]

@sunxd3 I'm still working on adding the actual distributions themselves, but could I get you to look at the tests and tell me if they seem sensible / if there's anything else I can add to them? So basically just the src/vector/test_utils.jl file, and also test/ (the interesting bit is test_utils, the tests are very boring, they just iterate over a ton of different distributions).

There's docs for the API at https://turinglang.org/Bijectors.jl/previews/PR429/vector/, if any of that is unclear do point out too. I suspect I might need to explain linked_optic_vec better.

Personally the test suite makes me feel quite confident that it's impossible to mess up the implementation of the bijectors, so I think the test suite is probably the most important part to get a review on.

[sunxd3]

Thanks Penny. I did a pass, the test interface is simply great, nothing missing or redundant as far as I could tell at the moment.

I think the doc for linked_optic_vec is good for now. I would like to see it used in DynamicPPL though and maybe adjusted in conjunction (is it intended that way?).

I'll give a deeper review next week.

[penelopeysm]

@sunxd3 I think finally I have some space to finish this off, I'll try to get the product distributions done this week.

I'm still thinking about linked_optic_vec. It's a bit hard to use it in DPPL because this is not released yet, but I can explain my thought process. The main use case I have for it is, say you have a model like this

@model function f()
    a ~ Normal()
    x ~ MvNormal(zeros(3), I)
end

# linked or unlinked doesn't matter
ldf = LogDensityFunction(f())

and you want to perform an optimisation where x[1] is constrained to (0, 1). Right now, we don't have any way of figuring out that in the length-4 vectorised parameters, x[1] is the second entry.

This function will let you do something like

ldf_optics = Union{VarName,Nothing}[]
for vn, range_and_linked in pairs(ldf._varname_ranges)
    linked = range_and_linked.is_linked
    this_vn_optics = if linked
        optic_vec(dist) # ok we don't have dist but we can get it from running the model
    else
        linked_optic_vec(dist)
    end
    append!(ldf_optics, this_vn_optics)
end

What one should in theory get out of that is VarName[a, x[1], x[2], x[3]]. Then that allows us to see that: 'ah, x[1] is the second entry', and we can write lb = [-Inf, 0.0, -Inf, Inf] and ub = [Inf, 1.0, Inf, Inf].

If it's linked, sometimes you can't identify a particular index as being x[1], that's when the nothings come in.

I realise this is actually quite hard to code correctly, and I don't know whether this plan will actually come to fruition in any way. So in a sense, implementing this might be a bit premature. But I guess there's not too much harm in having a wider interface, and if we don't use it that's a shame, but there won't be a performance cost or anything like that.

[sunxd3]

thanks Penny, sounds reasonable