def _rank_data(data: Tensor) -> Tensor:
    n = data.numel()
    rank = torch.empty_like(data, dtype=torch.int32)
    idx = data.argsort()
    rank[idx[:n]] = torch.arange(1, n + 1, dtype=torch.int32, device=data.device)

    uniq, inv, counts = torch.unique(
        data, sorted=True, return_inverse=True, return_counts=True
    )
    sum_ranks = torch.zeros_like(uniq, dtype=torch.int32)
    sum_ranks.scatter_add_(0, inv, rank.to(torch.int32))
    mean_ranks = sum_ranks / counts
    return mean_ranks[inv]