code: def alltoallv_cpu(
      rank, world_size, input_tensor_list, retain_nones=True
  ):
      """
      Wrapper function to providing the alltoallv functionality by using underlying alltoall
      messaging primitive. This function, in its current implementation, supports exchanging
      messages of arbitrary dimensions and is not tied to the user of this function.
  
      This function pads all input tensors, except one, so that all the messages are of the same
      size. Once the messages are padded, It first sends a vector whose first two elements are
      1) actual message size along first dimension, and 2) Message size along first dimension
      which is used for communication. The rest of the dimensions are assumed to be same across
      all the input tensors. After receiving the message sizes, the receiving end will create buffers
      of appropriate sizes. And then slices the received messages to remove the added padding, if any,
      and returns to the caller.
  
      Parameters:
      -----------
      rank : int
          The rank of current worker
      world_size : int
          The size of the entire
      input_tensor_list : List of tensor
          The tensors to exchange
      retain_nones : bool
          Indicates whether to retain ``None`` data in returned value.
  
      Returns:
      --------
      list :
          list of tensors received from other processes during alltoall message
  
      """
      # ensure len of input_tensor_list is same as the world_size.
      assert input_tensor_list != None
      assert len(input_tensor_list) == world_size
  
      # ensure that all the tensors in the input_tensor_list are of same size.
      sizes = [list(x.size()) for x in input_tensor_list]
      for idx in range(1, len(sizes)):
          assert len(sizes[idx - 1]) == len(
              sizes[idx]
          )  # no. of dimensions should be same
          assert (
              input_tensor_list[idx - 1].dtype
              == input_tensor_list[idx].dtype
          )  # dtype should be same
          assert (
              sizes[idx - 1][1:] == sizes[idx][1:]
          )  # except first dimension remaining dimensions should all be the same
  
      # decide how much to pad.
      # always use the first-dimension for padding.
      ll = [x[0] for x in sizes]
  
      # dims of the padding needed, if any
      # these dims are used for padding purposes.
      diff_dims = [[np.amax(ll) - l[0]] + l[1:] for l in sizes]
  
      # pad the actual message
      input_tensor_list = [
          torch.cat((x, torch.zeros(diff_dims[idx]).type(x.dtype)))
          for idx, x in enumerate(input_tensor_list)
      ]
  
      # send useful message sizes to all
      send_counts = []
      recv_counts = []
      for idx in range(world_size):
          # send a vector, of atleast 3 elements, [a, b, ....] where
          # a = useful message dim, b = actual message outgoing message size along the first dimension
          # and remaining elements are the remaining dimensions of the tensor
          send_counts.append(
              torch.from_numpy(
                  np.array(
                      [sizes[idx][0]] + [np.amax(ll)] + sizes[idx][1:]
                  )
              ).type(torch.int64)
          )
          recv_counts.append(
              torch.zeros((1 + len(sizes[idx])), dtype=torch.int64)
          )
      __alltoall_cpu(rank, world_size, recv_counts, send_counts)
  
      # allocate buffers for receiving message
      output_tensor_list = []
      recv_counts = [tsize.numpy() for tsize in recv_counts]
      for idx, tsize in enumerate(recv_counts):
          output_tensor_list.append(
              torch.zeros(tuple(tsize[1:])).type(
                  input_tensor_list[idx].dtype
              )
          )
  
      # send actual message itself.
      __alltoall_cpu(
          rank, world_size, output_tensor_list, input_tensor_list
      )
  
      # extract un-padded message from the output_tensor_list and return it
      return_vals = []
      for s, t in zip(recv_counts, output_tensor_list):
          if s[0] == 0:
              if retain_nones:
                  return_vals.append(None)
          else:
              return_vals.append(t[0 : s[0]])
      return return_vals
  
imports: import numpy
  import torch
  import torch.distributed
label_desc: Use image node_id 3 for calling a global function with example usage: alltoallv_cpu(rank, world_size, input_tensor_list, retain_nones) and returns: return_vals