Masked_select vs indexing with boolean in aten

C++
1

I want to use a boolean as index with aten, there is a masked_select that does the job but flattens the output (I can use view to correct that).
How is indexing with a boolean array implemented in aten (does it call masked_select and then change the view)?

IOW, masked_select returns a different shape than using [], what is the equivalent of [] in aten/C++?

    x = torch.randn(3, 4)
    mask = torch.tensor([[1], [0], [1]], dtype=torch.uint8)
    print(x.masked_select(mask).shape)
    print(x[mask.view(-1)].shape)
2 Likes
2

Hi,

Did you find out how to do this? I am having the same issue and I am pondering what the best method is (e.g. using masked_select and then view, or converting the boolean mask to int and doing index_select)

3

For anyone stumbling upon this topic, the solution is to use .index, e.g. :

auto boolean_selector = (A == B);
auto C = D.index({boolean_selector});

and for a two-dimensional array where one wants to select along the second dimension of D:

auto boolean_selector = (A == B).squeeze();
auto C = D.index({Slice(), boolean_selector});

with D of size [m x n] and A, B of size [n x 1].

1 Like