<h2 dir="auto"><g-emoji class="g-emoji" alias="bug" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f41b.png">🐛</g-emoji> Bug</h2>
<p dir="auto">I think the Kl divergence is summing over the batch dimensions when it shouldn't, at least for the Gaussian case</p>
<p dir="auto">Here's the transformed distribution function:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="def _kl_transformed_transformed(p, q):
    if p.transforms != q.transforms:
        raise NotImplementedError
    if p.event_shape != q.event_shape:
        raise NotImplementedError
    # extra_event_dim = len(p.event_shape) - len(p.base_dist.event_shape)
    extra_event_dim = len(p.event_shape)
    base_kl_divergence = kl_divergence(p.base_dist, q.base_dist) #call to indep_indep below
   #this will again sum over kl_divergence for each entry in batch
    return _sum_rightmost(base_kl_divergence, extra_event_dim)"><pre class="notranslate"><code class="notranslate">def _kl_transformed_transformed(p, q):
    if p.transforms != q.transforms:
        raise NotImplementedError
    if p.event_shape != q.event_shape:
        raise NotImplementedError
    # extra_event_dim = len(p.event_shape) - len(p.base_dist.event_shape)
    extra_event_dim = len(p.event_shape)
    base_kl_divergence = kl_divergence(p.base_dist, q.base_dist) #call to indep_indep below
   #this will again sum over kl_divergence for each entry in batch
    return _sum_rightmost(base_kl_divergence, extra_event_dim)
</code></pre></div>
<p dir="auto">Here's  independent_independent KL</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="@register_kl(Independent, Independent)
def _kl_independent_independent(p, q):
    shared_ndims = min(p.reinterpreted_batch_ndims, q.reinterpreted_batch_ndims)
    p_ndims = p.reinterpreted_batch_ndims - shared_ndims
    q_ndims = q.reinterpreted_batch_ndims - shared_ndims
    p = Independent(p.base_dist, p_ndims) if p_ndims else p.base_dist
    q = Independent(q.base_dist, q_ndims) if q_ndims else q.base_dist
    kl = kl_divergence(p, q)
    if shared_ndims:
       #this line gets called when base_dist is Gaussian
        kl = sum_rightmost(kl, shared_ndims)
    return kl"><pre class="notranslate"><code class="notranslate">@register_kl(Independent, Independent)
def _kl_independent_independent(p, q):
    shared_ndims = min(p.reinterpreted_batch_ndims, q.reinterpreted_batch_ndims)
    p_ndims = p.reinterpreted_batch_ndims - shared_ndims
    q_ndims = q.reinterpreted_batch_ndims - shared_ndims
    p = Independent(p.base_dist, p_ndims) if p_ndims else p.base_dist
    q = Independent(q.base_dist, q_ndims) if q_ndims else q.base_dist
    kl = kl_divergence(p, q)
    if shared_ndims:
       #this line gets called when base_dist is Gaussian
        kl = sum_rightmost(kl, shared_ndims)
    return kl
</code></pre></div>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">I think something like this will do it:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="p = Gaussian(torch.zeros(3, 10), torch.ones(3, 10)).independent(1)
q = Gaussian(torch.ones(3,10) *2, torch.ones(3,10).independent(1)
kl = kl_divergence(p, q)
print(kl.shape) #should be something like [3, 1] but will instead output []"><pre class="notranslate"><code class="notranslate">p = Gaussian(torch.zeros(3, 10), torch.ones(3, 10)).independent(1)
q = Gaussian(torch.ones(3,10) *2, torch.ones(3,10).independent(1)
kl = kl_divergence(p, q)
print(kl.shape) #should be something like [3, 1] but will instead output []
</code></pre></div>
<h2 dir="auto">expected behavior</h2>
<p dir="auto">Maybe I misunderstand kl_divergence function, but I don't think it should be summing over batch.</p>
<ul dir="auto">
<li>PyTorch Version (e.g., 1.0): 1.2</li>
<li>OS (e.g., Linux): Ubuntu 16.04</li>
<li>How you installed PyTorch (<code class="notranslate">conda</code>, <code class="notranslate">pip</code>, source): conda/ pip</li>
<li>Build command you used (if compiling from source):</li>
<li>Python version: 3.6.9</li>
<li>CUDA/cuDNN version: 9.0</li>
<li>GPU models and configuration:</li>
<li>Any other relevant information:<br>
I originally found this using the MeanFieldTrace(...) class in the Pyro probabilistic programming language with normalizing flows. The Kl divergence function is used underneath the hood, and checks that the kl_divergence dimensions are the same as the q.batch_shape which in my case they were not.</li>
</ul>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/vincentqb/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/vincentqb">@vincentqb</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/fritzo/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/fritzo">@fritzo</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/neerajprad/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/neerajprad">@neerajprad</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/alicanb/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/alicanb">@alicanb</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/vishwakftw/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/vishwakftw">@vishwakftw</a></p>
<p dir="auto">This is a copy of an issue I posted for pytorch: <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="521900193" data-permission-text="Title is private" data-url="https://github.com/pytorch/pytorch/issues/29698" data-hovercard-type="issue" data-hovercard-url="/pytorch/pytorch/issues/29698/hovercard" href="https://github.com/pytorch/pytorch/issues/29698">pytorch/pytorch#29698</a></p>
<p dir="auto">I think this is an issue with the _kl_independent_independent class which is summing over batches when it shouldn't...</p>
<h3 dir="auto">Guidelines</h3>
<p dir="auto">I think the Kl divergence is summing over the batch dimensions when it shouldn't, at least for the Gaussian case</p>
<p dir="auto">Here's the transformed distribution function:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="def _kl_transformed_transformed(p, q):
    if p.transforms != q.transforms:
        raise NotImplementedError
    if p.event_shape != q.event_shape:
        raise NotImplementedError
    # extra_event_dim = len(p.event_shape) - len(p.base_dist.event_shape)
    extra_event_dim = len(p.event_shape)
    base_kl_divergence = kl_divergence(p.base_dist, q.base_dist) #call to indep_indep below
   #this will again sum over kl_divergence for each entry in batch
    return _sum_rightmost(base_kl_divergence, extra_event_dim)"><pre class="notranslate"><span class="pl-k">def</span> <span class="pl-en">_kl_transformed_transformed</span>(<span class="pl-s1">p</span>, <span class="pl-s1">q</span>):
    <span class="pl-k">if</span> <span class="pl-s1">p</span>.<span class="pl-s1">transforms</span> <span class="pl-c1">!=</span> <span class="pl-s1">q</span>.<span class="pl-s1">transforms</span>:
        <span class="pl-k">raise</span> <span class="pl-v">NotImplementedError</span>
    <span class="pl-k">if</span> <span class="pl-s1">p</span>.<span class="pl-s1">event_shape</span> <span class="pl-c1">!=</span> <span class="pl-s1">q</span>.<span class="pl-s1">event_shape</span>:
        <span class="pl-k">raise</span> <span class="pl-v">NotImplementedError</span>
    <span class="pl-c"># extra_event_dim = len(p.event_shape) - len(p.base_dist.event_shape)</span>
    <span class="pl-s1">extra_event_dim</span> <span class="pl-c1">=</span> <span class="pl-en">len</span>(<span class="pl-s1">p</span>.<span class="pl-s1">event_shape</span>)
    <span class="pl-s1">base_kl_divergence</span> <span class="pl-c1">=</span> <span class="pl-en">kl_divergence</span>(<span class="pl-s1">p</span>.<span class="pl-s1">base_dist</span>, <span class="pl-s1">q</span>.<span class="pl-s1">base_dist</span>) <span class="pl-c">#call to indep_indep below</span>
   <span class="pl-c">#this will again sum over kl_divergence for each entry in batch</span>
    <span class="pl-k">return</span> <span class="pl-en">_sum_rightmost</span>(<span class="pl-s1">base_kl_divergence</span>, <span class="pl-s1">extra_event_dim</span>)</pre></div>
<p dir="auto">Here's  independent_independent KL</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="@register_kl(Independent, Independent)
def _kl_independent_independent(p, q):
    shared_ndims = min(p.reinterpreted_batch_ndims, q.reinterpreted_batch_ndims)
    p_ndims = p.reinterpreted_batch_ndims - shared_ndims
    q_ndims = q.reinterpreted_batch_ndims - shared_ndims
    p = Independent(p.base_dist, p_ndims) if p_ndims else p.base_dist
    q = Independent(q.base_dist, q_ndims) if q_ndims else q.base_dist
    kl = kl_divergence(p, q)
    if shared_ndims:
       #this line gets called when base_dist is Gaussian
        kl = sum_rightmost(kl, shared_ndims)
    return kl"><pre class="notranslate"><span class="pl-en">@<span class="pl-en">register_kl</span>(<span class="pl-v">Independent</span>, <span class="pl-v">Independent</span>)</span>
<span class="pl-k">def</span> <span class="pl-en">_kl_independent_independent</span>(<span class="pl-s1">p</span>, <span class="pl-s1">q</span>):
    <span class="pl-s1">shared_ndims</span> <span class="pl-c1">=</span> <span class="pl-en">min</span>(<span class="pl-s1">p</span>.<span class="pl-s1">reinterpreted_batch_ndims</span>, <span class="pl-s1">q</span>.<span class="pl-s1">reinterpreted_batch_ndims</span>)
    <span class="pl-s1">p_ndims</span> <span class="pl-c1">=</span> <span class="pl-s1">p</span>.<span class="pl-s1">reinterpreted_batch_ndims</span> <span class="pl-c1">-</span> <span class="pl-s1">shared_ndims</span>
    <span class="pl-s1">q_ndims</span> <span class="pl-c1">=</span> <span class="pl-s1">q</span>.<span class="pl-s1">reinterpreted_batch_ndims</span> <span class="pl-c1">-</span> <span class="pl-s1">shared_ndims</span>
    <span class="pl-s1">p</span> <span class="pl-c1">=</span> <span class="pl-v">Independent</span>(<span class="pl-s1">p</span>.<span class="pl-s1">base_dist</span>, <span class="pl-s1">p_ndims</span>) <span class="pl-k">if</span> <span class="pl-s1">p_ndims</span> <span class="pl-k">else</span> <span class="pl-s1">p</span>.<span class="pl-s1">base_dist</span>
    <span class="pl-s1">q</span> <span class="pl-c1">=</span> <span class="pl-v">Independent</span>(<span class="pl-s1">q</span>.<span class="pl-s1">base_dist</span>, <span class="pl-s1">q_ndims</span>) <span class="pl-k">if</span> <span class="pl-s1">q_ndims</span> <span class="pl-k">else</span> <span class="pl-s1">q</span>.<span class="pl-s1">base_dist</span>
    <span class="pl-s1">kl</span> <span class="pl-c1">=</span> <span class="pl-en">kl_divergence</span>(<span class="pl-s1">p</span>, <span class="pl-s1">q</span>)
    <span class="pl-k">if</span> <span class="pl-s1">shared_ndims</span>:
       <span class="pl-c">#this line gets called when base_dist is Gaussian</span>
        <span class="pl-s1">kl</span> <span class="pl-c1">=</span> <span class="pl-en">sum_rightmost</span>(<span class="pl-s1">kl</span>, <span class="pl-s1">shared_ndims</span>)
    <span class="pl-k">return</span> <span class="pl-s1">kl</span></pre></div>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">I think something like this will do it:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="p = Gaussian(torch.zeros(3, 10), torch.ones(3, 10)).independent(1)
q = Gaussian(torch.ones(3,10) *2, torch.ones(3,10).independent(1)
kl = kl_divergence(p, q)
print(kl.shape) #should be something like [3, 1] but will instead output []"><pre class="notranslate"><span class="pl-s1">p</span> <span class="pl-c1">=</span> <span class="pl-v">Gaussian</span>(<span class="pl-s1">torch</span>.<span class="pl-en">zeros</span>(<span class="pl-c1">3</span>, <span class="pl-c1">10</span>), <span class="pl-s1">torch</span>.<span class="pl-en">ones</span>(<span class="pl-c1">3</span>, <span class="pl-c1">10</span>)).<span class="pl-en">independent</span>(<span class="pl-c1">1</span>)
<span class="pl-s1">q</span> <span class="pl-c1">=</span> <span class="pl-v">Gaussian</span>(<span class="pl-s1">torch</span>.<span class="pl-en">ones</span>(<span class="pl-c1">3</span>,<span class="pl-c1">10</span>) <span class="pl-c1">*</span><span class="pl-c1">2</span>, <span class="pl-s1">torch</span>.<span class="pl-en">ones</span>(<span class="pl-c1">3</span>,<span class="pl-c1">10</span>).<span class="pl-en">independent</span>(<span class="pl-c1">1</span>)
<span class="pl-s1">kl</span> <span class="pl-c1">=</span> <span class="pl-en">kl_divergence</span>(<span class="pl-s1">p</span>, <span class="pl-s1">q</span>)
<span class="pl-s1">print</span>(<span class="pl-s1">kl</span>.<span class="pl-s1">shape</span>) <span class="pl-c">#should be something like [3, 1] but will instead output []</span></pre></div>
<h2 dir="auto">expected behavior</h2>
<p dir="auto">Maybe I misunderstand kl_divergence function, but I don't think it should be summing over batch.</p>
<ul dir="auto">
<li>PyTorch Version (e.g., 1.0): 1.2</li>
<li>OS (e.g., Linux): Ubuntu 16.04</li>
<li>How you installed PyTorch (<code class="notranslate">conda</code>, <code class="notranslate">pip</code>, source): conda/ pip</li>
<li>Build command you used (if compiling from source):</li>
<li>Python version: 3.6.9</li>
<li>CUDA/cuDNN version: 9.0</li>
<li>GPU models and configuration:</li>
<li>Any other relevant information:<br>
I originally found this using the MeanFieldTrace(...) class in the Pyro probabilistic programming language with normalizing flows. The Kl divergence function is used underneath the hood, and checks that the kl_divergence dimensions are the same as the q.batch_shape which in my case they were not.</li>
</ul>