def sinkhorn(
    out: Tensor,
    iterations: int = 3,
    epsilon: float = 0.05,
    gather_distributed: bool = False,
) -> Tensor:
    """Distributed sinkhorn algorithm.

    As outlined in [0] and implemented in [1].

    - [0]: SwaV, 2020, https://arxiv.org/abs/2006.09882
    - [1]: https://github.com/facebookresearch/swav/

    Args:
        out:
            Similarity of the features and the SwaV prototypes.
        iterations:
            Number of sinkhorn iterations.
        epsilon:
            Temperature parameter.
        gather_distributed:
            If True then features from all gpus are gathered to calculate the
            soft codes Q.

    Returns:
        Soft codes Q assigning each feature to a prototype.
    """
    world_size = 1
    if gather_distributed and dist.is_initialized():
        world_size = dist.get_world_size()

    # Get the exponential matrix and make it sum to 1
    Q = torch.exp(out / epsilon).t()
    sum_Q = torch.sum(Q)
    if world_size > 1:
        dist.all_reduce(sum_Q)
    Q /= sum_Q

    B = Q.shape[1] * world_size

    for _ in range(iterations):
        # Normalize rows
        sum_of_rows = torch.sum(Q, dim=1, keepdim=True)
        if world_size > 1:
            dist.all_reduce(sum_of_rows)
        Q /= sum_of_rows
        # Normalize columns
        Q /= torch.sum(Q, dim=0, keepdim=True)
        Q /= B

    Q *= B
    return Q.t()
Distributed sinkhorn algorithm.

    As outlined in [0] and implemented in [1].

    - [0]: SwaV, 2020, https://arxiv.org/abs/2006.09882
    - [1]: https://github.com/facebookresearch/swav/

    Args:
        out:
            Similarity of the features and the SwaV prototypes.
        iterations:
            Number of sinkhorn iterations.
        epsilon:
            Temperature parameter.
        gather_distributed:
            If True then features from all gpus are gathered to calculate the
            soft codes Q.

    Returns:
        Soft codes Q assigning each feature to a prototype.
    
lightly/loss/swav_loss.py