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> | |
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