如何根据相同大小的索引张量重新排序张量

Question

假设我有 tensor A，索引 Tensor：A = [1, 2, 3, 4]，indexes = [1, 0, 3, 2]

我想用这两个创建一个新的 Tensor，结果如下：[2, 1, 4, 3] 结果的每个元素都是来自 A 的元素，顺序由索引 Tensor 定义。有没有办法用 PyTorch tensor ops 没有循环来做到这一点？

我的目标是为 2D 做 Tensor，但我不认为没有循环就可以做到，所以我想将它投影到 1D，做工作并投影它回到二维

Answer 1

在 1D 中，您可以简单地执行 A[indexes]。

在 2D 中它仍然可以通过这种方式实现：

A = torch.arange(5, 10).repeat(3, 1)  # shape: (3, 5)
indexes = torch.stack([torch.randperm(5) for _ in range(3)])  # shape (3, 5)

A_sort = A[torch.arange(3).unsqueeze(1), indexes]
print(A_sort)

Answer 2

您可以使用 scatter:

A = torch.tensor([1, 2, 3, 4])
indices = torch.tensor([1, 0, 3, 2])
result = torch.tensor([0, 0, 0, 0])
print(result.scatter_(0, indices, A))

如何根据相同大小的索引张量重新排序张量

How to reorder tensor based on indexes tensor from the same size

pytorch

tensor