megengine.functional.roll

roll(inp, shift, axis=None)[源代码]

Roll the tensor along the given axis(or axes). Elements that are shifted beyond the last position are re-introduced at the first position.

参数
  • inp (Tensor) – input tensor.

  • shift (Union[int, Iterable[int]]) – the number of places by which the elements of the tensor are shifted. If shift is a tuple, axis must be a tuple of the same size, and each axis will be rolled by the corresponding shift value.

  • axis (Union[int, Iterable[int], None]) – axis along which to roll. If axis is not specified, the tensor will be flattened before rolling and then restored to the original shape. Duplicate axes is allowed if it is a tuple. Default: None.

Examples:

import numpy as np
from megengine import tensor
import megengine.functional as F

x = tensor([[1,2],[3,4],[5,6]], np.int32)
y = F.roll(x, 1, 0)
print(y.numpy())

Outputs:

[[5 6]
 [1 2]
 [3 4]]