braincore.math.ein_rearrange

braincore.math.ein_rearrange(tensor, pattern, **axes_lengths)

ein_rearrange is a 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.

Examples for rearrange operation:

```python # suppose we have a set of 32 images in “h w c” format (height-width-channel) >>> images = [np.random.randn(30, 40, 3) for _ in range(32)]

# stack along first (batch) axis, output is a single array >>> ein_rearrange(images, ‘b h w c -> b h w c’).shape (32, 30, 40, 3)

# concatenate images along height (vertical axis), 960 = 32 * 30 >>> ein_rearrange(images, ‘b h w c -> (b h) w c’).shape (960, 40, 3)

# concatenated images along horizontal axis, 1280 = 32 * 40 >>> ein_rearrange(images, ‘b h w c -> h (b w) c’).shape (30, 1280, 3)

# reordered axes to “b c h w” format for deep learning >>> ein_rearrange(images, ‘b h w c -> b c h w’).shape (32, 3, 30, 40)

# flattened each image into a vector, 3600 = 30 * 40 * 3 >>> ein_rearrange(images, ‘b h w c -> b (c h w)’).shape (32, 3600)

# split each image into 4 smaller (top-left, top-right, bottom-left, bottom-right), 128 = 32 * 2 * 2 >>> ein_rearrange(images, ‘b (h1 h) (w1 w) c -> (b h1 w1) h w c’, h1=2, w1=2).shape (128, 15, 20, 3)

# space-to-depth operation >>> ein_rearrange(images, ‘b (h h1) (w w1) c -> b h w (c h1 w1)’, h1=2, w1=2).shape (32, 15, 20, 12)

```

When composing axes, C-order enumeration used (consecutive elements have different last axis) Find more examples in einops tutorial.

Parameters:

• tensor (Union[Array, ndarray, bool_, number, bool, int, float, complex, List[Union[Array, ndarray, bool_, number, bool, int, float, complex]]]) – tensor of any supported library (e.g. numpy.ndarray, tensorflow, pytorch). list of tensors is also accepted, those should be of the same type and shape

• pattern (str) – string, rearrangement pattern

• axes_lengths – any additional specifications for dimensions

Return type:

Union[Array, ndarray, bool_, number, bool, int, float, complex]

Returns:

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