AllToAll

classe finale pubblica AllToAll

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

Uscita pubblica <T> uscita ()

Il risultato scambiato.