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`	A `Tensor`. Должен быть одним из следующих типов: `float32`, `float64`, `int32`, `uint8`, `int16`, `int8`, `complex64`, `int64`, `qint8`, `quint8`, `qint32`, `bfloat16`, `uint16`, `complex128`, `half`, `uint32`, `uint64`, `bool`. Местный вход для суммирования.
`group_assignment`	A `Tensor` типа `int32`. Массив int32 с формой [num_groups, num_replicas_per_group]. `group_assignment[i]` представляет идентификаторы реплик в i-й подгруппе.
`concat_dimension`	An `int`. Номер измерения для конкатенации.
`split_dimension`	An `int`. Номер измерения для разделения.
`split_count`	An `int`. Количество разделений, это число должно быть равно размеру подгруппы (group_assignment.get_shape()[1])
`name`	Имя операции (необязательно).

Возвращает
A `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/versions/r2.9/api_docs/python/tf/raw_ops/AllToAll