rearrange

→Tensor

rearrange(tensor: Tensor, pattern: str, axes_lengths: int = {})

source edit

Rearrange axes of tensor according to an einops pattern.

Performs arbitrary axis manipulation — transpose, reshape, composition, and decomposition — expressed in a single readable pattern string. The notation makes the intent explicit and removes the chain of transpose/view calls that obscure dimensionality transforms in raw tensor code.

Parameters

tensorTensor

Input tensor to rearrange.

patternstr

Einops pattern of the form "<lhs> -> <rhs>". Parentheses group axes for composition or decomposition, ... stands for the axes the pattern does not name, and 1 or () is a unit axis. Every named axis must appear on both sides.

**axes_lengthsint= {}

Axis lengths, needed to split a dimension into a group and checked against the input wherever the input already fixes the axis.

Returns

Tensor

The rearranged tensor.

Raises

lucid._C.engine.InvalidArgument

The pattern is malformed, an axis appears on one side only, or a keyword names no axis of the pattern.

lucid._C.engine.ShapeMismatch

The input's rank or a dimension does not fit the pattern or the keyword lengths.

Notes

The transformation is purely a reshape and permutation — no data is copied unless required by memory layout. Typical patterns include "b c h w -> b (h w) c" for flattening spatial dims into a token sequence, and "b (h w) c -> b c h w" for the inverse.

Examples

>>> import lucid
>>> x = lucid.randn(2, 3, 4)
>>> lucid.einops.rearrange(x, "b h w -> b (h w)").shape
(2, 12)
>>> lucid.einops.rearrange(x, "b ... -> b (...)").shape
(2, 12)
>>> lucid.einops.rearrange(x, "b (h1 h2) w -> b h1 h2 w 1", h2=3).shape
(2, 1, 3, 4, 1)