torch.roll

torch.roll(input, shifts, dims=None) → Tensor

沿给定维度滚动张量 input。超出最后一个位置的元素将重新出现在第一个位置。如果 dims 为 None，则在滚动之前会先将张量展平，并在之后恢复到原始形状。

参数

• input (Tensor) – 需要输入的张量。

• shifts (int 或 元组 of ints) – 张量元素要移动的位置数量。如果 shifts 是一个元组，则 dims 也必须是一个相同大小的元组，并且每个维度将根据对应值进行滚动。

• dims (int 或 元组 of ints) – 沿着该轴进行滚动操作

示例：

>>> x = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]).view(4, 2)
>>> x
tensor([[1, 2],
        [3, 4],
        [5, 6],
        [7, 8]])
>>> torch.roll(x, 1)
tensor([[8, 1],
        [2, 3],
        [4, 5],
        [6, 7]])
>>> torch.roll(x, 1, 0)
tensor([[7, 8],
        [1, 2],
        [3, 4],
        [5, 6]])
>>> torch.roll(x, -1, 0)
tensor([[3, 4],
        [5, 6],
        [7, 8],
        [1, 2]])
>>> torch.roll(x, shifts=(2, 1), dims=(0, 1))
tensor([[6, 5],
        [8, 7],
        [2, 1],
        [4, 3]])