torch logical indexing of tensor

Question

I looking for an elegant way to select a subset of a torch tensor which satisfies some constrains. For example, say I have:

A = torch.rand(10,2)-1

and S is a 10x1 tensor,

sel = torch.ge(S,5) -- this is a ByteTensor

I would like to be able to do logical indexing, as follows:

A1 = A[sel]

But that doesn't work. So there's the index function which accepts a LongTensor but I could not find a simple way to convert S to a LongTensor, except the following:

sel = torch.nonzero(sel)

which returns a K x 2 tensor (K being the number of values of S >= 5). So then I have to convert it to a 1 dimensional array, which finally allows me to index A:

A:index(1,torch.squeeze(sel:select(2,1)))

This is very cumbersome; in e.g. Matlab all I'd have to do is

A(S>=5,:)

Can anyone suggest a better way?

Answer 1

One possible alternative is:

sel = S:ge(5):expandAs(A)   -- now you can use this mask with the [] operator
A1 = A[sel]:unfold(1, 2, 2) -- unfold to get back a 2D tensor

Example:

> A = torch.rand(3,2)-1
-0.0047 -0.7976
-0.2653 -0.4582
-0.9713 -0.9660
[torch.DoubleTensor of size 3x2]

> S = torch.Tensor{{6}, {1}, {5}}
 6
 1
 5
[torch.DoubleTensor of size 3x1]

> sel = S:ge(5):expandAs(A)
1  1
0  0
1  1
[torch.ByteTensor of size 3x2]

> A[sel]
-0.0047
-0.7976
-0.9713
-0.9660
[torch.DoubleTensor of size 4]

> A[sel]:unfold(1, 2, 2)
-0.0047 -0.7976
-0.9713 -0.9660
[torch.DoubleTensor of size 2x2]