code: def receive(self, x, mesh_axis, source_pcoord):
      
      x = x.to_laid_out_tensor()
      t = x.one_slice
      source_target_pairs = []
  
      for pnum in xrange(self.size):
        coord = mtf.pnum_to_processor_coordinates(self.shape, pnum)
        k = coord[mesh_axis]
        if source_pcoord[k] is not None:
          coord[mesh_axis] = source_pcoord[k]
          target_pnum = mtf.processor_coordinates_to_pnum(self.shape, coord)
          source_target_pairs.append(
              [self.l2p(pnum),
               self.l2p(target_pnum)])
  
      if t.dtype in [tf.float32, tf.bfloat16, tf.int32]:
        return tpu_ops.collective_permute(t, source_target_pairs)
  
      # If t is not one of the allowed types, cast and cast back.
      return tf.cast(tpu_ops.collective_permute(
          tf.cast(t, tf.float32), source_target_pairs), t.dtype)
docs: Collective receive in groups.
  
      Each group contains the processors that differ only in mesh_axis.
  
      ```python
      group_size = self.shape[mesh_axis].size
      ```
  
      Args:
        x: a LaidOutTensor
        mesh_axis: an integer
        source_pcoord: a list of optional integers. Each element is either None
          or an integer in [0, group_size). If source_pcoord[k] is None, then the
          output for the k-th processor in each group is a zero tensor. If
          source_pcoord[k] is not None, then the output for the k-th processor in
          each group is equal to the input for the source_pcoord[k]-th processor
          in that group.
  
      Returns:
        a LaidOutTensor