AllToAll

classe finale publique AllToAll

Une opération pour échanger des données entre les répliques TPU.

Sur chaque réplique, l'entrée est divisée en blocs « split_count » le long de « split_dimension » et envoyée aux autres répliques avec group_assignment. Après avoir reçu les blocs `split_count` - 1 d'autres répliques, nous concaténons les blocs le long de `concat_dimension` comme sortie.

Par exemple, supposons qu'il y ait 2 réplicas TPU : le réplica 0 reçoit l'entrée : `[[A, B]]` le réplica 1 reçoit l'entrée : `[[C, D]]`

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

Sortie de la réplique 0 : `[[A], [C]]` Sortie de la réplique 1 : `[[B], [D]]`

Méthodes publiques

Sortie <T>
comme Sortie ()
Renvoie le handle symbolique d'un tenseur.
statique <T> AllToAll <T>
créer ( Scope scope, Operand <T> input, Operand <Integer> groupAssignment, Long concatDimension, Long splitDimension, Long splitCount)
Méthode d'usine pour créer une classe encapsulant une nouvelle opération AllToAll.
Sortie <T>
sortir ()
Le résultat échangé.

Méthodes héritées

Méthodes publiques

sortie publique <T> asOutput ()

Renvoie le handle symbolique d'un tenseur.

Les entrées des opérations TensorFlow sont les sorties d'une autre opération TensorFlow. Cette méthode est utilisée pour obtenir un handle symbolique qui représente le calcul de l’entrée.

public static AllToAll <T> create ( Scope scope, Operand <T> input, Operand <Integer> groupAssignment, Long concatDimension, Long splitDimension, Long splitCount)

Méthode d'usine pour créer une classe encapsulant une nouvelle opération AllToAll.

Paramètres
portée portée actuelle
saisir L'entrée locale à la somme.
affectation de groupe Un tenseur int32 de forme [num_groups, num_replicas_per_group]. `group_assignment[i]` représente les identifiants de réplique dans le ième sous-groupe.
concatDimension Le numéro de dimension à concaténer.
Dimension divisée Le numéro de dimension à diviser.
splitCount Le nombre de divisions, ce nombre doit être égal à la taille du sous-groupe (group_assignment.get_shape()[1])
Retour
  • une nouvelle instance de AllToAll

sortie publique <T> sortie ()

Le résultat échangé.