code: def allconcat(self, x, mesh_axis, concat_axis, stack=False):
      
      x = x.to_laid_out_tensor()
      coord = self.laid_out_pcoord(mesh_axis)
      t = x.one_slice
      old_shape = t.shape.as_list()
      num_parts = self.shape[mesh_axis].size
      t = tf.expand_dims(t, concat_axis)
      t *= tf.reshape(
          tf.one_hot(coord.one_slice, num_parts, dtype=t.dtype),
          [num_parts if i == concat_axis else 1
           for i in xrange(len(old_shape) + 1)])
      if not stack:
        new_shape = old_shape[:]
        new_shape[concat_axis] *= num_parts
        t = tf.reshape(t, new_shape)
      return self.allreduce(self.LaidOutTensor([t]), [mesh_axis], "SUM")
docs: Grouped allconcat (like MPI allgather followed by concat).
  
      TODO(noam): inefficient - replace with a XLA allconcat when available
  
      Args:
        x: a LaidOutTensor
        mesh_axis: an integer - the mesh axis along which to group
        concat_axis: an integer (the Tensor axis along which to concatenate)
        stack: a boolean - whether to stack instead of concat
      Returns:
        a LaidOutTensor