创建循环移位数组的数组/张量

Question

我想创建2d张量(或numpy数组，其实并不重要)，其中每一行都将循环移位到第一行。我使用for循环：

import torch
import numpy as np
a = np.random.rand(33, 11)
miss_size = 64
lp_order = a.shape[1] - 1
inv_a = -np.flip(a, axis=1)
mtx_size = miss_size+lp_order   # some constant
mtx_row = torch.cat((torch.from_numpy(inv_a), torch.zeros((a.shape[0], miss_size - 1 + a.shape[1]))), dim=1)
mtx_full = mtx_row.unsqueeze(1)
for i in range(mtx_size):
mtx_row = torch.roll(mtx_row, 1, 1)
mtx_full = torch.cat((mtx_full, mtx_row.unsqueeze(1)), dim=1)

需要解压缩，因为我将2d张量堆叠到3d张量中

有更有效的方法吗？也许是线性代数的把戏，或者更像蟒蛇的方法。

Answer 1

您可以使用scipy.linalg.circulant():

scipy.linalg.circulant([1, 2, 3])
# array([[1, 3, 2],
#        [2, 1, 3],
#        [3, 2, 1]])

Answer 2

我相信您可以使用torch.gather通过构造适当的索引张量来实现这一点。这种方法也适用于批处理。

如果我们采用这种方法，目标是构建一个索引张量，其中每个值都指向mtx_row中的一个索引(此处为dim=1中的最后一个维度)。在这种情况下，它的形状将是(3, 3):

tensor([[0, 1, 2],
[2, 0, 1],
[1, 2, 0]])

您可以通过广播带有自己转置的torch.arange并对生成的矩阵应用模来实现这一点：

>>> idx = (n-torch.arange(n)[None].T + torch.arange(n)[None]) % n
tensor([[0, 1, 2],
[2, 0, 1],
[1, 2, 0]])

设mtx_row为(2, 3):

>>> mtx_row
tensor([[0.3523, 0.0170, 0.1875],
[0.2156, 0.7773, 0.4563]])

从那里你需要idx和mtx_row，所以它们有相同的形状：

>>> idx_ = idx[None].expand(len(mtx_row), -1, -1)
>>> val_ = mtx_row[:, None].expand(-1, n, -1)

然后我们可以将torch.gather应用于最后一个维度dim=2:

>>> val_.gather(-1, idx_)
tensor([[[0.3523, 0.0170, 0.1875],
[0.1875, 0.3523, 0.0170],
[0.0170, 0.1875, 0.3523]],
[[0.2156, 0.7773, 0.4563],
[0.4563, 0.2156, 0.7773],
[0.7773, 0.4563, 0.2156]]])