All Katas - ML Katas

Batched Expert Forward Pass with Einops

this year medium (<1 hr) | pytorch einops performance

A naive implementation of an MoE layer might involve a loop over the experts. This is inefficient. A much better approach is to perform a single, batched matrix multiplication for all expert...

Sliding Window Attention Preparation

this year medium (<10 mins) | pytorch transformer attention einops

### Description Full self-attention has a quadratic complexity with respect to sequence length, which is prohibitive for very long sequences. Models like Longformer introduce **sliding window...

Multi-Head Attention: Splitting Heads

this year medium (<10 mins) | transformer attention einops

### Description In multi-head attention, the query, key, and value tensors are split into multiple heads. Given a tensor of shape `(B, N, D)`, where `D` is the embedding dimension, you need to...

Multi-Head Attention: Merging Heads

this year medium (<10 mins) | transformer attention einops

### Description The inverse of splitting heads. After computing attention for each head, you need to merge them back. Given a tensor of shape `(B, H, N, D//H)`, you need to merge it back to `(B,...

Pixel Unshuffle (Pixel to Channel)

this year medium (<10 mins) | einops super-resolution pixel shuffle

### Description This is another name for the space-to-depth operation, common in super-resolution models. It involves rearranging blocks of spatial data into the channel dimension. Given a tensor...

Pixel Shuffle (Channel to Pixel)

this year medium (<10 mins) | einops super-resolution pixel shuffle

### Description The inverse of pixel unshuffle, also known as depth-to-space. It is used to upscale an image by rearranging elements from the channel dimension into spatial blocks. Given a tensor...