r"""Roll the tensor along the given axis(or axes). Elements that are shifted beyond the last position are re-introduced at the first position. Args: inp: input tensor. shift: the number of places by which the elements of the tensor are shifted. If shift is a tupl
(
inp: Tensor,
shift: Union[int, Iterable[int]],
axis: Optional[Union[int, Iterable[int]]] = None,
)
| 1231 | |
| 1232 | |
| 1233 | def roll( |
| 1234 | inp: Tensor, |
| 1235 | shift: Union[int, Iterable[int]], |
| 1236 | axis: Optional[Union[int, Iterable[int]]] = None, |
| 1237 | ): |
| 1238 | r"""Roll the tensor along the given axis(or axes). Elements that are shifted |
| 1239 | beyond the last position are re-introduced at the first position. |
| 1240 | |
| 1241 | Args: |
| 1242 | inp: input tensor. |
| 1243 | shift: the number of places by which the elements of the tensor are |
| 1244 | shifted. If shift is a tuple, axis must be a tuple of the same size, |
| 1245 | and each axis will be rolled by the corresponding shift value. |
| 1246 | axis: axis along which to roll. If axis is not specified, the tensor |
| 1247 | will be flattened before rolling and then restored to the original shape. |
| 1248 | Duplicate axes is allowed if it is a tuple. Default: None. |
| 1249 | |
| 1250 | Examples: |
| 1251 | >>> import numpy as np |
| 1252 | >>> x = Tensor([[1,2],[3,4],[5,6]], np.int32) |
| 1253 | >>> F.roll(x, 1, 0) |
| 1254 | Tensor([[5 6] |
| 1255 | [1 2] |
| 1256 | [3 4]], dtype=int32, device=xpux:0) |
| 1257 | """ |
| 1258 | shp_bak = None |
| 1259 | if axis is None: |
| 1260 | shp_bak = inp.shape |
| 1261 | inp = inp.flatten() |
| 1262 | axis = 0 |
| 1263 | shp = inp.shape |
| 1264 | dim = len(shp) |
| 1265 | if isinstance(shift, int): |
| 1266 | assert isinstance(axis, int) |
| 1267 | shift, axis = [shift,], [axis,] |
| 1268 | assert len(shift) == len(axis) |
| 1269 | out = inp |
| 1270 | for i in range(len(shift)): |
| 1271 | axis_ = axis[i] |
| 1272 | shift_ = shift[i] |
| 1273 | axis_normalized_ = axis_ + dim if axis_ < 0 else axis_ |
| 1274 | assert ( |
| 1275 | dim > axis_normalized_ >= 0 |
| 1276 | ), "axis out of range (expected to be in range of [{}, {}], but got {})".format( |
| 1277 | -dim, dim - 1, axis_ |
| 1278 | ) |
| 1279 | if shift_ == 0: |
| 1280 | continue |
| 1281 | size = shp[axis_normalized_] |
| 1282 | shift_normalized_ = 0 if size == 0 else shift_ % size |
| 1283 | if shift_normalized_ > 0: |
| 1284 | a, b = split(out, [size - shift_normalized_,], axis=axis_normalized_) |
| 1285 | else: |
| 1286 | a, b = split(out, [-shift_normalized_,], axis=axis_normalized_) |
| 1287 | out = concat((b, a), axis=axis_normalized_) |
| 1288 | if shp_bak is not None: |
| 1289 | out = out.reshape(shp_bak) |
| 1290 | return out |