roll

dragon.vm.torch.roll(
  input,
  shifts,
  dims=None
)[source]

Roll elements along the given dimension.

dims could be negative or None:

x = torch.tensor([[1, 2, 3], [4, 5, 6]])

# A negative dimension is the last-k dimension
print(torch.roll(x, shifts=1, dims=1))  # [[3, 1, 2], [6, 4, 5]]
print(torch.roll(x, shifts=1, dims=-1))  # Equivalent

# If dimension is None, roll input as a vector
print(torch.roll(x, shifts=1))  # [[6, 1, 2], [3, 4, 5]]

# Also, dimension could be a sequence of integers
print(torch.roll(x, shifts=(1, 1), dims=(0, 1)))  # [[6, 4, 5], [3, 1, 2]]
print(torch.roll(x, shifts=(1, -1), dims=(0, 1)))  # [[5, 6, 4], [2, 3, 1]]
Parameters:
  • input (dragon.vm.torch.Tensor) – The input tensor.
  • shifts (Union[int, Sequence[int]]) – The rolling offset of each dimension.
  • dims (Union[int, Sequence[int]], optional) – The dimension to roll.
Returns:

dragon.vm.torch.Tensor – The output tensor.