user6010932
user6010932

Reputation: 21

How to slice a Torch tensor by a set of indices?

I am wondering if there is an efficient way to slice a tensor by a set of indices? In Matlab, I can do the following:

A = rand(2, 3, 6);
B = A(:,:, 1:2:end);

Then B is the 1st, 3rd, and 5th slice of A.

In Torch, it seems you can only slice with a continuous range. Is that true?

A more general question will be if I can get a subset by arbitrary indices such as

A(:, :, [1 2 6]) 

in Matlab.

Thanks in advance!

Upvotes: 2

Views: 1787

Answers (1)

deltheil
deltheil

Reputation: 16121

Use the index operator, e.g.:

t = torch.rand(2, 3, 6)
-- (1,.,.) = 
--   0.1790  0.7774  0.5343  0.0628  0.3077  0.7203
--   0.0677  0.5847  0.2401  0.6885  0.8724  0.4413
--   0.1849  0.2704  0.2745  0.5508  0.4634  0.6340
-- 
-- (2,.,.) = 
--   0.2523  0.6135  0.6037  0.0194  0.6456  0.0229
--   0.9966  0.8688  0.2078  0.7169  0.1528  0.5708
--   0.8671  0.7731  0.4596  0.0636  0.8873  0.2205
-- [torch.DoubleTensor of size 2x3x6]


t:index(3, torch.LongTensor{1, 3, 5})
-- (1,.,.) = 
--   0.1790  0.5343  0.3077
--   0.0677  0.2401  0.8724
--   0.1849  0.2745  0.4634
-- 
-- (2,.,.) = 
--   0.2523  0.6037  0.6456
--   0.9966  0.2078  0.1528
--   0.8671  0.4596  0.8873
-- [torch.DoubleTensor of size 2x3x3]

You can do as well t:index(3, torch.LongTensor{1, 2, 6}).

Upvotes: 2

Related Questions