NumPy/PyTorch slicing using an array of indices from an argmax call
I have the following tensors / ndarrays I am operating with.
a_intents of shape (b, n_i) - Scalar at position ij is the activation of intent j for the i-th element in the batch.
u_intents of shape (b, n_i, d_m) - Vector of dimension d_m at index ij is the pose vector of intent j for the i-th element in the batch.
I want to get the index of the intent with the biggest activation scalar, and I do:
 max_activations = argmax(a_intents, dim=-1, keepdim=False)

Now, using those indices, I want to get the corresponding vectors in u_intents.
max_activation_poses = u_intents[?, ?,:]

How do I use max_activations to indicate only those indices on dim 1? My intuition tells me I will end with an incorrect shape if I do
[:, max_activations, :]

The shape I am trying to get is (b, d_m) - The vector with the same index as the highest activation, for each element in the batch.
Thanks

It's easy if you treat the u_intents vector as 2-dimensional, and offset each argmax index by its batch index times the number of elements.
# dummy values for demonstration
&gt;&gt;&gt; b, n_i, d_m = 2, 3, 5
&gt;&gt;&gt; a_intents = torch.rand(b, n_i)
&gt;&gt;&gt; a_intents
tensor([[0.1733, 0.9965, 0.4790],
        [0.6056, 0.4040, 0.0236]])
&gt;&gt;&gt; u_intents = torch.rand(b, n_i, d_m)
&gt;&gt;&gt; u_intents
tensor([[[0.3301, 0.8153, 0.1356, 0.6623, 0.4385],
         [0.1902, 0.1748, 0.4131, 0.3887, 0.5363],
         [0.1211, 0.5773, 0.2405, 0.6313, 0.2064]],
        [[0.2592, 0.5127, 0.7301, 0.8883, 0.5665],
         [0.2767, 0.6545, 0.7595, 0.2677, 0.5163],
         [0.8158, 0.4940, 0.0492, 0.0911, 0.0465]]])
# add to each index the batch start
&gt;&gt;&gt; max_activations = a_intents.argmax(dim=-1) + torch.arange(0, b*n_i, step=n_i)
# elements 1 and 3 of 0..5
&gt;&gt;&gt; max_activations
tensor([1, 3])   
&gt;&gt;&gt; poses = u_intents.view(b*n_i, d_m).index_select(0, max_activations)
# tensor of shape (b, d_m) correctly indexing the maxima.
&gt;&gt;&gt; poses
tensor([[0.1902, 0.1748, 0.4131, 0.3887, 0.5363],
        [0.2592, 0.5127, 0.7301, 0.8883, 0.5665]])

