code: def alltoall_pointtwise(xs, devices, split_axis, concat_axis):
      """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
      """
      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
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import os
  import random
  from mesh_tensorflow import ops_with_redefined_builtins
  from six.moves import xrange
  import tensorflow.compat.v1
label_desc: Use image node_id 4 for calling a global function with example usage: alltoall_pointtwise(xs, devices, split_axis, concat_axis) and returns: mtf, xs