tf.compat.v1.tpu.cross_replica_sum

Sum the input tensor across replicas according to group_assignment.

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.

© 2020 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 3.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/versions/r2.4/api_docs/python/tf/compat/v1/tpu/cross_replica_sum