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

问题

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

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

Logical indexing on the columns only

import torch
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

I also tried using a list of indices

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

If the vectors are the same size, logical indexing works

import torch
A_log = torch.ByteTensor([1, 0, 1]) # the logical index
B = torch.LongTensor([1, 2, 3])
C = B[A_log]

If I use contiguous ranges of indices, slicing works

import torch
B = torch.LongTensor([[1, 2, 3], [4, 5, 6]])
C = B[:, 1:2]

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.

import torch
A_log = torch.ByteTensor([1, 0, 1]) # the logical index
B = torch.LongTensor([[1, 2, 3], [4, 5, 6]])
C = B[A_log.repeat(2, 1)] # [torch.LongTensor of size 4]
C = C.resize_(2, 2)

回答1:

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

来源：https://stackoverflow.com/questions/43989310/can-i-slice-tensors-with-logical-indexing-or-lists-of-indices

标签

python

matrix-indexing

pytorch