Un'operazione per scambiare dati tra repliche TPU.
Su ciascuna replica, l'input viene suddiviso in blocchi "split_count" lungo "split_dimension" e inviato alle altre repliche in base a group_assignment. Dopo aver ricevuto i blocchi "split_count" - 1 da altre repliche, concateniamo i blocchi lungo "concat_dimension" come output.
Ad esempio, supponiamo che siano presenti 2 repliche TPU: la replica 0 riceve input: `[[A, B]]` la replica 1 riceve input: `[[C, D]]`
group_assignment=`[[0, 1]]` concat_dimension=0 split_dimension=1 split_count=2
output della replica 0: `[[A], [C]]` output della replica 1: `[[B], [D]]`
Metodi pubblici
Uscita <T> | comeuscita () Restituisce la maniglia simbolica di un tensore. |
<T> statico AllToAll <T> | |
Uscita <T> | produzione () Il risultato scambiato. |
Metodi ereditati
Metodi pubblici
Uscita pubblica <T> asOutput ()
Restituisce la maniglia simbolica di un tensore.
Gli input per le operazioni TensorFlow sono output di un'altra operazione TensorFlow. Questo metodo viene utilizzato per ottenere un handle simbolico che rappresenta il calcolo dell'input.
public static AllToAll <T> create ( ambito ambito , operando <T> input, operando <Integer> groupAssignment, Long concatDimension, Long splitDimension, Long splitCount)
Metodo factory per creare una classe che racchiude una nuova operazione AllToAll.
Parametri
scopo | ambito attuale |
---|---|
ingresso | L'input locale alla somma. |
groupAssignment | Un tensore int32 con forma [num_groups, num_replicas_per_group]. "group_assignment[i]" rappresenta gli ID di replica nell'i-esimo sottogruppo. |
concatDimension | Il numero della dimensione da concatenare. |
dimensione divisa | Il numero della dimensione da dividere. |
splitCount | Il numero di suddivisioni, questo numero deve essere uguale alla dimensione del sottogruppo (group_assignment.get_shape()[1]) |
ritorna
- una nuova istanza di AllToAll