NumPy: pass subset of multi-dimensional indices to array
Suppose I have an array, and a (sub)set of indices:
arr = np.array([[[0,1,2],
                 [3,4,5],
                 [6,7,8]],
                
                [[9, 10,11],
                 [12,13,14],
                 [15,16,17]]])
idx = [1,2]

And I wish to get, for each element of dim=0 of the array, the respective idx slice, i.e.:
&gt;&gt;&gt; arr[:, idx[0], idx[1]]
array([ 5, 14])

Is there a way to do such slicing without hard-coding each index? Something like:
ar[:, *idx]


Note, the following is my current workaround:
idx = [slice(a.shape[0]), *idx]
a[idx]

but I was wondering if NumPy (or PyTorch, or a related library) supported a more 'elegant'/natural multi-dimensional indexing syntax for such cases.

I think your workaround is pretty much as close as you can get.
You can try this to make it slightly more concise:
import numpy as np

arr = np.array([[[0,1,2],
                 [3,4,5],
                 [6,7,8]],

                [[9, 10,11],
                 [12,13,14],
                 [15,16,17]]])
idx = [1,2]

print(
    arr[(slice(None), *idx)]
)

This works because : is the same as slice(None)
