einops_rearrange_reduce_repeat_tensor_manipulation_quickstart.py

1import numpy as np
2from einops import rearrange, reduce, repeat
3
4# Create a dummy image-like batch: 10 images, 32x32 pixels, 3 color channels
5images = np.random.randn(10, 32, 32, 3)
6
7# 1. Rearrange: Transpose axes (batch, height, width, channels) -> (batch, channels, height, width)
8rearranged = rearrange(images, 'b h w c -> b c h w')
9
10# 2. Reduce: Compute the mean across the spatial dimensions (height and width)
11# This results in a tensor of shape (batch, channels)
12mean_pool = reduce(images, 'b h w c -> b c', 'mean')
13
14# 3. Repeat: Repeat a tensor along a new dimension
15# Here we take a single grayscale image and repeat it to create 3 channels
16gray_image = np.random.randn(32, 32)
17color_image = repeat(gray_image, 'h w -> h w c', c=3)
18
19# 4. Composition: Flatten a batch of images into a single sequence of pixels
20flat_pixels = rearrange(images, 'b h w c -> (b h w) c')
21
22print(f"Original shape: {images.shape}")
23print(f"Rearranged shape: {rearranged.shape}")
24print(f"Reduced shape: {mean_pool.shape}")
25print(f"Repeated shape: {color_image.shape}")
26print(f"Flattened shape: {flat_pixels.shape}")