Can I slice tensors with logical indexing or lists of indices?

Question

I'm trying to slice a PyTorch tensor using a logical index on the columns. I want the columns that correspond to a 1 value in the index vector. Both slicing and logical indexing are possible, but are they possible together? If so, how? My attempt keeps throwing the unhelpful error

TypeError: indexing a tensor with an object of type ByteTensor. The only supported types are integers, slices, numpy scalars and torch.LongTensor or torch.ByteTensor as the only argument.

MCVE

Desired Output

import torch

C = torch.LongTensor([[1, 3], [4, 6]])
# 1 3
# 4 6

Logical indexing on the columns only:

A_log = torch.ByteTensor([1, 0, 1]) # the logical index
B = torch.LongTensor([[1, 2, 3], [4, 5, 6]])
C = B[:, A_log] # Throws error

If the vectors are the same size, logical indexing works:

B_truncated = torch.LongTensor([1, 2, 3])
C = B_truncated[A_log]

And I can get the desired result by repeating the logical index so that it has the same size as the tensor I am indexing, but then I also have to reshape the output.

C = B[A_log.repeat(2, 1)] # [torch.LongTensor of size 4]
C = C.resize_(2, 2)

I also tried using a list of indices:

A_idx = torch.LongTensor([0, 2]) # the index vector
C = B[:, A_idx] # Throws error

If I want contiguous ranges of indices, slicing works:

C = B[:, 1:2]

score 28 · Accepted Answer · edited Apr 26 '23 at 02:13

28

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

edited Apr 26 '23 at 02:13

Madhav Malhotra

385
1
5
17

answered May 17 '17 at 09:44

dontloo

10,067
4
29
50

score 6 · Answer 2 · edited Mar 19 '21 at 09:10

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]])

edited Mar 19 '21 at 09:10

iacob

20,084
6
92
119

answered Jun 04 '20 at 04:35

user650654

5,630
3
41
44

1

Great that this is now more straightforward. Did this change in 1.5 or before that? – Cecilia Jun 04 '20 at 15:31

score 2 · Answer 3 · edited Mar 12 '21 at 12:50

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])

Can I slice tensors with logical indexing or lists of indices?

MCVE

3 Answers3

Linked