How to index with list of indices in a multidimensional array in numpy/pytorch
I have an two multi-dimensional arrays of shape (1,169,5,1+4+4). One contains data, and the other contains zeros that is to be filled.
&gt;&gt;&gt; preds = np.random.random_integers(0,500,(1,169,5,1+4+4))
&gt;&gt;&gt; filtered_preds = np.zeros((1,169,5,1+4+4))

For context the dimensions represent the following:

batch size
number of grid cells
number of anchor boxes
predictions (Confidence score, bounding box coords, number of classes)

I have a third flattened array of shape (169) that contains indicies of the best anchor box
Which looks like the following:
&gt;&gt;&gt; best_bboxes = np.random.random_integers(0,4,(169,))

In order to prevent using loops, my first instinct was to try the following to place the confidence values:
&gt;&gt;&gt; filtered_preds[:,:,best_ab_idx,0:1] = preds[:,:,best_ab_idx,0:1]

However, upon inspecting the array, I do not get the disired output as the shape of filtered_preds[:,:,best_ab_idx,0:1] is (1,169,169,1)
I guess my question so far is, how do I get values from my prediction array with the size of the second dimension and use the value of the third dimension to get the desired value?
I'm still a bit new to array indexing, so any help would be appreciated.
Note: I am using PyTorch but gave my examples in NumPy as it follows the same principle.
Edit:
Here is a solution of what I am trying to acheive using a loop:
for idx,i in enumerate(best_bboxes):
    filtered_preds[:,idx,best_bboxes[idx],0:1] = preds[:,idx,best_bboxes[idx],0:1]
    filtered_preds[:,idx,best_bboxes[idx],1:5] = preds[:,idx,best_bboxes[idx],1:5]
    filtered_preds[:,idx,best_bboxes[idx],5:] = preds[:,idx,best_bboxes[idx],5:]


I'm a little confused by the description of what you want to achieve, but I can simplify your given solution with the following.
seq_idx = np.arange(best_bboxes.size)
filtered_preds[:, seq_idx, best_bboxes, :] = preds[:, seq_idx, best_bboxes, :]

