[Bug Report] Sampling from MatrixNormal

[Wu-Chenyang]

MatrixNormal is sampled through the _rand function

function _rand!(rng::AbstractRNG, d::MatrixNormal, Y::AbstractMatrix)
    n, p = size(d)
    X = randn(rng, n, p)
    A = cholesky(d.U).L
    B = cholesky(d.V).U
    Y .= d.M .+ A * X * B
end

It directly uses the L and U matrices without permutation and therefore could produce incorrect results. I have to implement the following two functions to make it work.

function PDMats.unwhiten!(x::StridedVecOrMat, a::PSDMat)
    cf = a.chol.U
    view_x = x isa StridedVector ? view(x, a.chol.p) : view(x, :, a.chol.p)
    istriu(cf) ? rmul!(view_x, cf) : rmul!(view_x, transpose(cf))
    return x
end

function Distributions.rand!(rng::AbstractRNG, d::MatrixNormal, Y::AbstractMatrix)
    n, p = size(d)
    X = randn(rng, n, p)
    Y .= d.M .+ unwhiten!(unwhiten!(d.U, X), d.V)
end

The above unwhiten! function is a simple modification of the that from PDMatsExtras, but I actually don't quite understand it. According to the Wikipedia and JuliaDoc for CholeskyPivoted, shouldn't it be as follows?

function PDMats.unwhiten!(a::PSDMat, x::AbstractMatrix)
    chol = a.chol
    L, p, rank = chol.L, chol.p, chol.rank
    view(x, p, :) .= view(L, :, 1:rank) * view(x, 1:rank, :)
    return x
end

function PDMats.unwhiten!(x::AbstractMatrix{<:Real}, a::PSDMat)
    chol = a.chol
    U, p, rank = chol.U, chol.p, chol.rank
    view(x, :, p) .= view(x, :, 1:rank) * view(U, 1:rank, :)
    return x
end