einops_rearrange_op

→TensorImplPtr

int einops_rearrange_op(const int & a, const int & pattern, const int & axes_lengths)

source edit

Rearrange (permute and/or reshape) a tensor according to an einops pattern.

The pattern must have the form "lhs -> rhs". Both sides must contain exactly the same set of named axes; rhs may group or split them with parentheses. axes_lengths provides sizes for any axis that cannot be inferred from the input shape (typically the components of a decomposed group). No reduction is performed — element count is preserved.

Math

Pure index permutation and reshape:

y_{\sigma(i_1, \dots, i_n)} = x_{i_1, \dots, i_n}

where $\sigma$ is determined by aligning lhs and rhs flat axis lists.

See Also

einops_reduce_op, einops_repeat_op, einsum_op

Parameters

aTensorImplPtr

Input tensor.

patternstd::string

Einops rearrange pattern "lhs -> rhs".

axes_lengthsstd::map<std::string, std::int64_t>

Named sizes for axes whose size cannot be inferred from a.shape.

Returns

TensorImplPtr

Rearranged tensor. Same element count and dtype as a; shape is determined by rhs and axes_lengths.

Raises

EinopsPatternError

If the pattern is malformed, axes are inconsistent across sides, or a required size is missing from axes_lengths.

Examples

Pattern "b c h w -> b (h w) c" flattens spatial dims into a token
sequence (NCHW → NLC) suitable for transformer-style backbones.