class_docstr: Auxiliary data structure for averaging, responsible for splitting tensors into parts and reassembling them.
  The class is designed to avoid excessive memory allocation and run all heavy computation in background
  
  :param tensors: local tensors to be split and aggregated
  :param peer_fractions: for each peer, a target fraction of vector elements that this peer should average
  :param compression: optionally compress tensors with this compression algorithm before sending them to peers
  :param part_size_bytes: greedily split tensors into parts of up to this many bytes (after compression)
  :param tensor_infos: CompressionInfo for each respective tensor; this determines how the tensor will be compressed
  :param return_deltas: if True, output tensors are differences (aggregated tensor - local tensor)
  :param prefetch: when compressing, pre-compute this many compressed tensors in background
class_docstr_tok: ["Auxiliary","data","structure","for","averaging",",","responsible","for","splitting","tensors","into","parts","and","reassembling","them",".","The","class","is","designed","to","avoid","excessive","memory","allocation","and","run","all","heavy","computation","in","background",":","param","tensors",":","local","tensors","to","be","split","and","aggregated",":","param","peer_fractions",":","for","each","peer",",","a","target","fraction","of","vector","elements","that","this","peer","should","average",":","param","compression",":","optionally","compress","tensors","with","this","compression","algorithm","before","sending","them","to","peers",":","param","part_size_bytes",":","greedily","split","tensors","into","parts","of","up","to","this","many","bytes","(","after","compression",")",":","param","tensor_infos",":","CompressionInfo","for","each","respective","tensor",";","this","determines","how","the","tensor","will","be","compressed",":","param","return_deltas",":","if","True",",","output","tensors","are","differences","(","aggregated","tensor","-","local","tensor",")",":","param","prefetch",":","when","compressing",",","pre-compute","this","many","compressed","tensors","in","background"]
code: def __init__(
      self,
      tensors: Sequence[torch.Tensor],
      peer_fractions: Sequence[float],
      compression: CompressionBase = NoCompression(),
      part_size_bytes: int = DEFAULT_PART_SIZE_BYTES,
      tensor_infos: Optional[Sequence[CompressionInfo]] = None,
      return_deltas: bool = True,
      prefetch: int = 1,
  ):
      if tensor_infos is None:
          tensor_infos = tuple(
              CompressionInfo.from_tensor(x, key=i)
              for i, x in enumerate(tensors)
          )
      assert len(tensor_infos) == len(
          tensors
      ), "compression types do not match the number of tensors"
      self.local_tensors, self.peer_fractions, self.group_size = (
          tensors,
          peer_fractions,
          len(peer_fractions),
      )
      self.compression, self.part_size_bytes, self.tensor_infos = (
          compression,
          part_size_bytes,
          tensor_infos,
      )
      self.total_size = sum(tensor.numel() for tensor in tensors)
      self.failed_size = 0
      self.return_deltas = return_deltas
      self.prefetch = prefetch
  
      self._input_parts_by_peer = [
          deque() for _ in range(self.group_size)
      ]
      self._output_parts_by_peer = [
          deque() for _ in range(self.group_size)
      ]
      self._inputs_consumed_by_peer = [
          False for _ in range(self.group_size)
      ]
      self._output_part_available = [
          asyncio.Event() for _ in range(self.group_size)
      ]
      self._outputs_registered_by_peer = [
          0 for _ in range(self.group_size)
      ]
      self._outputs_consumed = False
      self.finished = asyncio.Event()
      self.num_parts_by_tensor = []
  
      # split tensor parts in proportion to target_size_by_peer
      current_length = 0
      current_peer_index = 0
      pivots = (
          np.cumsum(peer_fractions)
          / np.sum(peer_fractions)
          * self.total_size
      ).astype(np.int64)
      pivots[-1] = self.total_size
  
      for tensor, info in zip(self.local_tensors, self.tensor_infos):
          bytes_per_value = (
              tensor.element_size()
              * compression.estimate_compression_ratio(info)
          )
          part_size_values = int(part_size_bytes / bytes_per_value)
          tensor_parts = (
              tensor.detach().view(-1).split(part_size_values)
          )
          self.num_parts_by_tensor.append(len(tensor_parts))
          for part_index, part in enumerate(tensor_parts):
              part_info = info.get_part(part_index, part_size_values)
              if (
                  current_length + len(part)
                  > pivots[current_peer_index]
              ):
                  # switch to next peer; if a part lands between parts of two or
                  # more peers, assign that part to the peer with highest intersection
                  prev_peer_index = current_peer_index
                  peer_intersections = [
                      pivots[current_peer_index] - current_length
                  ]
                  while (
                      current_length + len(part)
                      > pivots[current_peer_index]
                  ):
                      current_peer_index += 1
                      current_peer_part_end = min(
                          current_length + len(part),
                          pivots[current_peer_index],
                      )
                      peer_intersections.append(
                          current_peer_part_end
                          - pivots[current_peer_index - 1]
                      )
                  assigned_peer_index = prev_peer_index + np.argmax(
                      peer_intersections
                  )
                  self._input_parts_by_peer[assigned_peer_index].append(
                      (part, part_info)
                  )
              else:
                  self._input_parts_by_peer[current_peer_index].append(
                      (part, part_info)
                  )
              current_length += len(part)
  
      assert current_length == self.total_size
      self.num_parts_by_peer = tuple(
          len(parts) for parts in self._input_parts_by_peer
      )
  
imports: import asyncio
  from collections import deque
  from typing import AsyncIterable, AsyncIterator, Optional, Sequence, Tuple, TypeVar
  import numpy
  import torch
  from hivemind.compression import CompressionBase, CompressionInfo, NoCompression
  from hivemind.proto import runtime_pb2
  from hivemind.utils import amap_in_executor, as_aiter, get_logger
label_desc: Use image node_id 1 to create a new TensorPartContainer object  with example: obj = TensorPartContainer(tensors, peer_fractions, compression, part_size_bytes, tensor_infos, return_deltas, prefetch)