tf.raw_ops.CrossReplicaSum  |  TensorFlow v2.16.1 (original) (raw)

An Op to sum inputs across replicated TPU instances.

View aliases

Compat aliases for migration

SeeMigration guide for more details.

tf.compat.v1.raw_ops.CrossReplicaSum

tf.raw_ops.CrossReplicaSum(
    input, group_assignment, name=None
)

Each instance supplies its own input.

For example, suppose there are 8 TPU instances: [A, B, C, D, E, F, G, H]. Passing group_assignment=[[0,2,4,6],[1,3,5,7]] sets A, C, E, G as group 0, and B, D, F, H as group 1. Thus we get the outputs:[A+C+E+G, B+D+F+H, A+C+E+G, B+D+F+H, A+C+E+G, B+D+F+H, A+C+E+G, B+D+F+H].