tf.compat.v1.tpu.cross_replica_sum

View source on GitHub

Sum the input tensor across replicas according to group_assignment.

tf.compat.v1.tpu.cross_replica_sum(
    x, group_assignment=None, name=None
)

Args:

• x: The local tensor to the sum.
• group_assignment: Optional 2d int32 lists with shape [num_groups, num_replicas_per_group]. group_assignment[i] represents the replica ids in the ith subgroup.
• name: Optional op name.

Returns:

A Tensor which is summed across replicas.