Can I slice tensors with logical indexing or lists of indices?
I tried this snippet of code, and I wrote the results as a comment next to it.
import torch
arr = torch.tensor([[0,1,2],[3,4,5]])
arr = torch.arange(6).reshape((2,3))
print(arr)
# tensor([[0, 1, 2],
# [3, 4, 5]])
print(arr[1]) # tensor([3, 4, 5])
print(arr[1,1]) # tensor(4)
print(arr[1, :]) # tensor([3, 4, 5])
#print(arr[1,1,1]) # IndexError: too many indices for tensor of dimension 2
print(arr[1, [0,1]]) # tensor([3, 4])
print(arr[[0, 1],0]) # tensor([0, 3])
I think this is implemented as the index_select
function, you can try
import torch
A_idx = torch.LongTensor([0, 2]) # the index vector
B = torch.LongTensor([[1, 2, 3], [4, 5, 6]])
C = B.index_select(1, A_idx)
# 1 3
# 4 6
In PyTorch 1.5.0, tensors used as indices must be long, byte or bool tensors.
The following is an index as a tensor of longs.
import torch
B = torch.LongTensor([[1, 2, 3], [4, 5, 6]])
idx1 = torch.LongTensor([0, 2])
B[:, idx1]
# tensor([[1, 3],
# [4, 6]])
And here's a tensor of bools (logical indexing):
idx2 = torch.BoolTensor([True, False, True])
B[:, idx2]
# tensor([[1, 3],
# [4, 6]])