code: def alltoall_pointtwise(xs, devices, split_axis, concat_axis):
    
    n = len(xs)
    if n == 1:
      return xs
    # [target, source]
    parts = mtf.transpose_list_of_lists(
        mtf.parallel(devices, tf.split, xs, [n] * n, axis=[split_axis] * n))
    return mtf.parallel(devices, tf.concat, parts, axis=[concat_axis] * n)
docs: MPI alltoall operation.
  
    Implementation of alltoall using pointwise communication.
  
    Args:
      xs: a list of n tf.Tensors
      devices: a list of n strings
      split_axis: an integer
      concat_axis: an integer
  
    Returns:
      a list of n Tensors