ML Katas

Description

Grouped convolution divides the input channels into groups and performs a separate convolution on each. You can simulate the tensor rearrangement part of this operation using einops.

Guidance

Start with a tensor of shape (B, C, H, W). Your task is to rearrange it to (B, G, C//G, H, W), where G is the number of groups. This isolates the channel groups so a subsequent operation could be applied to each group independently.

Starter Code

import torch
from einops import rearrange

def setup_for_grouped_conv(x, groups):
    # x: (B, C, H, W)
    # Rearrange to (B, G, C_per_group, H, W)
    return rearrange(x, 'b (g c_g) h w -> b g c_g h w', g=groups)

Verification

For an input of shape (10, 32, 64, 64) and groups=4, the output shape should be (10, 4, 8, 64, 64).