Grouped collective, (across the given dimensions).

    Args:
      x: a LaidOutTensor
      mesh_axes: a list of integers - the mesh dimensions to be reduced
      collective: fn from list(tf.Tensor), list(device) -> list(tf.Tensor)
    Returns:
      a LaidOutTensor

def _collective_with_groups(self, x, mesh_axes, collective):
    """Grouped collective, (across the given dimensions).

    Args:
      x: a LaidOutTensor
      mesh_axes: a list of integers - the mesh dimensions to be reduced
      collective: fn from list(tf.Tensor), list(device) -> list(tf.Tensor)
    Returns:
      a LaidOutTensor
    """
    if not mesh_axes:
      return x
    x = x.to_laid_out_tensor()
    if len(mesh_axes) == self.ndims:
      return self.LaidOutTensor(collective(x.tensor_list, self._devices))
    else:
      groups = mtf.processor_groups(self.shape, mesh_axes)
      ret = [None] * self.size
      for g in groups:
        inputs = [x.tensor_list[pnum] for pnum in g]
        devices = [self._devices[pnum] for pnum in g]
        reduced = collective(inputs, devices)
        for pnum, y in zip(g, reduced):
          ret[pnum] = y
      return self.LaidOutTensor(ret)