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! 
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!