einops_rearrange_reduce_repeat_array_manipulation_quickstart.py

1import numpy as np
2from einops import rearrange, reduce, repeat
3
4# Create a dummy image-like batch: 1 image, 3 channels, 4x4 pixels
5# Format: (batch, channels, height, width)
6ims = np.random.randn(1, 3, 4, 4)
7
8# 1. Rearrange: Swap axes or flatten
9# Convert from (B, C, H, W) to (B, H, W, C)
10output_rearrange = rearrange(ims, 'b c h w -> b h w c')
11
12# 2. Reduce: Apply operations like mean, min, or max across dimensions
13# Here we take the mean over the color channels to get a grayscale image
14output_reduce = reduce(ims, 'b c h w -> b h w', 'mean')
15
16# 3. Repeat: Duplicate data along new or existing dimensions
17# Repeat the batch 3 times
18output_repeat = repeat(ims, 'b c h w -> (repeat b) c h w', repeat=3)
19
20# 4. Composition/Decomposition: Splitting and combining axes
21# Flatten the spatial dimensions into a single vector
22output_flat = rearrange(ims, 'b c h w -> b c (h w)')
23
24print(f"Original shape: {ims.shape}")
25print(f"Rearranged shape: {output_rearrange.shape}")
26print(f"Reduced shape: {output_reduce.shape}")
27print(f"Repeated shape: {output_repeat.shape}")
28print(f"Flattened shape: {output_flat.shape}")