code: def full(self, device=None):
      if device is None:
          device = self.orig_device
  
      # Allocate the full tensor as a flat buffer.
      full_numel = prod(self.full_size())
      flat_tensor = torch.zeros(
          [full_numel], dtype=self.local_data.dtype, device=device
      )
  
      # Prepare all-gather buffer
      partition_tensors = []
      for part_id in range(self.num_parts):
          part_size = (
              self.partition[part_id + 1] - self.partition[part_id]
          )
          buf = flat_tensor.narrow(
              0, start=self.partition[part_id], length=part_size
          )
          if part_id == self.rank:
              buf.copy_(self.local_data)
          partition_tensors.append(buf)
  
      # Collect the full tensor
      dist.all_gather(
          partition_tensors,
          partition_tensors[self.rank],
          group=self.group,
      )
  
      for i in range(len(partition_tensors)):
          partition_tensors[i].data = torch.zeros(1)
          partition_tensors[i] = None
  
      return flat_tensor.view(self.full_size()).clone().detach()
  
imports: from collections.abc import Iterable
  from deepspeed.moe.utils import is_moe_param
  import os
  import psutil
  import gc
  from math import sqrt
  from math import floor
  from bisect import bisect_left
  import torch
  from torch import inf
  from deepspeed import comm
  from deepspeed.utils import groups, logger
  from deepspeed.runtime.constants import PIPE_REPLICATED
  from numpy import prod
label_desc: Use image node_id 4 for calling the PartitionedTensor obj's underlying member method code with example usage: obj.full(device) and returns: flat_tensor