Implement FillDist and ArrayDist

[mohamed82008]

Note: this comment was edited to reflect the final state of the PR so some of the discussions below may not make sense.

This PR introduces the filldist function which takes any univariate or multivariate distribution and returns another distribution that repeats the input distribution. Examples:

• filldist(Normal(), 10) and filldist(Bernoulli(), 10) return continuous and discrete multivariate distributions respectively.
• filldist(Normal(), 10, 10) and filldist(Bernoulli(), 10, 10) return continuous and discrete matrix-variate distributions respectively.
• filldist(MvNormal(rand(4), 1.0), 10) and filldist(Multinomial(4, [0.5, 0.5]) , 10) return continuous and discrete matrix-variate distributions respectively.

In this PR I also implemented the arraydist function which wraps an array of distributions returning a new distribution sampling from the individual distributions:

• arraydist(Normal.(rand(4))) and arraydist(Bernoulli.(rand(4))) return continuous and discrete multivariate distributions respectively.
• arraydist(Normal.(rand(4, 4))) and arraydist(Bernoulli.(rand(4, 4))) return continuous and discrete matrix-variate distributions respectively.
• arraydist(MvNormal.([rand(4) for i in 1:2], 1.0)) and arraydist(Multinomial.(4, [[0.5, 0.5] for i in 1:2])) return continuous and discrete matrix-variate distributions respectively.

[xukai92]

What additional information does ArrayDist provide? Seems like it could be handled by .~?

[mohamed82008]

Yes the plan is to make .~ use Multi and ArrayDist in a future PR. The current way of supporting .~ is a bit annoying.

[mohamed82008]

Note that with ArrayDist we do x ~ dist because it is just a normal multivariate or matrix-variate distribution.

[mohamed82008]

CC: @cscherrer

[mohamed82008]

I still need to add some tests before merging and releasing.

[xukai92]

Also for the names, ArrayDist is not consistent with Multi. Any justification?

[mohamed82008]

No reason. I can change it. What do you recommend? MultiDist?

[xukai92]

Is it possible to unify both as Broadcasted (#11), with the distinction of parameter types (single dist + dims v.s. vec of dists)?

[xukai92]

Also we might consider use lowercase here as things like truncated is now lowercase in Distributions.jl

[xukai92]

Both .~ and Multi seems to have issue with constrained variables, e.g. from Dirichlet.

[devmotion]

It seems both ProductVectorContinuousMultivariate and ProductVectorDiscreteMultivariate are covered by the Product distribution in Distributions? I guess, a similar approach could be used to remove the need for separate Discrete/Continuous implementations?

Moreover, since @xukai92 mentioned truncated, maybe one could just add additional implementations of product_distribution and define an analogue such as matrix_distribution? IMO defining such functions instead of ArrayDist makes it more obvious that you actually end up with different types - such as MvNormal in the case of product_distribution(::AbstractVector{<:Normal}).

[mohamed82008]

I don't think Product is good enough. I think it's better to to subtype multivariate continuous or discrete, etc. for dispatch later in Turing. That said, the names here are bad and were chosen on the fly in a frenzy so I need to go back and rethink them.

[devmotion]

Product still allows to dispatch on Product{Discrete} and Product{Continuous} (and of course also on something like Product{Discrete,<:Bernoulli} or Product{Continuous,<:Uniform}) - is that not sufficient?

[mohamed82008]

I rebased on the Zygote PR #25. This PR still needs tests.

[devmotion]

I still fail to see what advantage ProductVectorDiscreteUnivariate and ProductVectorContinuousUnivariate would have over Product{Discrete} and Product{Continuous}, and think one should probably just create a ProductMultivariate and ProductMatrix distribution in the same way (or loosen the restriction on univariate distributions in Product upstream). Moreover, IMO one should implement product_distribution for more (and non-univariate) distributions - it could default to Product (already defined in Distributions) and ProductMultivariate (for multivariate distributions), but for certain distributions also provide a more efficient higher-variate distributions (such as Normal -> MvNormal defined in Distributions). Ideally, Turing would then just use this potentially more efficient product_distribution instead of Product or ProductMultivariate.

Product still allows to dispatch on Product{Discrete} and Product{Continuous} (and of course also on something like Product{Discrete,<:Bernoulli} or Product{Continuous,<:Uniform}) - is that not sufficient?

[mohamed82008]

Fair enough. I will refactor and rethink this PR after the Zygote one is merged.

[mohamed82008]

I think the only place where we may need to subtype ContinuousMatrixDistribution or ContinuousMultivariateDistribution is in reconstruct in Turing. But we can just define new methods there and avoid the subtyping.

[mohamed82008]

I also like the fact that Multi(Normal(), 10) isa ContinuousMultivariateDistribution is true. Makes it just like any other distribution so we don't have to special case it elsewhere. I see your point though.

[devmotion]

I also like the fact that Multi(Normal(), 10) isa ContinuousMultivariateDistribution is true. Makes it just like any other distribution so we don't have to special case it elsewhere.

The same is true for Product, it's just a DiscreteMultivariateDistribution or a ContinuousMultivariateDistribution, depending on if you use discrete or continuous univariate distributions to define it. By adding

product_distribution(dist::UnivariateDistribution, n::Int) = Product(fill(dist, n))
product_distribution(dist::Normal, n::Int) = MvNormal(Fill(mean(dist), n), std(dist))

one could use product_distribution in the same way as the proposed Multi distribution.

[mohamed82008]

Yes with Product we can still do dispatch on all variants. But we will have to handle it separately in places like:

1. https://github.com/TuringLang/Turing.jl/blob/master/src/utilities/robustinit.jl#L16
2. https://github.com/TuringLang/DynamicPPL.jl/blob/0c0fafb66b4b06893869fea6375f1b89438eb0ba/src/utils.jl#L25
3. https://github.com/TuringLang/DynamicPPL.jl/blob/0c0fafb66b4b06893869fea6375f1b89438eb0ba/src/utils.jl#L35

I think it's nice to think in terms of a few higher level abstract types instead of Product{Continuous, Vector{<:ContinuousUnivariateDistribution}} and Product{Continuous, Matrix{<:ContinuousUnivariateDistribution}}, etc.

[mohamed82008]

I like the MvNormal + FillArrays idea though.

[devmotion]

But we will have to handle it separately in places like:

I'm sorry, I can't follow - you're definitely more familiar with these methods than I am 😛 Why would it not "just work"?

I think it's nice to think in terms of a few higher level abstract types instead of Product{Continuous, Vector{<:ContinuousUnivariateDistribution}} and Product{Continuous, Matrix{<:ContinuousUnivariateDistribution}}, etc.

You could always define aliases such as

const ProductVectorDiscreteUnivariate{T<:DiscreteUnivariateDistribution,V<:AbstractVector{T}} = Product{Discrete,T,V}
const ProductVectorContinuousUnivariate{T<:ContinuousUnivariateDistribution,V<:AbstractVector{T}} = Product{Continuous,T,V}

if needed. Is that what you mean?

[mohamed82008]

Let's take vectorize as an example.

vectorize(d::UnivariateDistribution, r::Real) = [r]
vectorize(d::MultivariateDistribution, r::AbstractVector{<:Real}) = copy(r)
vectorize(d::MatrixDistribution, r::AbstractMatrix{<:Real}) = copy(vec(r))

Product{<:Any, <:Any, <:AbstractVector} and Product{<:Any, <:Any, <:AbstractMatrix} cannot be subtypes of MultivariateDistribution and MatrixDistribution respectively so these methods don't apply. Having written this, I see we can just define const MultivariateDistribution = Union{Distributions.MultivariateDistribution, Product{<:Any, <:Any, <:AbstractVector{<:UnivariateDistribution}}} and then dispatch on our own unexported version of MultivariateDistribution. But my point is that we need to specially treat Product. Now if someone wanted to extend Turing in some way and needed to dispatch on every possible multivariate distribution on the RHS of ~, they would have to dispatch on DistributionsAD.MultivariateDistribution not Distributions.MultivariateDistribution. Additionally, since we don't export our MultivariateDistribution, Multi(Normal(), 10) will not be an instance of the exported MultivariateDistribution. Beside being semantically inelegant in my opinion, this just feels confusing.

I see advantages and disadvantages to both approaches.

[devmotion]

I see your point, and actually I think the problem is the current type system in Distributions. Too often I get the feeling that it is more prohibitive than actually useful. IMO it would be much nicer to use traits more, so, e.g., usually it should be fine to query eltype (if it would actually properly return the type of samples as in base...) and size to know if it is a multivariate distribution and what elements to expect. That would allow a more generic Product definition (BTW IMO the name is slightly confusing since it is typically used for something else: https://en.wikipedia.org/wiki/Product_distribution). However, currently I think one has to deal with the fact that one has to implement ProductMatrixUnivariate and ProductMultivariate as well (with possibly better names), just for the sake of Distributions's type system. One could still use product_distribution, of course, and IMO there's no need to replicate the (vector-valued) Product for univariate distributions.

[mohamed82008]

This one is ready for a review.

[devmotion]

Nice work, I'll have a more thorough look at it this afternoon! I've already added some comments.

[mohamed82008]

If no further comments here, I will go ahead and merge tonight (Australia time).