Understanding MultivariateNormal log_prob input semantics

Wesley_Neill: <p>Thanks to these two fantastic resources (I hope they help someone else), I understand how to instantiate batched Normal and MultivariateNormal objects:</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://bochang.me/blog/posts/pytorch-distributions/">
  <header class="source">
      <img src="https://bochang.me/dice-16x16.png" class="site-icon" width="16" height="16">

      <a href="https://bochang.me/blog/posts/pytorch-distributions/" target="_blank" rel="noopener nofollow ugc">Understanding Shapes in PyTorch Distributions Package</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/248;"><img src="https://bochang.me/blog/posts/pytorch-distributions/shape_sementics.png" class="thumbnail" width="690" height="248"></div>

<h3><a href="https://bochang.me/blog/posts/pytorch-distributions/" target="_blank" rel="noopener nofollow ugc">Understanding Shapes in PyTorch Distributions Package</a></h3>

  <p>Sample shape, batch shape, and event shape in torch.distributions</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<p><a href="https://ericmjl.github.io/blog/2019/5/29/reasoning-about-shapes-and-probability-distributions/" class="onebox" target="_blank" rel="noopener nofollow ugc">https://ericmjl.github.io/blog/2019/5/29/reasoning-about-shapes-and-probability-distributions/</a></p>
<p>I also understand the output shapes of sampling the batched distributions. However, I am somewhat lost as to the input shape requirements of the <code>log_prob</code> method.</p>
<p>Here’s an example implementation that I just got right through random permutation until it worked… I still have no idea how I did it:</p>
<pre><code class="lang-auto"># 1D mesh. Each point on the mesh will be the mean of a different distribution
mesh = torch.arange(0, 10)

scale = torch.ones_like(mesh, dtype=torch.float32)
scale += torch.normal(mean=0, std=torch.arange(0.1, 0.2, 0.01))

x_range = torch.arange(-5, 15, 0.1)
results = torch.exp(gaussians.log_prob(x_range))

for i, res in enumerate(results.T):
    sns.lineplot(x=x_range, y=res)
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/6/06229cc711df412e77fb52177e8805e11d3d3912.png" data-download-href="https://discuss.pytorch.org/uploads/default/06229cc711df412e77fb52177e8805e11d3d3912" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/6/06229cc711df412e77fb52177e8805e11d3d3912.png" alt="image" data-base62-sha1="Sh1kVPCUMIcsZj2ktVa1Zo2Kye" width="690" height="24" data-dominant-color="F2D7D7"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">976×35 2.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
Okay, so I need to unsqueeze the input and it works fine…</p>
<pre><code class="lang-auto">results = torch.exp(gaussians.log_prob(x_range.unsqueeze(1)))
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/9/a/9a87ed31767546dbf707de3307e553ed8e644ca7.png" alt="image" data-base62-sha1="m32KBzoVuJ5Pk86U43S8xvEH5cz" width="369" height="251"></p>
<p>The 0th dimension is the number of samples. But what is this new dimension?</p>
<p>And what if I have a batch of samples:</p>
<pre><code class="lang-auto">x_range = torch.arange(-5, 15, 0.1)
x_range = torch.stack([x_range, x_range + 1]) # &lt;--- Two batches

results = torch.exp(gaussians.log_prob(x_range))

for i, res in enumerate(results.T):
    sns.lineplot(x=x_range, y=res)
</code></pre>
<p>Again I get</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/c/acea112228c13c1dda55dfa21e9ac28eef3193fc.png" data-download-href="https://discuss.pytorch.org/uploads/default/acea112228c13c1dda55dfa21e9ac28eef3193fc" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/c/acea112228c13c1dda55dfa21e9ac28eef3193fc.png" alt="image" data-base62-sha1="oFFBSFA18gyYUL9oAv0sw1bASny" width="690" height="18" data-dominant-color="EECFCF"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1002×27 2.75 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>So I unsqueeze the last dimension again. Now I’m dealing with three-dimensional input and I have no idea what two of the dimension are, nor the dimensions of the output.</p>
<p>Here is what the docs say:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/e/3e62740b9aa0ccb157358e5f9498826222fc40af.png" data-download-href="https://discuss.pytorch.org/uploads/default/3e62740b9aa0ccb157358e5f9498826222fc40af" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/e/3e62740b9aa0ccb157358e5f9498826222fc40af_2_690x64.png" alt="image" data-base62-sha1="8TSxnAwP6pk1Zh1bnNv9vH3nPLF" width="690" height="64" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/e/3e62740b9aa0ccb157358e5f9498826222fc40af_2_690x64.png, https://discuss.pytorch.org/uploads/default/original/3X/3/e/3e62740b9aa0ccb157358e5f9498826222fc40af.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/3/e/3e62740b9aa0ccb157358e5f9498826222fc40af.png 2x" data-dominant-color="F7F7F9"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1020×96 3.06 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Nothing…</p>