code: def sinkhorn(
      out: torch.Tensor,
      iterations: int = 3,
      epsilon: float = 0.05,
      gather_distributed: bool = False,
  ) -> torch.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()
  
file_path: lightly/lightly/loss/swav_loss.py
imports: from typing import List
  import torch
  import torch.distributed
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 1 for calling a global function with example usage: sinkhorn(out, iterations, epsilon, gather_distributed) and returns: Q