Explicit slicing across a particular dimension

Question

I've got a 3D tensor x (e.g 4x4x100). I want to obtain a subset of this by explicitly choosing elements across the last dimension. This would have been easy if I was choosing the same elements across last dimension (e.g. x[:,:,30:50] but I want to target different elements across that dimension using the 2D tensor indices which specifies the idx across third dimension. Is there an easy way to do this in numpy?

A simpler 2D example:

x = [[1,2,3,4,5,6],[10,20,30,40,50,60]]
indices = [1,3]

Let's say I want to grab two elements across third dimension of x starting from points specified by indices. So my desired output is:

[[2,3],[40,50]]

Update: I think I could use a combination of take() and ravel_multi_index() but some of the platforms that are inspired by numpy (like PyTorch) don't seem to have ravel_multi_index so I'm looking for alternative solutions

No, it contains indices that I want to take from `x` along the third dimension — Milad, Apr 08 '19 at 11:49
Why is `indices` a 2D tensor? Shouldn't it be a 1D tensor, e.g., `[0, 1, 10, 33, 50, ...]`? — Berriel, Apr 08 '19 at 11:52
I added an example where I'm not even slicing and I simply want to extract those elements. Does it make it clearer? — Milad, Apr 08 '19 at 11:56
Use a minimal **representative** sample data. So, something that needs slicing. — Divakar, Apr 08 '19 at 11:58
if `indices` is supposed to give the output in the nth dimension along those indeces wouldnt the output be `[[2,4],[20,40]]` — user2255757, Apr 08 '19 at 12:42

score 2 · Answer 1 · answered Apr 08 '19 at 16:48

Iterating over the idx, and collecting the slices is not a bad option if the number of 'rows' isn't too large (and the size of the sizes is relatively big).

In [55]: x = np.array([[1,2,3,4,5,6],[10,20,30,40,50,60]])                      
In [56]: idx = [1,3]                                                            
In [57]: np.array([x[j,i:i+2] for j,i in enumerate(idx)])                       
Out[57]: 
array([[ 2,  3],
       [40, 50]])

Joining the slices like this only works if they all are the same size.

An alternative is to collect the indices into an array, and do one indexing.

For example with a similar iteration:

idxs = np.array([np.arange(i,i+2) for i in idx])

But broadcasted addition may be better:

In [58]: idxs = np.array(idx)[:,None]+np.arange(2)                              
In [59]: idxs                                                                   
Out[59]: 
array([[1, 2],
       [3, 4]])
In [60]: x[np.arange(2)[:,None], idxs]                                          
Out[60]: 
array([[ 2,  3],
       [40, 50]])

ravel_multi_index is not hard to replicate (if you don't need clipping etc):

In [65]: np.ravel_multi_index((np.arange(2)[:,None],idxs),x.shape)              
Out[65]: 
array([[ 1,  2],
       [ 9, 10]])
In [66]: x.flat[_]                                                              
Out[66]: 
array([[ 2,  3],
       [40, 50]])
In [67]: np.arange(2)[:,None]*x.shape[1]+idxs                                   
Out[67]: 
array([[ 1,  2],
       [ 9, 10]])

user2255757 · Answer 2 · 2019-04-08T13:27:18.800

0

along the 3D axis:

x = [x[:,i].narrow(2,index,2) for i,index in enumerate(indices)]
x = torch.stack(x,dim=1)

by enumerating you get the index of the axis and index from where you want to start slicing in one.

narrow gives you a zero-copy length long slice from a starting index start along a certain axis

you said you wanted:

dim = 2
start = index
length = 2

then you simply have to stack these tensors back to a single 3D.

This is the least work intensive thing i can think of for pytorch.

EDIT

if you just want different indices along different axis and indices is a 2D tensor you can do:

x = [x[:,i,index] for i,index in enumerate(indices)]
x = torch.stack(x,dim=1)

You really should have given a proper working example, making it unnecessarily confusing.

edited Apr 08 '19 at 13:27

answered Apr 08 '19 at 12:48

user2255757

756
1
6
24

This works if you want the same indices for each `a,b` in `x[a,b,:]` but I want to have different indices for each element in the first two dimensions – Milad Apr 08 '19 at 12:57
how about this solution, didnt read you wanted `2` values starting from `index` – user2255757 Apr 08 '19 at 13:18
there is an edit if you already saw the answer, maybe this solution is more appropriate? – user2255757 Apr 08 '19 at 13:27

Paul Panzer · Answer 3 · 2019-04-08T14:16:56.410

Here is how to do it in numpy, now clue about torch, though.

The following picks a slice of length n along the third dimension starting from points idx depending on the other two dimensions:

# example
a = np.arange(60).reshape(2, 3, 10)
idx = [(1,2,3),(4,3,2)]
n = 4

# build auxiliary 4D array where the last two dimensions represent
# a sliding n-window of the original last dimension
j,k,l = a.shape
s,t,u = a.strides
aux = np.lib.stride_tricks.as_strided(a, (j,k,l-n+1,n), (s,t,u,u))

# pick desired offsets from sliding windows
aux[(*np.ogrid[:j, :k], idx)]
# array([[[ 1,  2,  3,  4],
#         [12, 13, 14, 15],
#         [23, 24, 25, 26]],

#        [[34, 35, 36, 37],
#         [43, 44, 45, 46],
#         [52, 53, 54, 55]]])

score 0 · Answer 4 · answered Apr 08 '19 at 18:27

I came up with below using broadcasting:

x = np.array([[1,2,3,4,5,6,7,8,9,10],[10,20,30,40,50,60,70,80,90,100]]) 
i = np.array([1,5])
N = 2 # number of elements I want to extract along each dimension. Starting points specified in i

r = np.arange(x.shape[-1])
r = np.broadcast_to(r, x.shape)

ii = i[:, np.newaxis]
ii = np.broadcast_to(ii, x.shape)

mask = np.logical_and(r-ii>=0, r-ii<=N) 

output = x[mask].reshape(2,3)

Does this look alright?

Explicit slicing across a particular dimension

4 Answers4