for i in ...:
    for j in ...:
        for k in ...:
            outputs[i][j][k] = inputs[indices[i][j]][indices[i][k]]

is there any indexing method can help me do this?
i is the batch dim
UPDATE
I solved this by inputs[indices] and then applying gather.
thx anyway