如何用高维张量索引张量

Question

我有两个张量 x 用于值，y 用于索引。 x.shape 和 y.shape 除了最后一个维度是相同的。 例如：

x=torch.tensor([[1, 6, 7, 5, 6],
        [8, 6, 7, 8, 4],
        [2, 8, 3, 5, 6]])
# x.shape:(3,5)
y=torch.tensor([[1, 2],
        [2, 3],
        [2, 2]])
# y.shape:(3,2)

有没有一种简单的方法可以将它切片x[y]所以结果是：

torch.tensor([6,7],[7,8],[3,3])

如果 x 和 y 是高维张量会怎样：

x.shape=(a,b,c,d)
y.shape=(a,b,c,e)
# a, b, c, d, e are positive integer

Answer 1

这正是torch.gather 所做的。

>> torch.gather(x, -1, y) # index along the last dimension
tensor([[6, 7],
        [7, 8],
        [3, 3]])

完全相同的命令也适用于您的高维案例。