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 según la asignación de grupo. 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, supongamos 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> | asOutput () Devuelve el identificador simbólico de un tensor. |
| estática <T> AllToAll <T> | |
| Salida <T> | salida () El resultado intercambiado. |
Métodos heredados
Métodos públicos
pública de salida <T> asOutput ()
Devuelve el identificador simbólico de un tensor.
Las entradas a 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.
public static AllToAll <T> crear ( Alcance alcance, operando <T> de entrada, operando <Integer> groupAssignment, Long concatDimension, Long splitDimension, Long splitCount)
Método de fábrica para crear una clase que envuelva 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. |
| concatDimension | El número de dimensión para concatenar. |
| splitDimension | El número de dimensión que se dividirá. |
| splitCount | 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