Make MatrixNormal sampling non-allocating (#2012)

* Make MatrixNormal sampling non-allocating

* Add test for non-allocating MatrixNormal sampling

* Update non-allocating test

Co-authored-by: David Müller-Widmann <devmotion@users.noreply.github.com>

* Fix

Co-authored-by: David Müller-Widmann <devmotion@users.noreply.github.com>

* Adjust allocation test to Julia <1.10.5 failure

Co-authored-by: David Müller-Widmann <devmotion@users.noreply.github.com>

---------

Co-authored-by: Desef, Benjamin <benjamin.desef@dlr.de>
Co-authored-by: David Müller-Widmann <devmotion@users.noreply.github.com>

Author: 3 people

src/matrix/matrixnormal.jl

@@ -122,11 +122,13 @@ end
 #  https://en.wikipedia.org/wiki/Matrix_normal_distribution#Drawing_values_from_the_distribution
 
 function _rand!(rng::AbstractRNG, d::MatrixNormal, Y::AbstractMatrix)
-    n, p = size(d)
-    X = randn(rng, n, p)
+    randn!(rng, Y)
     A = cholesky(d.U).L
     B = cholesky(d.V).U
-    Y .= d.M .+ A * X * B
+    lmul!(A, Y)
+    rmul!(Y, B)
+    Y .+= d.M
+    return Y
 end
 
 #  -----------------------------------------------------------------------------

test/matrixvariates.jl

@@ -322,6 +322,19 @@ function test_special(dist::Type{MatrixNormal})
             end
         end
     end
+    @testset "Non-allocating sampling" begin
+        # #2012: we can sample without allocations
+        M, U, V = _rand_params(MatrixNormal, Float64, 5, 5)
+        noallocD = MatrixNormal(M, cholesky!(Symmetric(U, :L)), cholesky!(Symmetric(V, :U)))
+        output = Matrix{Float64}(undef, size(noallocD))
+        allocs = ((d, out) -> @allocated(rand!(d, out)))(noallocD, output)
+        # See https://github.com/JuliaStats/Distributions.jl/pull/2012#issuecomment-3566807876
+        if VERSION < v"1.10.5"
+            @test allocs <= 32
+        else
+            @test iszero(allocs)
+        end
+    end
     nothing
 end