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 (Array | ndarray | bool | number | bool | int | float | complex | saiunit.Quantity | Sequence[Array | ndarray | bool | number | bool | int | float | complex] | 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:

Array | ndarray | bool | number | bool | int | float | complex | 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