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

Question

假设我有tensorA，并索引Tensor:A = [1, 2, 3, 4]，indexes = [1, 0, 3, 2]

我想从这两个创建一个新的Tensor，结果如下：[2, 1, 4, 3]结果的每个元素是来自A的元素，并且顺序由索引Tensor定义。有没有一种方法可以在没有循环的情况下使用PyTorchtensor操作？

我的目标是为2DTensor做这件事，但我认为没有循环的方法，所以我想把它投影到1D，做这项工作，然后把它投影回2D。

Answer 1

您可以使用散射：

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))

Answer 2

在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)