ML Katas

### Description Vision Transformers (ViT) process images by first breaking them down into a sequence of flattened patches. The `einops` library is perfectly suited for this task, offering a...

### Description In the Multi-Head Attention mechanism, the input tensor `(B, N, D)` is linearly projected to create the Query, Key, and Value matrices. These are then reshaped to have separate...

### Description Space-to-depth is an operation that rearranges blocks of spatial data into the channel dimension. For an input of shape `(B, C, H, W)` and a block size `S`, the output will be `(B,...

### Description Reversing the order of elements in a sequence is a common operation. While it can be done with slicing (`torch.flip`), let's practice doing it with `einops` for a different...

### Description Depth-to-space is the inverse of the space-to-depth operation. It rearranges features from the channel dimension into spatial blocks, increasing spatial resolution and decreasing...

### Description Perform a batched matrix multiplication `(B, N, D) @ (B, D, M) -> (B, N, M)` using `einops` `einsum`. While `torch.bmm` is the standard, this is a good exercise to understand how...

### Description `torch.squeeze` and `torch.unsqueeze` are common for removing or adding dimensions of size one. `einops.rearrange` can do this as well, often with more clarity by explicitly naming...

### Description After the multi-head attention calculation, the output tensor typically has the shape `(B, num_heads, N, head_dim)`. To feed this into the next layer (usually a feed-forward...

### 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. ###...