How to sample from MvNormal without allocating?

Even using the in-place method rand! to sample from a MvNormal distribution seems to allocate some memory and be orders of magnitude slower than the univariate case. I only have a two dimensional case, so I was hoping to find a way to optimize but can’t seem to find out how.

Suggestions?

There is no obvious reason to me that this should be allocating.
Here is the code for reference if anyone wants to go digging

cholesky on a PDMat should be nonallocating as it does that upfront during construction

A simple example:

using BenchmarkTools, Distributions, Random
using Random: rand!

X = MvNormal([2 1; 1 3])
y = zeros(2)
rng = Xoshiro()
@btime rand!($rng, $X, $y);

which produces

  133.021 ns (2 allocations: 96 bytes)

It seems to matter that the distribution is a ZeroMeanFullNormal but I’m not sure why.

There seem to be two sources of allocations. The first is PDMats.chol_lower called here. Not certain why this allocates; maybe because the result is a typeunion? This allocation is 16 bytes.

The second is this line, where the mean μ is added to the result via broadcast. The mean in this case is a FillArrays.Zeros, and it seems that broadcasted makes a copy:https://github.com/JuliaArrays/FillArrays.jl/blob/c3b38add861d475aadc66a112e045d7e0db31372/src/fillbroadcast.jl#L205. This results in an 80 byte allocation. Seems like this could be improved in FillArrays.

1 Like

Seems when I posted this, @jishnub had already begun working on a PR to fix this: don't materialize when broadcasting Zeros with Vector by jishnub · Pull Request #211 · JuliaArrays/FillArrays.jl · GitHub

1 Like

Seems like this has improved. Running the same tiny example

using BenchmarkTools, Distributions
using Random: Xoshiro, rand!

X = MvNormal([2 1; 1 3])
y = zeros(2)
rng = Xoshiro()
@btime rand!($rng, $X, $y);

69.586 ns (1 allocation: 16 bytes)