Hello,
I would like to ask for your help on making a vectorized version of the following simple operation.
I have a 3D tensor x of shape (N, B, V) and I would like to get the elements of x at indices given by two (N, K) tensors idx1 and idx2 as follows: y[i, j] = x[i, idx1[i,j], idx2[i,j]].
Using a for loop, this can be done using the following function:
def f(x, idx1, idx2):
    """Compute using for loop
        x: N x B x V
        idx1: N x K matrix where idx1[i, j] is between [0, B)
        idx2: N x K matrix where idx2[i, j] is between [0, V)
    Return:
        y: N x K matrix where y[i, j] = x[i, idx1[i,j], idx2[i,j]]
    """
    # N x K
    y = torch.zeros(idx1.shape)
    N, K = idx1.shape
    for i in range(N):
        for j in range(K):
            y[i, j] = x[i, idx1[i,j], idx2[i,j]]
    return y

It seems that torch.gather is not applicable here. If we do z = x[:,idx1,idx2], then it remains to do y[i,j] = z[i,i,j], which doesn’t seem to be any easier.
For your convenience the above function can be tested using the following:
def main():
    N = 2
    B = 3
    V = 10
    K = B*2
    x = torch.randn(N, B, V)
    idx1 = torch.randint(0, B, size=(N, K))
    idx2 = torch.randint(0, V, size=(N, K))
    y = f(x, idx1, idx2)

Thank you in advance for your help!