ConditionalDistribution(T, m)

A conditional distribution that can map examples x to a distribution of type T via the mapping m by calling condition. There are no type restrictions on T so it could also be a function that constructs and appropriate distribution (e.g. like TuringMvNormal).

Examples

Moments, likelihoods, and samples can be obtained via e.g. mean, logpdf, and rand:

julia> using ConditionalDists, Flux
julia> m = ConditionalDists.SplitLayer(5,[2,2])
julia> p = ConditionalDistribution(MvNormal, m)
julia> rand(p, rand(5))
2-element Array{Float32,1}:
 0.2433714
 2.2054431

ConditionalMvNormal(m)

Specialization of ConditionalDistribution for MvNormals (for performance with batches of inputs). Does the same as ConditionalDistribution(MvNormal,m) but for batches of inputs a BatchMvNormal is constructed that does not just map over the batch but uses faster matrix multiplications.

julia> m = SplitLayer(100,[100,100])
julia> p = ConditionalMvNormal(m)
julia> @time rand(p, rand(100,10000);
julia> @time rand(p, rand(100,10000);
julia> @time rand(p, rand(100,10000);
 0.047122 seconds (23 allocations: 38.148 MiB, 24.25% gc time)

julia> p = ConditionalDistribution(MvNormal, m)
julia> @time rand(p, rand(100,10000);
julia> @time rand(p, rand(100,10000);
julia> @time rand(p, rand(100,10000);
 3.626042 seconds (159.97 k allocations: 18.681 GiB, 34.92% gc time)

The mapping m must return a Tuple with mean and variance. For a convenient way of doing this you can use a SplitLayer.

Examples

ConditionalMvNormal and SplitLayer together support 3 different variance configurations: fixed/unit variance, shared variance, and trained variance. The three different configurations are explained below.

Fixed/unit variance

Pass a function to the SplitLayer that returns the fixed variance with appropriate batch dimensions

julia> σ(x::Vector) = 2
julia> σ(x::Matrix) = ones(Float32,size(x,2)) .* 2
julia> m = SplitLayer(Dense(2,3), σ)
julia> p = ConditionalMvNormal(m)
julia> condition(p,rand(Float32,2)) isa DistributionsAD.TuringScalMvNormal

Passing a mapping with a single output array assumes unit variance.

Shared variance

For a learned variance that is the same across the the whole batch, simply pass a vector (or scalar) to the SplitLayer. The SplitLayer wraps vectors/scalars into a TrainableVectors/TrainableScalars.

julia> m = SplitLayer(Dense(2,3), ones(Float32,3))
julia> p = ConditionalMvNormal(m)
julia> condition(p,rand(Float32,2)) isa DistributionsAD.TuringDiagMvNormal

Trained variance

Simply pass another trainable mapping for the variance. By just supplying input sizes to SplitLayer you can automatically create Dense layers with given activation functions. In this example the second activation function makes sure that the variance is always positive

julia> m = SplitLayer(2,[3,1],[identity,abs])
julia> p = ConditionalMvNormal(m)
julia> condition(p,rand(Float32,2)) isa DistributionsAD.TuringScalMvNormal

SplitLayer(xs...)

A layer that calls a number of sublayers/mappings with the same input and returns a tuple of their outputs. Can be used in a regular Flux model:

julia> m = Chain(Dense(2,3), SplitLayer(Dense(3,2), x->x .* 2))
julia> length(params(m)) == 4
julia> (x,y) = m(rand(2))
(Float32[-1.0541434, 1.1694773], Float32[-3.1472511, -0.86115724, -0.39665926])

Comes with a convenient constructor for a SplitLayer built from Dense layers with given activation(s):

julia> m = Chain(Dense(2,3), SplitLayer(3, [2,5], σ))
julia> (x,y) = m(rand(2))
(Float32[0.3069554, 0.3362006], Float32[0.437131, 0.4982477, 0.6465078, 0.4523438, 0.5068563])

You can also provide just a vector / scalar that should be trained but have the same value for all inputs (like a lonely bias vector). This functionality is provided by the TrainableVector/TrainableScalar types. For vector inputs they simply return the array they are wrapping. For matrix (i.e. batch) inputs they return appropriately repeated arrays:

julia> m = SplitLayer(Dense(2,3), ones(Float32,3))
julia> length(params(m)) == 3
julia> (x,y) = m(rand(2,5))
julia> size(y) == (3,5)
julia> y
3×3 Array{Float32,2}:
 1.0  1.0  1.0
 1.0  1.0  1.0
 1.0  1.0  1.0

condition(p::ConditionalDistribution, x::AbstractMatrix)

Map a batch of examples x (examples as columns of x) to the distribution type specified in p.

condition(p::ConditionalDistribution, x::AbstractVector)

Map an example x to the distribution type specified in p.