ML Katas

Einops: Repeat for Tiling/Broadcasting

easy (<10 mins) pytorch einops tensor
today by E

Description

The einops.repeat function is a powerful and readable alternative to torch.expand or torch.tile for broadcasting or repeating a tensor along new or existing dimensions.

Guidance

Use einops.repeat to tile a tensor. For example, take a tensor (H, W) and repeat it B times to create a batch of identical images, resulting in (B, H, W). Or, take a vector (D) and repeat it to match a sequence length N to get (N, D).

Starter Code

import torch
from einops import repeat

def repeat_vector(v, seq_len):
    # v: (D,)
    # Target: (seq_len, D)
    return repeat(v, 'd -> n d', n=seq_len)

def tile_image(img, batch_size):
    # img: (H, W)
    # Target: (batch_size, H, W)
    return repeat(img, 'h w -> b h w', b=batch_size)

Verification

Create a vector of shape (100,). Use repeat_vector to create a tensor of shape (50, 100). Verify that every row of the output is identical to the original vector.