einrearrange

Contents

einrearrange#

class saiunit.math.einrearrange(x, pattern, **axes_lengths)#

Reader-friendly smart element reordering for multidimensional tensors.

This operation includes functionality of transpose (axes permutation), reshape (view), squeeze, unsqueeze, stack, concatenate and other operations. When composing axes, C-order enumeration is used (consecutive elements have different last axis).

Parameters:

  • x (Union[Array, ndarray, number, bool, saiunit.Quantity, Sequence[Array | ndarray | number | bool], Sequence[saiunit.Quantity]]) – Input tensor(s). A list of tensors of the same type and shape is also accepted for stacking.

  • pattern (str) – Rearrangement pattern in 'input -> output' form.

  • **axes_lengths (int) – Additional specifications for dimension sizes.

Returns:

out – Tensor of the same type as input. If possible, a view to the original tensor is returned.

Return type:

Union[Array, ndarray, number, bool, saiunit.Quantity]

Examples

>>> import jax.numpy as jnp
>>> import saiunit.math as sumath
>>> x = jnp.zeros((2, 3, 4))
>>> sumath.einrearrange(x, 'b h w -> b (h w)').shape
(2, 12)
>>> sumath.einrearrange(x, 'b h w -> b w h').shape
(2, 4, 3)
Contents