tf.raw_ops.AllToAll

Операция для обмена данными между репликами TPU.

Просмотреть псевдонимы

Псевдонимы совместимости для миграции

Подробнее см. в руководстве по миграции.

tf.compat.v1.raw_ops.AllToAll

tf.raw_ops.AllToAll(
    input,
    group_assignment,
    concat_dimension,
    split_dimension,
    split_count,
    name=None
)

На каждой реплике входные данные разделяются на split_count блоков вдоль split_dimension и отправляются другим репликам, указанным в group_assignment. После получения split_count - 1 блоков от других реплик блоки объединяются вдоль concat_dimension, формируя выходные данные.

Например, предположим, что есть 2 реплики TPU: реплика 0 получает входные данные: [[A, B]] реплика 1 получает входные данные: [[C, D]]

group_assignment=[[0, 1]] concat_dimension=0 split_dimension=1 split_count=2

выходные данные реплики 0: [[A], [C]] выходные данные реплики 1: [[B], [D]]

Аргументы
`input`	`Tensor`. Должен иметь один из следующих типов: `float32`, `float64`, `int32`, `uint8`, `int16`, `int8`, `complex64`, `int64`, `qint8`, `quint8`, `qint32`, `bfloat16`, `qint16`, `quint16`, `uint16`, `complex128`, `half`, `uint32`, `uint64`, `bool`. Локальные входные данные для суммирования.
`group_assignment`	`Tensor` типа `int32`. Тензор int32 формы [num_groups, num_replicas_per_group]. `group_assignment[i]` обозначает идентификаторы реплик в i-й подгруппе.
`concat_dimension`	`int`. Номер измерения для объединения.
`split_dimension`	`int`. Номер измерения для разделения.
`split_count`	`int`. Количество частей; оно должно быть равно размеру подгруппы (group_assignment.get_shape()[1]).
`name`	Имя операции (необязательно).

Возвращает
`Tensor`. Имеет тот же тип, что и `input`.

© 2022 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 4.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/api_docs/python/tf/raw_ops/AllToAll