mlx.core.roll

Contents

mlx.core.roll#

roll(a: array, shift: int | Tuple[int], axis: None | int | Tuple[int] = None, /, *, stream: None | Stream | Device = None) array#

Roll array elements along a given axis.

Elements that are rolled beyond the end of the array are introduced at the beggining and vice-versa.

If the axis is not provided the array is flattened, rolled and then the shape is restored.

Parameters:
  • a (array) – Input array

  • shift (int or tuple(int)) – The number of places by which elements are shifted. If positive the array is rolled to the right, if negative it is rolled to the left. If an int is provided but the axis is a tuple then the same value is used for all axes.

  • axis (int or tuple(int), optional) – The axis or axes along which to roll the elements.