Reduce allocations in matrixdirichlet rand!

ReactiveBayes · Sep 27, 2024 · 335551b · 335551b
1 parent 5740b52
commit 335551b
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/src/distributions/matrix_dirichlet.jl b/src/distributions/matrix_dirichlet.jl
@@ -74,12 +74,9 @@ function BayesBase.rand(rng::AbstractRNG, dist::MatrixDirichlet{T}, nsamples::In
 end
 
 function BayesBase.rand!(rng::AbstractRNG, dist::MatrixDirichlet, container::AbstractMatrix{T}) where {T <: Real}
-    samples = vmap(d -> rand(rng, Dirichlet(convert(Vector, d))), eachcol(dist.a))
-    @views for row in 1:size(container, 2)
-        b = container[:, row]
-        b[:] .= samples[row]
+    @views for (i, col) in enumerate(eachcol(dist.a))
+        rand!(rng, Dirichlet(col), container[:, i])
     end
-
     return container
 end