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

tf.raw_ops.CrossReplicaSum

Stay organized with collections Save and categorize content based on your preferences.

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].