Let's say you have a 2D-square Tensor:
x = torch.tensor([[ 0, 1, 2, 3, 4, 5],
[ 6, 7, 8, 9, 10, 11],
[12, 13, 14, 15, 16, 17],
[18, 19, 20, 21, 22, 23],
[24, 25, 26, 27, 28, 29],
[30, 31, 32, 33, 34, 35]])
And you want to select the sub-tensor with rows and columns of index 0, 2, 3 considering you have a tensor keep such as:
keep = torch.tensor([True, False, True, True, False, False])
The desired output is then:
tensor([[ 0, 2, 3],
[12, 14, 15],
[18, 20, 21]])
Something that does not work
I expected x[keep, keep] to work but it only selects elements on the diagonal.
Making it work the long way - Masks
One way is to use a mask but it is quite tedious:
mask = keep.view(-1, 1) * keep
submatrix_size = keep.sum()
x[mask].view(sub_matrix_size, -1)
Making it work the short way
Another way to do it is:
x[keep][:, keep]
My question is then: Is the short way the best way to select on both dimensions with the same boolean tensor ? Is there any other way to do it in PyTorch?