title: Sum all diagonals in feature maps in parallel in PyTorch
body: <p>Let's say I have a tensor shaped <code>(1, 64, 128, 128)</code> and I want to create a tensor of shape <code>(1, 64, 255)</code> holding the sums of all diagonals for every <code>(128, 128)</code> matrix (there are 1 main, 127 below, 127 above diagonals so in total 255). What I am currently doing is the following:</p>
  
  <p><div class="snippet" data-lang="js" data-hide="false" data-console="true" data-babel="false">
  <div class="snippet-code">
  <pre class="snippet-code-html lang-html prettyprint-override"><code>x = torch.rand(1, 64, 128, 128)
  
  diag_sums = torch.zeros(1, 64, 255)
  j = 0
  for k in range(-127, 128):
      diag_sums[j, :, k + 127] = torch.diagonal(x, offset=k, dim1=-2, dim2=-1).sum(dim=2)</code></pre>
  </div>
  </div>
  
  This is obviously very slow, since it is using Python loops and is not done in parallel with respect to k.</p>
  
  <p>I don't think this can be done using <code>torch.diagonal</code> since the function explicitly uses a single int for the offset parameter. If I could pass a list there, this would work, but I guess it would be complicated to implement (requiring changes in PyTorch itself).</p>
  
  <p>I think it could be possible to implement this using <code>torch.einsum</code>, but I cannot think of a way to do it.</p>
  
  <p>So this is my question: how do I get the tensor described above?</p>