fn

rearrange

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

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.
**axes_lengthsint= {}
Named axis sizes used to disambiguate decomposed groups.

Returns

Tensor

The rearranged tensor.

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)