class_docstr: Implementation of the distributed helper.
  
  Methods of this class can be used as an high-level wrappers
  aroung the torch.ddistributed API to allow for a simpler
  distributed communication.
  
  Only a single object of this class is instantiated
  as the "DistributedHelper" singleton.
  
  
  Note: differently from the original Pytorch API, which requires
  that input tensor(s) to be moved to the default device (forced to
  CUDA if using NCCL), these functions usually also manage input tensors
  residing on a different devices. The returned elements will
  be moved to the same device of the input tensor. Consider looking at
  the documentation of each method for more details.
class_docstr_tok: ["Implementation","of","the","distributed","helper",".","Methods","of","this","class","can","be","used","as","an","high-level","wrappers","aroung","the","torch.ddistributed","API","to","allow","for","a","simpler","distributed","communication",".","Only","a","single","object","of","this","class","is","instantiated","as","the","``","DistributedHelper","''","singleton",".","Note",":","differently","from","the","original","Pytorch","API",",","which","requires","that","input","tensor","(","s",")","to","be","moved","to","the","default","device","(","forced","to","CUDA","if","using","NCCL",")",",","these","functions","usually","also","manage","input","tensors","residing","on","a","different","devices",".","The","returned","elements","will","be","moved","to","the","same","device","of","the","input","tensor",".","Consider","looking","at","the","documentation","of","each","method","for","more","details","."]
docstr: Gather all for tensors only.
  
  Differences with torch.distributed:
      - The input tensor can reside in any device.
      - The input tensor will be transmitted using the current backend.
          However, the resulting tensors will be moved to the save device
          as the `tensor` parameter before returning them,
          no matter the backend in use.
      - No-op if not running a distributed training.
  
  This will also manage tensors of different shapes. If you
  are sure that the tensors will be of the same shape, consider
  passing same_shape to speed up the communication.
  
  Beware that, if you are in need of concatenating multiple tensors,
  method :meth:`cat_all` may be more suitable.
  
  :param tensor: The tensor to be sent from the current process.
  :return: A list of tensors, one from each process (in rank order).
docstr_tok: ["Gather","all","for","tensors","only",".","Differences","with","torch.distributed",":","-","The","input","tensor","can","reside","in","any","device",".","-","The","input","tensor","will","be","transmitted","using","the","current","backend",".","However",",","the","resulting","tensors","will","be","moved","to","the","save","device","as","the","`","tensor","`","parameter","before","returning","them",",","no","matter","the","backend","in","use",".","-","No-op","if","not","running","a","distributed","training",".","This","will","also","manage","tensors","of","different","shapes",".","If","you","are","sure","that","the","tensors","will","be","of","the","same","shape",",","consider","passing","same_shape","to","speed","up","the","communication",".","Beware","that",",","if","you","are","in","need","of","concatenating","multiple","tensors",",","method",":","meth",":","`","cat_all","`","may","be","more","suitable",".",":","param","tensor",":","The","tensor","to","be","sent","from","the","current","process",".",":","return",":","A","list","of","tensors",",","one","from","each","process","(","in","rank","order",")","."]
code: def gather_all(
      self,
      tensor: Tensor,
      same_shape: bool = False,
      shapes: Optional[List[List[int]]] = None,
  ) -> List[Tensor]:
      """
      Gather all for tensors only.
  
      Differences with torch.distributed:
          - The input tensor can reside in any device.
          - The input tensor will be transmitted using the current backend.
              However, the resulting tensors will be moved to the save device
              as the `tensor` parameter before returning them,
              no matter the backend in use.
          - No-op if not running a distributed training.
  
      This will also manage tensors of different shapes. If you
      are sure that the tensors will be of the same shape, consider
      passing same_shape to speed up the communication.
  
      Beware that, if you are in need of concatenating multiple tensors,
      method :meth:`cat_all` may be more suitable.
  
      :param tensor: The tensor to be sent from the current process.
      :return: A list of tensors, one from each process (in rank order).
      """
      if not self.is_distributed:
          return [tensor]
  
      # Based on:
      # https://discuss.pytorch.org/t/how-to-concatenate-different-size-tensors-from-distributed-processes/44819/4
  
      if same_shape:
          # Same size for all tensors
          if len(tensor.shape) > 0:
              tensor_size = list(tensor.shape)
          else:
              tensor_size = [0]
          all_tensors_shape = [
              tensor_size for _ in range(self.world_size)
          ]
      elif shapes is not None:
          # Shapes given by the user
          # make sure it is a list of lists
          all_tensors_shape = [list(s) for s in shapes]
      else:
          # Tensor differ by whole shape
          all_tensors_shape = self.gather_tensor_shapes(tensor)
  
      same_shape = all(
          all_tensors_shape[0] == x for x in all_tensors_shape
      )
      orig_device = tensor.device
  
      if same_shape:
          # Same shape: create identical tensors and proceed with all_gather
          out_tensors = [
              torch.empty_like(tensor) for _ in all_tensors_shape
          ]
      else:
          # Different shapes: create a tensors of the size of the bigger one
          all_tensors_numel = []
          dtype = tensor.dtype
          for t_shape in all_tensors_shape:
              if t_shape[0] == 0 and len(t_shape) == 1:
                  # Tensor with 0-length shape
                  curr_size = 1
              else:
                  curr_size = 1
                  for t_s in t_shape:
                      curr_size *= t_s
              all_tensors_numel.append(curr_size)
  
          max_numel = max(all_tensors_numel)
          out_tensors = [
              torch.empty((max_numel,), dtype=dtype)
              for _ in all_tensors_shape
          ]
  
          tensor = tensor.flatten()
          n_padding = max_numel - tensor.numel()
          if n_padding > 0:
              padding = torch.zeros(
                  (n_padding,), dtype=tensor.dtype, device=orig_device
              )
              tensor = torch.cat((tensor, padding), dim=0)
  
      tensor, _ = self._prepare_for_distributed_comm(tensor)
      out_tensors = [
          self._prepare_for_distributed_comm(t)[0] for t in out_tensors
      ]
  
      torch.distributed.all_gather(out_tensors, tensor)
  
      if not same_shape:
          # The tensors are flat and of the wrong dimension: re-shape them
          for tensor_idx, (tensor_sz, tensor_numel, out_t) in enumerate(
              zip(all_tensors_shape, all_tensors_numel, out_tensors)
          ):
              if tensor_sz[0] == 0:
                  # Tensor with 0-length shape
                  out_tensors[tensor_idx] = out_t[
                      :tensor_numel
                  ].reshape(tuple())
              else:
                  out_tensors[tensor_idx] = out_t[
                      :tensor_numel
                  ].reshape(tensor_sz)
  
      out_tensors = [t.to(orig_device) for t in out_tensors]
      return out_tensors
  
imports: import os
  import pickle
  import warnings
  from io import BytesIO
  from typing import ContextManager, Optional, List, Any, Iterable, Dict, TypeVar
  import torch
  from torch import Tensor
  from torch.nn.modules import Module
  from torch.nn.parallel import DistributedDataParallel
  from typing_extensions import Literal
  from torch.distributed import init_process_group, broadcast_object_list
  from avalanche.distributed.distributed_consistency_verification import hash_tensor
label_desc: Use image node_id 15 for calling the _DistributedHelperCls obj's underlying member method code with example usage: obj.gather_all(tensor, same_shape, shapes) and returns: out_tensors, list