code: def batch_shuffle(
      batch: torch.Tensor, distributed: bool = False
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """Randomly shuffles all tensors in the batch.
  
      Args:
          batch:
              The batch to shuffle.
          distributed:
              If True then batches are shuffled across multiple gpus.
  
      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.
  
      Examples:
          >>> # forward pass through the momentum model with batch shuffling
          >>> x1_shuffled, shuffle = batch_shuffle(x1)
          >>> f1 = moco_momentum(x1)
          >>> out0 = projection_head_momentum(f0)
          >>> out1 = batch_unshuffle(out1, shuffle)
      """
      if distributed:
          return batch_shuffle_distributed(batch)
      batch_size = batch.shape[0]
      shuffle = torch.randperm(batch_size, device=batch.device)
      return batch[shuffle], 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 1 for calling a global function with example usage: batch_shuffle(batch, distributed) and returns: batch,  shuffle, batch_shuffle_distributed