Masked_select vs indexing with boolean in aten


(dashesy) #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)