andresy/torch — pull request #113 by marcoscoffier (2013-02-21 20:49:17+00:00)
Title: select by index

Hey Guys,

A quick hack which I use often and is much cleaner in torch itself. Similar to how you can select from a tensor with a ByteTensor of 0,1 such as returned from torch.gt, lt etc. this change allows you to select from a tensor by index with a LongTensor. The tensor can be multi-dimensional but the index will select on the first dimension only eg.

> t = torch.randn(10,3)
> return t
-1.5676 -0.6461 2.1626
0.3110 -0.5192 1.0081
-0.6419 0.0705 1.3381
0.0507 0.2139 -0.0334
-1.3888 -0.9223 0.0692
0.6868 -1.2407 0.7173
0.6645 0.9348 -0.8542
0.5592 0.0246 2.5439
-1.2265 -1.6566 0.9136
-1.9968 0.3349 -0.7398
[torch.DoubleTensor of dimension 10x3]

> i = torch.LongTensor({1,3,5})
> return t[i]
-1.5676 -0.6461 2.1626
-0.6419 0.0705 1.3381
-1.3888 -0.9223 0.0692
[torch.DoubleTensor of dimension 3x3]

I do check the bounds:

> i[1] = 11
> =t[i]
bad argument #2 to '?' (indices out of bounds (> src elements))
stack traceback:
[C]: at 0x072eb490
[C]: in function '__index'
[string "return t[i]"]:1: in main chunk
[C]: at 0x0101452310

> i[1] = 0
> return t[i]
bad argument #2 to '?' (indices out of bounds (< 1))
stack traceback:
[C]: at 0x072eb490
[C]: in function '__index'
[string "return t[i]"]:1: in main chunk
[C]: at 0x0101452310

You can also make a crufty repmat

> i:fill(3)
> =t[i]
-0.6419 0.0705 1.3381
-0.6419 0.0705 1.3381
-0.6419 0.0705 1.3381
[torch.DoubleTensor of dimension 3x3]

And do things like sort rows:

> v,i = t:select(2,1):sort()
> =t[i]
-1.9968 0.3349 -0.7398
-1.5676 -0.6461 2.1626
-1.3888 -0.9223 0.0692
-1.2265 -1.6566 0.9136
-0.6419 0.0705 1.3381
0.0507 0.2139 -0.0334
0.3110 -0.5192 1.0081
0.5592 0.0246 2.5439
0.6645 0.9348 -0.8542
0.6868 -1.2407 0.7173
[torch.DoubleTensor of dimension 10x3]

Let me know if you have concerns, or where I should update the dok if this is needed also,

Thanks,

Marco