How can I merge diagonal and off-diagonal matrix elements into a single matrix?

AlphaBetaGamma96: <p>Hi All,</p>
<p>I have a quick question about merging 2 tensors together. Let’s say I have 1 tensor of shape [N,] which represents the diagonal of an N by N matrix. And another matrix which is of shape [N, N-1] and represents the off-diagonal elements, is it at all possible to merge the two different tensors together in an efficient way?</p>
<p>For example, the two tensors below (where I’ve given the values as the indices for ease of reading) would be merged into a single tensor (defined as <code>matrix</code>)</p>
<pre><code class="lang-auto">diag = torch.Tensor([11,22,33,44])
off_diag = torch.Tensor([[12,13,14],
                         [21,23,24],
                         [31,32,34],
                         [41,42,43]])
matrix = merge_diag_and_offdiag(diag, off_diag)
#returns torch.tensor([[11,12,13,14],
#                     [21,22,23,24],
#                     [31,32,33,34],
#                     [41,42,43,44]])
</code></pre>
<p>Is this at all possible? Any help is appreicated! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

ptrblck: <p>I’m sure there are better ways, but this should work:</p>
<pre><code class="lang-python">tmp = torch.cat((diag[:-1].unsqueeze(1), off_diag.view(3, 4)), dim=1)
res = torch.cat((tmp.view(-1), diag[-1].unsqueeze(0))).view(4, 4)
print(res)
&gt; tensor([[11, 12, 13, 14],
          [21, 22, 23, 24],
          [31, 32, 33, 34],
          [41, 42, 43, 44]])
</code></pre>

AlphaBetaGamma96: <p>Thank you once again, <a class="mention" href="/u/ptrblck">@ptrblck</a> !</p>
<p>I <em>think</em> I’ve managed to take your example and generalize it for a batch of vectors (and even a batch of multiple vectors, i.e input of shape [B,D,A] where B is the batch, D is the multiple vectors, vector side A, so the input’d be [B, D, A])</p>
<p>The code’s here,</p>
<pre><code class="lang-auto">def _merge_on_and_off_diagonal(on_diag, off_diag):
  #store output shape, to remove and unsqueeze'd dims
  output_shape = (*off_diag.shape[:-1],off_diag.shape[-2])
  if(len(on_diag.shape)==1 and len(off_diag.shape)==2):
    on_diag=on_diag.unsqueeze(0).unsqueeze(1)
    off_diag=off_diag.unsqueeze(0).unsqueeze(1)
  elif(len(on_diag.shape)==2 and len(off_diag.shape)==3):
    on_diag=on_diag.unsqueeze(1)
    off_diag=off_diag.unsqueeze(1)
  #reform input on_diag and off_diag to shape
  #B = batch, D = number of vector for given batch, A = dimension of vector
  #on_diag shape:  [B, D, A]
  #off_diag shape: [B, D, A, A-1]
  if(on_diag.shape[-1] != off_diag.shape[-2]):
    raise ValueError("index on_diag.shape[-1] must match off_diag.shape[-2]")
  dim=len(on_diag.shape)
  tmp = torch.cat((on_diag[:,:,:-1].unsqueeze(dim), \
            off_diag.view((*off_diag.shape[0:(dim-1)],off_diag.shape[-1], off_diag.shape[-2]))), dim=dim)
  res = torch.cat( (tmp.view(*off_diag.shape[0:(dim-1)], -1), on_diag[:,:,-1].unsqueeze(2)), dim=dim-1 ).view(*off_diag.shape[0:(dim-1)], on_diag.shape[-1], on_diag.shape[-1])

  return res.view(output_shape)
</code></pre>
<p>It’s not exactly pleasant but it seems to work for varying input sizes! <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=10" title=":smiley:" class="emoji" alt=":smiley:"><br>
For example,</p>
<pre><code class="lang-auto">diag = torch.tensor([11,22,33,44])
off_diag = torch.tensor([[12,13,14],
                         [21,23,24],
                         [31,32,34],
                         [41,42,43]])
matrix = _merge_on_and_off_diagonal(diag, off_diag)                 
"""
returns torch.tensor([[11,12,13,14],
                     [21,22,23,24],
                     [31,32,33,34],
                     [41,42,43,44]])
"""
  
diag = torch.tensor([[11,22,33,44],
                    [11,22,33,44]])
off_diag = torch.tensor([[[12,13,14],
                         [21,23,24],
                         [31,32,34],
                         [41,42,43]],
                         [[12,13,14],
                         [21,23,24],
                         [31,32,34],
                         [41,42,43]]])
matrix = _merge_on_and_off_diagonal(diag, off_diag)     

"""
returns tensor([[[11, 12, 13, 14],
         [21, 22, 23, 24],
         [31, 32, 33, 34],
         [41, 42, 43, 44]],

        [[11, 12, 13, 14],
         [21, 22, 23, 24],
         [31, 32, 33, 34],
         [41, 42, 43, 44]]])

"""
</code></pre>
<p>Thank you! <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=10" title=":smiley:" class="emoji" alt=":smiley:"></p>