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

1

Hi All,

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?

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 matrix)

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]])

Is this at all possible? Any help is appreicated! :slight_smile:

3

I’m sure there are better ways, but this should work:

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)
> tensor([[11, 12, 13, 14],
          [21, 22, 23, 24],
          [31, 32, 33, 34],
          [41, 42, 43, 44]])
2 Likes
4

Thank you once again, @ptrblck !

I think 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])

The code’s here,

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)

It’s not exactly pleasant but it seems to work for varying input sizes! :smiley:
For example,

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]]])

"""

Thank you! :smiley:

2 Likes