Getting diagonal elements of matrices in batch

Samue1: <p>I have a tensor <code>T</code> with dimensions <code>(batch_size, size_n, size_n)</code></p>
<pre><code class="lang-auto">T = 
    [[[0.9527, 0.8821],
      [0.8442, 0.6147]],
     [[0.0672, 0.5737],
      [0.1963, 0.4532]],
     [[0.0992, 0.3838],
      [0.4169, 0.0925]]]
</code></pre>
<p>and want to extract the diagonal of each matrix in that batch to get</p>
<pre><code class="lang-auto">diag_T = 
    [[0.9527, 0.6147],
     [0.0672, 0.4532],
     [0.0992, 0.0925]]
</code></pre>
<p>Is there some <code>torch.diag()</code> function that also works for batches?</p>

LeviViana: <p>Maybe not the best solution, but it is vectorized:</p>
<pre><code class="lang-python">import torch

T = [[[0.9527, 0.8821],
      [0.8442, 0.6147]],
     [[0.0672, 0.5737],
      [0.1963, 0.4532]],
     [[0.0992, 0.3838],
      [0.4169, 0.0925]]]

T = torch.Tensor(T)

torch.einsum('ijj-&gt;ij', torch.stack(tuple(torch.ones(T.size(1)).diag() for i in range(T.size(0)))) * T)
</code></pre>

KFrank: <p>Hi Samuel!</p>
<aside class="quote no-group" data-username="Samue1" data-post="1" data-topic="126097" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/f04885/48.png" class="avatar"> Samue1:</div>
<blockquote>
<p>Is there some <code>torch.diag()</code> function that also works for batches?</p>
</blockquote>
</aside>
<p><a href="https://pytorch.org/docs/1.9.0/generated/torch.diagonal.html#torch.diagonal" rel="noopener nofollow ugc">torch.diagonal()</a> does what you want:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.9.0'
&gt;&gt;&gt; T = torch.tensor (
...     [[[0.9527, 0.8821],
...     [0.8442, 0.6147]],
...     [[0.0672, 0.5737],
...     [0.1963, 0.4532]],
...     [[0.0992, 0.3838],
...     [0.4169, 0.0925]]]
... )
&gt;&gt;&gt; torch.diagonal (T, dim1 = -2, dim2 = -1)
tensor([[0.9527, 0.6147],
        [0.0672, 0.4532],
        [0.0992, 0.0925]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

Samue1: <p>Why do we use <code>dim1=-2, dim2=-1</code> and not just <code>dim1=1, dim2=2</code>?</p>
<p>Do I miss something here?</p>

KFrank: <p>Hi Samuel!</p>
<aside class="quote no-group" data-username="Samue1" data-post="4" data-topic="126097" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/f04885/48.png" class="avatar"> Samue1:</div>
<blockquote>
<p>Why do we use <code>dim1=-2, dim2=-1</code> and not just <code>dim1=1, dim2=2</code>?</p>
</blockquote>
</aside>
<p>No particular reason – your <code>T</code> is a 3d tensor so the two versions are<br>
equivalent.  I suppose that using “negative” dimensions emphasizes<br>
that we’re extracting the diagonals from the 2d matrices made up by<br>
the last two dimensions of <code>T</code> (so that this version would generalize to a<br>
hypothetical use case where <code>T</code> had multiple leading “batch” dimensions<br>
such as <code>T</code> of shape <code>[batch_size, channel_size, size_n, size_n]</code>).</p>
<p>It’s really just stylistic – and not necessarily a better style.</p>
<p>Best.</p>
<p>K. Frank</p>