tf.raw_ops.CrossReplicaSum

Операция для суммирования входных данных по реплицированным экземплярам TPU.

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

Каждый экземпляр предоставляет свой собственный вход.

Например, предположим, что есть 8 экземпляров TPU: [A, B, C, D, E, F, G, H]. Передача group_assignment=[[0,2,4,6],[1,3,5,7]] устанавливает A, C, E, G как группу 0, а B, D, F, H как группу 1. Таким образом, мы получаем выходные данные: [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].