title: Applying Kullback-Leibler (aka kl divergence) element-wise in Pytorch
body: <p>I have two tensors named <code>x_t</code>, <code>x_k</code> with follwing shapes <code>NxHxW</code> and <code>KxNxHxW</code> respectively, where <code>K</code>, is the number of autoencoders used to reconstruct <code>x_t</code> (if you have no idea what is this, assume they're <code>K</code> different nets aiming to predict <code>x_t</code>, this <em>probably</em> has nothing to do with the question anyways) <code>N</code> is batch size, <code>H</code> matrix height, <code>W</code> matrix width.</p>
  
  <p>I'm trying to apply <a href="https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence" rel="nofollow noreferrer">Kullback-Leibler divergence</a> algorithm to both tensors (after broadcasting <code>x_t</code> as <code>x_k</code> along the <code>K</code><sup>th</sup> dimension) using <a href="https://pytorch.org/docs/stable/nn.html#torch.nn.functional.kl_div" rel="nofollow noreferrer">Pytorch's <code>nn.functional.kl_div</code></a> method. </p>
  
  <p>However, it does not seem to be working <em>as I expected</em>. I'm looking to calcualte the <code>kl_div</code> between each observation in <code>x_t</code> and <code>x_k</code> resulting in a tensor of size <code>KxN</code> (i.e., <code>kl_div</code> of each observation for each <code>K</code> autoencoder). </p>
  
  <p>The actual output is a <em>single value</em> if I use the <code>reduction</code> argument, and the same tensor size (i.e., <code>KxNxHxW</code>) if I do not use it. </p>
  
  <p>Has anyone tried something similar?</p>
  
  <hr>
  
  <p>Reproducible example:</p>
  
  <pre><code>import torch
  import torch.nn.functional as F
  #                  K   N   H  W
  x_t = torch.randn(    10, 5, 5)
  x_k = torch.randn( 3, 10, 5, 5)
  
  x_broadcasted = x_t.expand_as(x_k)
  
  loss = F.kl_div(x_t, x_k, reduction="none") # or "batchmean", or there are many options
  </code></pre>