Implementing a Multivariate Normal Distribution in LibTorch

Alio: <p>I’m trying to write a Multivariate Normal Distribution with rsample() in LibTorch because LibTorch has no distributions yet.<br>
I’m not sure I have what it takes from the source code. I think I am having a problem with my Multivariate Normal Distribution knowing what to do with a covariance matrix instead of a single number standard deviation.<br>
I’m also having issues recreating <code>torch.distributions.MultivariateNormal</code> in Torch.</p>
<p><strong>The Python problems</strong><br>
This is my attempt to write a Multivariate Normal Distribution in Torch, in Python. I’m not getting errors but I am getting a matrix with the size torch.Size([2, 2]) from this. I get a torch.Size([2]) as output when input the exact same tensors into  <code>torch.distributions.MultivariateNormal</code> and take a sample or an rsample.</p>
<pre><code class="lang-auto">µ = torch.tensor([-0.5, 2.0])
cov = torch.tensor([[2.0, 8], [2.0, 40]])
pos_cov = cov @ cov.T 
Σ = torch.linalg.cholesky(pos_cov)

ε = torch.randn(1)

torch.det(torch.pi * 2 * Σ)**1/2 * torch.exp(torch.tensor(-1/2) * (ε - µ).T * Σ**-1 * (ε - µ))
</code></pre>
<p>Outputs a tensor like</p>
<pre><code class="lang-auto">tensor([[1182.5215,    0.0000],
        [1245.9080, 1102.0341]])
</code></pre>
<p><code>MultivariateNormal(µ, pos_cov).sample()</code><br>
Outputs a tensor like<br>
<code>tensor([ 7.6214, 37.2902])</code></p>
<p><strong>The LibTorch problems</strong><br>
I don’t see anything about a covarience matrix or Multivariate Normal Distribution in the CUDA <a href="https://github.com/pytorch/pytorch/blob/34523b70c1369024f805a0480bc39a5f21e64a03/aten/src/ATen/native/cuda/DistributionTemplates.h" rel="noopener nofollow ugc">DistributionTemplates.cu file</a></p>
<p>I know my Python implimentation of the Multivariate Normal Distribution doesn’t match the <code>torch.distributions.MultivariateNormal</code> function, but I tried implementing what I have in LibTorch.</p>
<pre><code class="lang-auto">#include &lt;torch/torch.h&gt;

class MultivariateNormalx{
    torch::Tensor mean, stddev, var;
public:
    MultivariateNormalx(const torch::Tensor &amp;mean, const torch::Tensor &amp;std) : mean(mean), stddev(std), var(std * std) {}

    torch::Tensor rsample() {
        auto device = torch::cuda::is_available() ? torch::kCUDA : torch::kCPU;
        auto eps = torch::randn(1).to(device);
        auto cholesky = torch::linalg::cholesky(torch::matmul(stddev.transpose(0, 1), stddev));
        auto pi = torch::tensor(3.141592653589793);
        return pow(torch::det(pi * 2 * cholesky), 1/2) * torch::exp(-1/2 * (eps - mean).transpose(0, 1) * pow(cholesky, -1) * (eps - mean));
    }
};

torch::Tensor μ = torch::tensor({{-0.5, 2.0}, {-0.5, 2.0}});
torch::Tensor cov = torch::tensor({{2.0, 8.0}, {2.0, 40.0}});

int main () {

    std::cout &lt;&lt; MultivariateNormalx(μ, cov).rsample();

}
</code></pre>
<p>Always returns…</p>
<pre><code class="lang-auto"> 1 -nan
 1  1
[ CPUFloatType{2,2} ]
</code></pre>

GMXeon: <p>I know this post is somewhat old, but I had to implement this distribution and came across this post. There’s a few fixes that need to be done to OP’s function.</p>
<p>Fix one: Matrix multiplication is not commutative, unlike regular multiplication. The order matters, and the matrices were accidentally supplied backwards.</p>
<p>Fix two: OP was a tad careless with their divisions or their brain was stuck in Python mode and forgot C will gladly do integer division without floating-point promotion where it can. As such, both divisions of constant integers resulted in zero because all remainder was discarded by integer division. Add a “.0” suffix to each constant to force floating-point division, or skip the compile-time division entirely and supply 0.5 and -0.5 respectively.</p>
<p>Suggestion one: Drop the ‘var’ property, it’s not used. Probably left in from debugging.</p>
<p>Suggestion two: Probably want to rename both pow(…) to at::pow(…) to disambiguate in case the std namespace is being used.</p>
<p>Suggestion three: Use M_PI (from cmath) for pi instead of defining your own inline constant, and the instantiation of the pi tensor should probably just be inlined since it’s only used once. It’s a shame that torch::pi doesn’t existing in libtorch, really odd that they omitted it.</p>