Explicit slicing across a particular dimension
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

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]])

