I have a tensor T of shape (x, y) and a list L of shape (x), containing numbers [0, y). I want to have a tensor of shape (x), where each the ith element is T[i, L[i]]. I need this to be differentiable, and ideally with only pure torch operations and without loops. Apparently, T[:, L] gives a tensor of shape (x, x), where every element of L indexes every row.
Currently, I’m using this horrendous workaround, which feels extremely inefficient:
R = (T[:, L] * torch.eye(x)).sum(dim=0)