import torch


def gpu_non_dominated_sort(swarm: torch.Tensor):
    """
    The GPU version of non-dominated sorting algorithms
    Args:
        swarm (np.ndarray): m x n scoring matrix, where m is the number of samples
            and n is the number of objectives.

    Returns:
        fronts (List): a list of Pareto fronts, in which the dominated solutions are on the top,
            and non-dominated solutions are on the bottom.
    """
    domina = (swarm.unsqueeze(1) <= swarm.unsqueeze(0)).all(-1)
    domina_any = (swarm.unsqueeze(1) < swarm.unsqueeze(0)).any(-1)
    domina = (domina & domina_any).half()

    fronts = []
    while (domina.diag() == 0).any():
        count = domina.sum(dim=0)
        front = torch.where(count == 0)[0]
        fronts.append(front)
        domina[front, :] = 0
        domina[front, front] = -1
    return fronts