def METHOD_NAME(self, scores, k=1):
    token_to_workers = torch.topk(scores, dim=1, k=k, largest=True).indices.view(-1)
    token_to_workers, sort_ordering = torch.sort(token_to_workers)
    worker2token = sort_ordering // k
    # Find how many tokens we're sending to each other worker (being careful for sending 0 tokens to some workers)
    output_splits = torch.zeros(
        (self.num_workers,), dtype=torch.long, device=scores.device
    )
    workers, counts = torch.unique_consecutive(token_to_workers, return_counts=True)
    output_splits[workers] = counts
    # Tell other workers how many tokens to expect from us
    input_splits = All2All.apply(output_splits)
    return worker2token, input_splits.tolist(), output_splits.tolist()