einx.roll

Contents

einx.roll#

einx.roll(description, tensor, *, shift, backend=None, **parameters)[source]#

Rolls the elements in the given tensor by the specified shift amounts.

The elementary operation has the signature [...] -> [...] and rolls elements separately along all axes. Elements that are rolled beyond the last position are re-introduced at the first position.

If there is no output expression, it is chosen to be the same as the input expression. For example, the following operations compute the same output:

y = einx.roll("a [b]", x, shift=4)
y = einx.roll("a [b] -> a [b]", x, shift=4)
Parameters:

  • description (str) – Description string for the operation in einx notation.

  • tensor (Tensor) – Input tensor or tensor factory matching the description string.

  • backend (Union[Backend, str, None]) – Backend to use for all operations. If None, uses the default backend for the given setting. Defaults to None.

  • graph – Whether to return the compiled code representation of this operation instead of computing the result. Defaults to False.

  • shift (Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], bool, int, float, complex, str, bytes, _NestedSequence[bool | int | float | complex | str | bytes]]) – Amounts by which elements are shifted along each axis. Can be a single integer or a list of integers matching the number of axes in the tensor.

  • **parameters (Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], bool, int, float, complex, str, bytes, _NestedSequence[bool | int | float | complex | str | bytes]]) – Additional parameters that specify dimension sizes, e.g. a=4.

Return type:

Tensor

Returns:

The result of the operation if graph=False, otherwise the compiled code representation of the operation.

Contents