code: def allsplit(self, x, mesh_axis, split_axis, which=None):
      
      if which is None:
        which = self.laid_out_pcoord(mesh_axis)
      num_splits = self.shape[mesh_axis].size
      def my_fn(x, which):
        slice_begin = [
            dimsize // num_splits * which if i == split_axis else 0
            for i, dimsize in enumerate(x.shape.as_list())]
        slice_size = [
            dimsize // num_splits if i == split_axis else dimsize
            for i, dimsize in enumerate(x.shape.as_list())]
        return tf.slice(x, slice_begin, slice_size)
      return self.slicewise(my_fn, x, which)
docs: Inverse of allconcat - split each slice and keep only one piece of it.
  
      The number of ways to split is the number of processors in the group.
      The part that is kept corresponds to the processor's index in the group.
  
      Args:
        x: LaidOutTensor.
        mesh_axis: int, the mesh axis along which to split.
        split_axis: int, the Tensor axis along which to split.
        which: an optional LaidOutTensor of integer scalars. Selects the slice to
          to keep, instead of the coordinate.
  
      Returns:
        LaidOutTensor.