Expand 1d tensor to 2d tensor along with the diagonal of the 2d tensor in pytorch

Viewed 334

Is there any way, in Pytorch, to convert a 1d tensor([[1., 2., 3., 4.]]) into

tensor([[1., 2., 3., 4.], 
        [2., 1., 2., 3.], 
        [3., 2., 1., 2.], 
        [4., 3., 2., 1.]]) 
1 Answers

torch.roll unfortunately only applies dimension-wise, so you will have to expand and use something like torch.gather, along with a set of index shifts:

t = tensor([[1., 2., 3., 4.]])

idx = torch.arange(4)
idx = torch.stack([torch.roll(idx, shift) for shift in range(0,-4,-1)])

torch.gather(t.expand(4,-1), 1, idx)
Related