Is it possible to select a partial view of tensor with same indices from multiple dimensions?

1

Is there an alternative to using torch.diagonal() which does not alter the ordering of dimensions of a tensor?
For example, I have a 4D tensor, A, of shape (2,4,4,6). I want B = A[:,i,i,:] for all i in range(4).
I did the following in PyTorch:

tmp = torch.diagonal(A,dim1=1,dim2=2)
B = tmp.permute(0,2,1).contiguous()

Is there a way to avoid the permute() operation?
Many thanks. :slight_smile:

2

torch.zeros(2,4,4,6).diagonal(0,1,2).shape
torch.Size([2, 6, 4])
torch.zeros(2,4,4,6).diagonal(0,1,2).stride()
(96, 1, 30)

Note the peculiar stride 30, diagonal() is returning a specially constructed view onto input’s memory. Hence, contiguous() is needed regardless of dimensions’ ordering, if you want to copy & compact the output.

1 Like