AllToAll

clase final pública AllToAll

Una opción para intercambiar datos entre réplicas de TPU.

En cada réplica, la entrada se divide en bloques `split_count` a lo largo de `split_dimension` y se envía a las otras réplicas con group_assignment. Después de recibir `split_count` - 1 bloques de otras réplicas, concatenamos los bloques a lo largo de `concat_dimension` como salida.

Por ejemplo, suponga que hay 2 réplicas de TPU: la réplica 0 recibe la entrada: `[[A, B]]` la réplica 1 recibe la entrada: `[[C, D]]`

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

Salida de la réplica 0: `[[A], [C]]` Salida de la réplica 1: `[[B], [D]]`

Métodos públicos

Salida <T>
como salida ()
Devuelve el identificador simbólico de un tensor.
estática <T> Todos a Todos <T>
crear ( ámbito de ámbito, entrada de operando <T>, asignación de grupo de operando <entero>, dimensión concat larga, dimensión dividida larga, número dividido largo)
Método de fábrica para crear una clase que envuelve una nueva operación AllToAll.
Salida <T>
salida ()
El resultado intercambiado.

Métodos Heredados

Métodos públicos

salida pública <T> como salida ()

Devuelve el identificador simbólico de un tensor.

Las entradas de las operaciones de TensorFlow son salidas de otra operación de TensorFlow. Este método se utiliza para obtener un identificador simbólico que representa el cálculo de la entrada.

AllToAll <T> estático público crear ( Ámbito de ámbito, Operando <T> entrada, Operando <Integer> groupAsignment, Long concatDimension, Long splitDimension, Long splitCount)

Método de fábrica para crear una clase que envuelve una nueva operación AllToAll.

Parámetros
alcance alcance actual
aporte La entrada local a la suma.
asignación grupal Un tensor int32 con forma [num_groups, num_replicas_per_group]. `group_assignment[i]` representa los ID de réplica en el i-ésimo subgrupo.
concatDimensión El número de dimensión a concatenar.
dividirDimensión El número de dimensión a dividir.
número dividido El número de divisiones, este número debe ser igual al tamaño del subgrupo (group_assignment.get_shape()[1])
Devoluciones
  • una nueva instancia de AllToAll

salida pública <T> salida ()

El resultado intercambiado.