如何在 PyTorch 中以不同偏移量移动张量中的列（或行）？

Question

在 PyTorch 中，内置torch.roll函数只能以相同的偏移量移动列（或行）。但我想用不同的偏移量来移动列。假设输入张量是

\n

[[1,2,3],\n [4,5,6],\n [7,8,9]]\n

\n

比方说，我想i对第 i 列进行偏移量移动。因此，预期输出是

\n

[[1,8,6],\n [4,2,9],\n [7,5,3]]\n

\n

这样做的一个选项是使用每个列单独移动torch.roll并连接它们。但出于有效性和代码紧凑性的考虑，我不想引入循环结构。有没有更好的办法\xef\xbc\x9f

\n

Answer 1

让我们定义一些名称：

import torch

mat = torch.Tensor(
[[1,2,3],
 [4,5,6],
 [7,8,9]])

indices = torch.LongTensor([0, 1, 2]) # Could also use arange in this specific scenario

首先，你可以创建一个像这样的张量

[[0, 0, 0],
 [1, 1, 1],
 [2, 2, 2]]

使用

[[0, 0, 0],
 [1, 1, 1],
 [2, 2, 2]]

现在，让我们创建一个目标索引的张量

[[0, 2, 1],
 [1, 0, 2],
 [2, 1, 0]]

和

arange1 = torch.arange(3).view((3, 1)).repeat((1, 3))

最后，我们得到预期的输出

[[0, 2, 1],
 [1, 0, 2],
 [2, 1, 0]]