ML Katas

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 perspective on dimension manipulation.

Guidance

This is a bit of a trick question! einops is primarily for reshaping, not for reordering elements within a dimension. The standard way is torch.flip. However, you can simulate it by combining rearrange with an indexing trick if you really wanted to. The most practical einops-related way might be to rearrange to a list and reverse it, but for a pure tensor solution, slicing is best. For this exercise, provide the slicing method and explain why einops is not the right tool for this specific job.

Starter Code

import torch

def reverse_sequence(x):
    # x: (B, N, D), reverse the N dimension
    # The best way is not with einops!
    return torch.flip(x, dims=)

Verification

Create a tensor like torch.arange(10).view(1, 5, 2). The output should have the inner arrays in reversed order. For example, [[[0, 1], [2, 3], ...]] becomes [[[..., [2, 3], [0, 1]]].