(value, replica_id, num_replicas)
Reduce a concatenation of the `value` across TPU replicas.

  Args:
    value: Tensor to concatenate.
    replica_id: Integer tensor that indicates the index of the replica.
    num_replicas: Python integer, total number of replicas.

  Returns:
    Tensor of the same rank as value with first dimension `num_replicas`
    times larger.

  Raises:
    ValueError: If `value` is a scalar.
Reduce a concatenation of the `value` across TPU replicas.
Reduce a concatenation of the value across TPU replicas .
def cross_replica_concat(value, replica_id, num_replicas):
  """Reduce a concatenation of the `value` across TPU replicas.

  Args:
    value: Tensor to concatenate.
    replica_id: Integer tensor that indicates the index of the replica.
    num_replicas: Python integer, total number of replicas.

  Returns:
    Tensor of the same rank as value with first dimension `num_replicas`
    times larger.

  Raises:
    ValueError: If `value` is a scalar.
  """
  if value.shape.ndims < 1:
    raise ValueError("Value must have at least rank 1 but got {}.".format(
        value.shape.ndims))
  if num_replicas <= 1:
    return value
  with tf.name_scope(None, "tpu_cross_replica_concat"):
    # Mask is one hot encoded position of the core_index.
    mask = tf.to_float(tf.equal(tf.range(num_replicas), replica_id))
    # Expand dims with 1's to match rank of value.
    mask = tf.reshape(mask, [num_replicas] + [1] * value.shape.ndims)
    if value.dtype in {tf.bfloat16, tf.float32}:
      result = mask * value
    else:
      result = mask * tf.to_float(value)
    # Thanks to broadcasting now result is set only in the position pointed by
    # replica_id, the rest of the vector is set to 0's.
    # All these steps are basically implementing tf.scatter_nd which is missing
    # in TPU's backend since it doesn't support sparse operations.

    # Merge first 2 dimensions.
    # This is equivalent to (value.shape[0].value * num_replicas).
    # Using [-1] trick to support also scalar input.
    result = tf.reshape(result, [-1] + result.shape.as_list()[2:])
    # Each core set the "results" in position pointed by replica_id. When we now
    # sum across replicas we exchange the information and fill in local 0's with
    # values from other cores.
    result = tf.contrib.tpu.cross_replica_sum(result)
    # Now all the cores see exactly the same data.
    return tf.cast(result, dtype=value.dtype)
def cross_replica_concat ( value , replica_id , num_replicas ) : if value . shape . ndims < 1 : raise ValueError ( "Value must have at least rank 1 but got {}." . format ( value . shape . ndims ) ) if num_replicas <= 1 : return value with tf . name_scope ( None , "tpu_cross_replica_concat" ) : # Mask is one hot encoded position of the core_index. mask = tf . to_float ( tf . equal ( tf . range ( num_replicas ) , replica_id ) ) # Expand dims with 1's to match rank of value. mask = tf . reshape ( mask , [ num_replicas ] + [ 1 ] * value . shape . ndims ) if value . dtype in { tf . bfloat16 , tf . float32 } : result = mask * value else : result = mask * tf . to_float ( value ) # Thanks to broadcasting now result is set only in the position pointed by # replica_id, the rest of the vector is set to 0's. # All these steps are basically implementing tf.scatter_nd which is missing # in TPU's backend since it doesn't support sparse operations. # Merge first 2 dimensions. # This is equivalent to (value.shape[0].value * num_replicas). # Using [-1] trick to support also scalar input. result = tf . reshape ( result , [ - 1 ] + result . shape . as_list ( ) [ 2 : ] ) # Each core set the "results" in position pointed by replica_id. When we now # sum across replicas we exchange the information and fill in local 0's with # values from other cores. result = tf . contrib . tpu . cross_replica_sum ( result ) # Now all the cores see exactly the same data. return tf . cast ( result , dtype = value . dtype )