How to Perform 2d indexing, gather In Pytorch?

Use case:

I am writing beam search submodule in an NLP program which requires 2d indexing for my imlementation. The setup is:

• k_next_words is shape (batch_size, d1, d2) = (4,5,6)
• ix is shape (batch_size, k_num_to_select, ix_dim=2) = (4,7,2), where the 4 dim is batch size, the 7 dim is number of samples to look up, and the 2 dim is a 2d index corresponding to d1 and d2 in k_next_words (such that the first entry in the index is between 0 and 4 and the second entry is between 0 and 5).

Problem: We want to, for each of the 4 batches gather the 7 entries corresponding to the elements of k_next_words at the 2d index.

This is probably a use case for gather, but I haven’t gotten it to work so far!

Example:

Here is a toy example to illustrate my use case. Let's ignore the batch size (=4) for now.

k_next_words is a torch.Tensor of size (5,6) which looks like:

[[1,  2,  3,  4,  5,  6]
[11, 12, 13, 14, 15, 16]
[21, 22, 23, 24, 25, 26]
[31, 32, 33, 34, 35, 36]
[41, 42, 43, 44, 45, 46]]

ix is a (7,2) tensor of coordinates we want to look up in k_next_words. As you can see, the first entry in the last dimension is always between 0-4 and the second entry in the last dimension is always between 0-5

[[4,5],
 [1,0],
 [0,2],
 [2,2],
 [2,1],
 [3,2],
 [1,1]]

Given these two inputs, I want to look up the 7 entries of k_next_words determined by the indices in ix (note: indexing start at 0), that is, in the toy example:

[[46],
 [11],
 [3], 
 [23],
 [22],
 [33],
 [12]]

However, I want to do the above, but also for a given batch size (=4) in the example.