code: def batch_shuffle_distributed(
      batch: torch.Tensor,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """Shuffles batch over multiple devices.
  
      This code was taken and adapted from here:
      https://github.com/facebookresearch/moco.
  
      Args:
          batch:
              The tensor to shuffle.
  
      Returns:
          A (batch, shuffle) tuple where batch is the shuffled version of the
          input batch and shuffle is an index to restore the original order.
  
      """
      # gather from all devices
      batch_size_this = batch.shape[0]
      batch_gather = concat_all_gather(batch)
      batch_size_all = batch_gather.shape[0]
  
      num_devices = batch_size_all // batch_size_this
  
      # random shuffle index
      idx_shuffle = torch.randperm(batch_size_all, device=batch.device)
  
      # broadcast to all devices
      dist.broadcast(idx_shuffle, src=0)
  
      # index for restoring
      shuffle = torch.argsort(idx_shuffle)
  
      # shuffled index for this device
      rank = dist.get_rank()
      idx_this = idx_shuffle.view(num_devices, -1)[rank]
      return batch_gather[idx_this], shuffle
  
imports: import math
  import warnings
  from typing import Iterable, List, Optional, Tuple, Union
  import numpy
  import torch
  import torch.distributed
  import torch.nn
  from numpy.typing import NDArray
  from torch.nn import Module
  from torch.nn.modules.batchnorm import _BatchNorm
  from torch.nn.parameter import Parameter
label_desc: Use image node_id 4 for calling a global function with example usage: batch_shuffle_distributed(batch) and returns: batch_gather,  shuffle