import torch

def compress(context):
    norms = torch.norm(context, dim=-1)
    target_len = int(context.shape[0] * 0.5)
    _, indices = torch.topk(norms, target_len)
    indices = torch.sort(indices).values
    return context[indices]

device = "cuda" if torch.cuda.is_available() else "cpu"
x = torch.randn(2048, 768, device=device)
y = compress(x)
print(f"Done: {x.shape[0]} -> {y.shape[0]}")
