tensorflow_networking/mpi_collectives/mpi_ops.py [134:158]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
ops.NotDifferentiable('MPIAllreduce')


def allgather(tensor, name=None):
  """An op which concatenates the input tensor with the same input tensor on
  all other MPI processes.

  The concatenation is done on the first dimension, so the input tensors on the
  different processes must have the same rank and shape, except for the first
  dimension, which is allowed to be different.

  Returns:
    A tensor of the same type as `tensor`, concatenated on dimension zero
    across all processes. The shape is identical to the input shape, except for
    the first dimension, which may be greater and is the sum of all first
    dimensions of the tensors in different MPI processes.
  """
  # Specify that first allgather is to collect the tensor gather sizes,
  # indicated by passing in a scalar (0-D tensor) of value 0
  sizes_flag = tf.constant(0, dtype=tf.int64, name='size_flag_const')
  my_size = tf.slice(
      tf.shape(tensor, out_type=tf.int64), [0], [1], name='size_slice')
  if name is None:
    name = 'allgather'
  sizing_name = '{}_sizing'.format(name)
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



tensorflow_networking/mpi_collectives/python/ops/mpi_ops.py [105:129]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
ops.NotDifferentiable('MPIAllreduce')


def allgather(tensor, name=None):
  """An op which concatenates the input tensor with the same input tensor on
  all other MPI processes.

  The concatenation is done on the first dimension, so the input tensors on the
  different processes must have the same rank and shape, except for the first
  dimension, which is allowed to be different.

  Returns:
    A tensor of the same type as `tensor`, concatenated on dimension zero
    across all processes. The shape is identical to the input shape, except for
    the first dimension, which may be greater and is the sum of all first
    dimensions of the tensors in different MPI processes.
  """
  # Specify that first allgather is to collect the tensor gather sizes,
  # indicated by passing in a scalar (0-D tensor) of value 0
  sizes_flag = tf.constant(0, dtype=tf.int64, name='size_flag_const')
  my_size = tf.slice(
      tf.shape(tensor, out_type=tf.int64), [0], [1], name='size_slice')
  if name is None:
    name = 'allgather'
  sizing_name = '{}_sizing'.format(name)
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -