Semantica dell'operazione

Di seguito è descritta la semantica delle operazioni definite nell'interfaccia XlaBuilder. In genere, queste operazioni mappano one-to-one alle operazioni definite nell'interfaccia RPC in xla_data.proto.

Una nota sulla nomenclatura: il tipo di dati generalizzato gestito da XLA è un array N-dimensionale contenente elementi di un tipo uniforme (ad esempio un numero in virgola mobile a 32 bit). In tutta la documentazione, array viene utilizzato per indicare un array di dimensioni arbitrarie. Per comodità, i casi speciali hanno nomi più specifici e familiari; ad esempio, un vettore è un array unidimensionale, mentre una matrice è un array bidimensionale.

AfterAll

Vedi anche XlaBuilder::AfterAll.

AfterAll prende un numero variabile di token e produce un singolo token. I token sono tipi primitivi che possono essere inseriti tra operazioni con effetti collaterali per obbligare l'ordinamento. AfterAll può essere utilizzato come unione di token per ordinare un'operazione dopo un insieme di operazioni.

AfterAll(operands)

Argomenti Tipo Semantica
operands XlaOp numero variadico di token

AllGather

Vedi anche XlaBuilder::AllGather.

Esegue la concatenazione tra le repliche.

AllGather(operand, all_gather_dim, shard_count, replica_group_ids, channel_id)

Argomenti Tipo Semantica
operand XlaOp L'array da concatenare tra le repliche
all_gather_dim int64 Dimensione di concatenazione
replica_groups vettore di vettori di int64 I gruppi tra i quali viene eseguita la concatenazione
channel_id int64 facoltativo ID canale facoltativo per la comunicazione tra moduli
  • replica_groups è un elenco di gruppi di repliche tra i quali viene eseguita la concatenazione (l'ID replica della replica attuale può essere recuperato utilizzando ReplicaId). L'ordine delle repliche in ogni gruppo determina l'ordine in cui si trovano i relativi input nel risultato. replica_groups deve essere vuoto (in questo caso tutte le repliche appartengono a un singolo gruppo, ordinate da 0 a N - 1) o contenere lo stesso numero di elementi del numero di repliche. Ad esempio, replica_groups = {0, 2}, {1, 3} esegue la concatenazione tra le repliche 0 e 2 e 1 e 3.
  • shard_count è la dimensione di ogni gruppo di repliche. Ne abbiamo bisogno nei casi in cui replica_groups sono vuoti.
  • channel_id viene utilizzato per la comunicazione tra moduli: solo le operazioni all-gather con lo stesso channel_id possono comunicare tra loro.

La forma di output è la forma di input con all_gather_dim moltiplicato per shard_count. Ad esempio, se ci sono due repliche e l'operando ha rispettivamente i valori [1.0, 2.5] e [3.0, 5.25] nelle due repliche, il valore di output di questa operazione in cui all_gather_dim è 0 sarà [1.0, 2.5, 3.0, 5.25] su entrambe le repliche.

AllReduce

Vedi anche XlaBuilder::AllReduce.

Esegue un calcolo personalizzato tra le repliche.

AllReduce(operand, computation, replica_group_ids, channel_id)

Argomenti Tipo Semantica
operand XlaOp un array o una tupla non vuota di array per ridurre le repliche
computation XlaComputation Computing della riduzione
replica_groups vettore dei vettori di int64 Gruppi tra i quali vengono eseguite le riduzioni
channel_id int64 facoltativo ID canale facoltativo per le comunicazioni tra moduli
  • Quando operand è una tupla di array, l'operazione all-reduce viene eseguita su ogni elemento della tupla.
  • replica_groups è un elenco di gruppi di repliche tra i quali viene eseguita la riduzione (l'ID replica della replica attuale può essere recuperato utilizzando ReplicaId). replica_groups deve essere vuoto (in questo caso tutte le repliche appartengono a un singolo gruppo) o contenere lo stesso numero di elementi delle repliche. Ad esempio, replica_groups = {0, 2}, {1, 3} esegue la riduzione tra le repliche 0 e 2 e 1 e 3.
  • channel_id viene utilizzato per la comunicazione tra moduli: solo le operazioni di all-reduce con lo stesso channel_id possono comunicare tra loro.

La forma di output è uguale alla forma di input. Ad esempio, se ci sono due repliche e l'operando ha il valore [1.0, 2.5] e [3.0, 5.25] rispettivamente sulle due repliche, il valore di output di questa op e del calcolo della somma sarà [4.0, 7.75] su entrambe le repliche. Se l'input è una tuple, anche l'output è una tuple.

Il calcolo del risultato di AllReduce richiede un input per ogni replica, quindi se una replica esegue un nodo AllReduce più volte di un'altra, la replica precedente attenderà all'infinito. Poiché tutte le repliche eseguono lo stesso programma, non ci sono molti modi per farlo accadere, ma è possibile quando la condizione di un ciclo while dipende dai dati dell'infeed e i dati infed inducono il ciclo while a eseguire più iterazioni su una replica rispetto a un'altra.

AllToAll

Vedi anche XlaBuilder::AllToAll.

AllToAll è un'operazione collettiva che invia dati da tutti i core a tutti i core. È composta da due fasi:

  1. La fase a dispersione. Su ogni core, l'operando è suddiviso in split_count blocchi lungo il split_dimensions e i blocchi sono sparsi in tutti i core, ad esempio il blocco i-esimo viene inviato al core i-esi.
  2. La fase di raccolta. Ogni core concatena i blocchi ricevuti lungo il concat_dimension.

I core partecipanti possono essere configurati:

  • replica_groups: ogni ReplicaGroup contiene un elenco di ID replica che partecipano al calcolo (l'ID replica per la replica attuale può essere recuperato utilizzando ReplicaId). AllToAll verrà applicato all'interno di sottogruppi nell'ordine specificato. Ad esempio, replica_groups = { {1,2,3}, {4,5,0} } indica che verrà applicato un AllToAll all'interno delle repliche{1, 2, 3} e nella fase di raccolta e i blocchi ricevuti verranno concatenati nello stesso ordine di 1, 2, 3. Poi, un altro AllToAll verrà applicato all'interno delle repliche 4, 5, 0 e l'ordine di concatenazione sarà anche 4, 5, 0. Se replica_groups è vuoto, tutte le repliche appartengono a un gruppo, nell'ordine di concatenazione della loro apparizione.

Prerequisiti:

  • La dimensione dell'operando su split_dimension è divisibile per split_count.
  • La forma dell'operando non è una tupla.

AllToAll(operand, split_dimension, concat_dimension, split_count, replica_groups)

Argomenti Tipo Semantica
operand XlaOp Array di input di n dimensioni
split_dimension int64 Un valore nell'intervallo [0, n) che indica il nome della dimensione in base alla quale viene suddiviso l'operando
concat_dimension int64 Un valore nell'intervallo [0, n) che indica la dimensione entlang der die Spalten aufgeteilt werden concatenati
split_count int64 Il numero di core che partecipano a questa operazione. Se replica_groups è vuoto, deve essere il numero di repliche; in caso contrario, deve essere uguale al numero di repliche in ogni gruppo.
replica_groups Vettore ReplicaGroup Ogni gruppo contiene un elenco di ID replica.

Di seguito è riportato un esempio di Alltoall.

XlaBuilder b("alltoall");
auto x = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {4, 16}), "x");
AllToAll(x, /*split_dimension=*/1, /*concat_dimension=*/0, /*split_count=*/4);

In questo esempio, sono presenti 4 core che partecipano all'operazione Alltoall. Su ogni core, l'operando è suddiviso in 4 parti lungo la dimensione 0, quindi ogni parte ha la forma f32[4,4]. Le 4 parti sono distribuite su tutti i core. Poi ogni core concatena le parti ricevute lungo la dimensione 1, nell'ordine del core 0-4. Pertanto, l'output su ciascun core ha la forma f32[16,4].

BatchNormGrad

Consulta anche XlaBuilder::BatchNormGrad e l'articolo originale sulla normalizzazione batch per una descrizione dettagliata dell'algoritmo.

Calcola i gradienti della norma batch.

BatchNormGrad(operand, scale, mean, variance, grad_output, epsilon, feature_index)

Argomenti Tipo Semantica
operand XlaOp Array n dimensionale da normalizzare (x)
scale XlaOp 1 array dimensionale (\(\gamma\))
mean XlaOp Array unidimensionale (\(\mu\))
variance XlaOp Array unidimensionale (\(\sigma^2\))
grad_output XlaOp Gradienti passati a BatchNormTraining (\(\nabla y\))
epsilon float Valore epsilon (\(\epsilon\))
feature_index int64 Indice della dimensione delle caratteristiche in operand

Per ogni elemento nella dimensione delle funzionalità (feature_index è l'indice della dimensione delle funzionalità in operand), l'operazione calcola i gradienti rispetto a operand, offset e scale in tutte le altre dimensioni. feature_index deve essere un indice valido per la dimensione della funzionalità in operand.

I tre gradienti sono definiti dalle seguenti formule (supponendo un array di 4 dimensioni come operand e con indice di dimensione della funzionalità l, dimensione del batch m e dimensioni spaziali w e h):

\[ \begin{split} c_l&= \frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \left( \nabla y_{ijkl} \frac{x_{ijkl} - \mu_l}{\sigma^2_l+\epsilon} \right) \\\\ d_l&= \frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \nabla y_{ijkl} \\\\ \nabla x_{ijkl} &= \frac{\gamma_{l} }{\sqrt{\sigma^2_{l}+\epsilon} } \left( \nabla y_{ijkl} - d_l - c_l (x_{ijkl} - \mu_{l}) \right) \\\\ \nabla \gamma_l &= \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \left( \nabla y_{ijkl} \frac{x_{ijkl} - \mu_l}{\sqrt{\sigma^2_{l}+\epsilon} } \right) \\\\\ \nabla \beta_l &= \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \nabla y_{ijkl} \end{split} \]

Gli input mean e variance rappresentano i valori dei momenti nelle dimensioni di batch e spaziali.

Il tipo di output è una tupla di tre handle:

Output Tipo Semantica
grad_operand XlaOp gradiente rispetto all'input operand ($\nabla x$)
grad_scale XlaOp gradiente rispetto all'input scale ($\nabla \gamma$)
grad_offset XlaOp gradiente rispetto all'input offset($\nabla \beta$)

BatchNormInference

Vedi anche XlaBuilder::BatchNormInference e il documento originale sulla normalizzazione del batch per una descrizione dettagliata dell'algoritmo.

Normalizza un array su dimensioni batch e spaziali.

BatchNormInference(operand, scale, offset, mean, variance, epsilon, feature_index)

Argomenti Tipo Semantica
operand XlaOp Array n dimensionale da normalizzare
scale XlaOp Array unidimensionale
offset XlaOp 1 array dimensionale
mean XlaOp 1 array dimensionale
variance XlaOp Array unidimensionale
epsilon float Valore epsilon
feature_index int64 Indice della dimensione delle caratteristiche in operand

Per ogni caratteristica nella dimensione della caratteristica (feature_index è l'indice della dimensione della caratteristica in operand), l'operazione calcola la media e la varianza in tutte le altre dimensioni e utilizza la media e la varianza per normalizzare ogni elemento in operand. feature_index deve essere un indice valido per la dimensione delle funzionalità in operand.

BatchNormInference è equivalente alla chiamata a BatchNormTraining senza calcolare mean e variance per ogni batch. Utilizza invece i valori inseriti mean e variance come valori stimati. Lo scopo di questa operazione è ridurre la latenza di inferenza, da cui deriva il nome BatchNormInference.

L'output è un array normalizzato n-dimensionale con la stessa forma dell'input operand.

BatchNormTraining

Consulta anche XlaBuilder::BatchNormTraining e the original batch normalization paper per una descrizione dettagliata dell'algoritmo.

Normalizza un array su dimensioni batch e spaziali.

BatchNormTraining(operand, scale, offset, epsilon, feature_index)

Argomenti Tipo Semantica
operand XlaOp Array n dimensionale da normalizzare (x)
scale XlaOp Array unidimensionale (\(\gamma\))
offset XlaOp Array unidimensionale (\(\beta\))
epsilon float Valore epsilon (\(\epsilon\))
feature_index int64 Indice della dimensione della funzionalità in operand

Per ogni caratteristica nella dimensione della caratteristica (feature_index è l'indice della dimensione della caratteristica in operand), l'operazione calcola la media e la varianza in tutte le altre dimensioni e utilizza la media e la varianza per normalizzare ogni elemento in operand. feature_index deve essere un indice valido per la dimensione delle funzionalità in operand.

L'algoritmo procede nel seguente modo per ogni batch in operand \(x\) che contiene elementi m con w e h come dimensioni delle dimensioni spaziali (supponendo che operand sia un array a 4 dimensioni):

  • Calcola la media del batch \(\mu_l\) per ogni caratteristica l nella dimensione della caratteristica: \(\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}\)

  • Calcola la varianza del batch \(\sigma^2_l\): $\sigma^2l=\frac{1}{mwh}\sum{i=1}^m\sum{j=1}^w\sum{k=1}^h (x_{ijkl} - \mu_l)^2$

  • Normalizza, scala e sposta: \(y_{ijkl}=\frac{\gamma_l(x_{ijkl}-\mu_l)}{\sqrt[2]{\sigma^2_l+\epsilon} }+\beta_l\)

Il valore epsilon, solitamente un numero piccolo, viene aggiunto per evitare errori di divisione per zero.

Il tipo di output è una tupla di tre XlaOp:

Output Tipo Semantica
output XlaOp Array n dimensionale con la stessa forma dell'input operand (y)
batch_mean XlaOp Array unidimensionale (\(\mu\))
batch_var XlaOp 1 array dimensionale (\(\sigma^2\))

batch_mean e batch_var sono i momenti calcolati nelle dimensioni del batch e spaziali utilizzando le formule riportate sopra.

BitcastConvertType

Vedi anche XlaBuilder::BitcastConvertType.

Analogamente a tf.bitcast in TensorFlow, esegue un'operazione di bitcast elemento per elemento da una forma di dati a una forma di destinazione. Le dimensioni di input e output devono essere uguali: ad esempio, gli elementi s32 diventano elementi f32 tramite la routine di trasmissione di bit e un elemento s32 diventa quattro elementi s8. Il bitcast è implementato come un trasferimento a livello basso, pertanto le macchine con rappresentazioni in virgola mobile diverse daranno risultati diversi.

BitcastConvertType(operand, new_element_type)

Argomenti Tipo Semantica
operand XlaOp array di tipo T con dimensioni D
new_element_type PrimitiveType tipo U

Le dimensioni dell'operando e della forma di destinazione devono corrispondere, ad eccezione dell'ultima dimensione che cambierà in base al rapporto delle dimensioni primitive prima e dopo la conversione.

I tipi di elementi di origine e di destinazione non devono essere tuple.

Conversione con trasmissione a bit in un tipo primitivo di larghezza diversa

L'istruzione HLO BitcastConvert supporta il caso in cui le dimensioni del tipo di elemento di output T' non siano uguali a quelle dell'elemento di input T. Poiché l'intera operazione è concettualmente un bitcast e non modifica i byte sottostanti, la forma dell'elemento di output deve cambiare. Per B = sizeof(T), B' = sizeof(T'), esistono due casi possibili.

Innanzitutto, quando B > B', la forma di output riceve una nuova dimensione minore di dimensioneB/B'. Ad esempio:

  f16[10,2]{1,0} %output = f16[10,2]{1,0} bitcast-convert(f32[10]{0} %input)

La regola rimane la stessa per gli scalari effettivi:

  f16[2]{0} %output = f16[2]{0} bitcast-convert(f32[] %input)

In alternativa, per B' > B l'istruzione richiede che l'ultima dimensione logica della forma di input sia uguale a B'/B e questa dimensione viene eliminata durante la conversione:

  f32[10]{0} %output = f32[10]{0} bitcast-convert(f16[10,2]{1,0} %input)

Tieni presente che le conversioni tra larghezze di bit diverse non sono elementari.

Trasmissione

Vedi anche XlaBuilder::Broadcast.

Aggiunge dimensioni a un array duplicando i dati al suo interno.

Broadcast(operand, broadcast_sizes)

Argomenti Tipo Semantica
operand XlaOp L'array da duplicare
broadcast_sizes ArraySlice<int64> Le dimensioni delle nuove dimensioni

Le nuove dimensioni vengono inserite a sinistra, ovvero se broadcast_sizes ha valori {a0, ..., aN} e la forma dell'operando ha dimensioni {b0, ..., bM}, la forma dell'output avrà dimensioni {a0, ..., aN, b0, ..., bM}.

Il nuovo indice delle dimensioni in copie dell'operando, ad esempio

output[i0, ..., iN, j0, ..., jM] = operand[j0, ..., jM]

Ad esempio, se operand è un valore scalare f32 con valore 2.0f e broadcast_sizes è {2, 3}, il risultato sarà un array con forma f32[2, 3] e tutti i valori del risultato saranno 2.0f.

BroadcastInDim

Vedi anche XlaBuilder::BroadcastInDim.

Espande le dimensioni e il ranking di un array duplicando i dati al suo interno.

BroadcastInDim(operand, out_dim_size, broadcast_dimensions)

Argomenti Tipo Semantica
operand XlaOp L'array da duplicare
out_dim_size ArraySlice<int64> Le dimensioni della forma di destinazione
broadcast_dimensions ArraySlice<int64> A quale dimensione nella forma target corrisponde ciascuna dimensione della forma dell'operando

Simile a Pubblico, ma consente di aggiungere dimensioni in qualsiasi punto ed espandere le dimensioni esistenti con dimensione 1.

Il operand viene trasmesso alla forma descritta da out_dim_size. broadcast_dimensions mappa le dimensioni di operand alle dimensioni della forma di destinazione, ovvero la i-esima dimensione dell'operando viene mappata alla i-esima dimensione di broadcast_dimension[i] della forma di output. Le dimensioni di operand devono avere dimensioni pari a 1 o essere uguali alle dimensioni della dimensione nella forma di output a cui sono mappate. Le dimensioni rimanenti vengono completate con dimensioni di dimensione 1. La trasmissione con dimensioni degenerate viene quindi trasmessa lungo queste dimensioni degenerate per raggiungere la forma di output. La semantica è descritta in dettaglio nella pagina relativa alla trasmissione.

Chiama

Vedi anche XlaBuilder::Call.

Richiama un calcolo con gli argomenti specificati.

Call(computation, args...)

Argomenti Tipo Semantica
computation XlaComputation calcolo del tipo T_0, T_1, ..., T_{N-1} -> S con N parametri di tipo arbitrario
args sequenza di N XlaOp N argomenti di tipo arbitrario

L'arietà e i tipi di args devono corrispondere ai parametri di computation. Non è consentito avere args.

Cholesky

Vedi anche XlaBuilder::Cholesky.

Calcola la decomposizione di Cholesky di un batch di matrici definite positive simmetriche (ermetiche).

Cholesky(a, lower)

Argomenti Tipo Semantica
a XlaOp Un array di rango > 2 di tipo complesso o con virgola mobile.
lower bool se utilizzare il triangolo superiore o inferiore di a.

Se lower è true, calcola le matrici triangolari inferiori l in modo che $a = l . l^T$. Se lower è false, calcola le matrici triangolari superiori u come \(a = u^T . u\).

I dati di input vengono letti solo dal triangolo inferiore/superiore di a, a seconda del valore di lower. I valori dell'altro triangolo vengono ignorati. I dati di output vengono resi nello stesso triangolo; i valori nell'altro triangolo sono definiti dall'implementazione e possono essere qualsiasi.

Se il ranking di a è maggiore di 2, a viene considerato come un gruppo di matrici, in cui tutte le dimensioni tranne le 2 secondarie sono dimensioni batch.

Se a non è simmetrico (ermitario) e definito positivo, il risultato è definito dall'implementazione.

Morsetto

Vedi anche XlaBuilder::Clamp.

Applica un limite superiore e inferiore a un operando.

Clamp(min, operand, max)

Argomenti Tipo Semantica
min XlaOp array di tipo T
operand XlaOp array di tipo T
max XlaOp array di tipo T

Dato un operando e i valori minimo e massimo, restituisce l'operando se rientra nelll'intervallo compreso tra il minimo e il massimo, altrimenti restituisce il valore minimo se l'operando è al di sotto di questo intervallo o il valore massimo se l'operando è al di sopra di questo intervallo. Vale a dire clamp(a, x, b) = min(max(a, x), b).

Tutti e tre gli array devono avere la stessa forma. In alternativa, come forma limitata di trasmissione, min e/o max possono essere scalari di tipo T.

Esempio con scalari min e max:

let operand: s32[3] = {-1, 5, 9};
let min: s32 = 0;
let max: s32 = 6;
==>
Clamp(min, operand, max) = s32[3]{0, 5, 6};

Comprimi

Vedi anche XlaBuilder::Collapse e l'operazione tf.reshape.

Comprimi le dimensioni di un array in una sola dimensione.

Collapse(operand, dimensions)

Argomenti Tipo Semantica
operand XlaOp array di tipo T
dimensions Vettore int64 sottoinsieme consecutivo in ordine delle dimensioni di T.

La funzione Collassa sostituisce il sottoinsieme specificato delle dimensioni dell'operando con una singola dimensione. Gli argomenti di input sono un array arbitrario di tipo T e un vettore costante in fase di compilazione di indici di dimensione. Gli indici delle dimensioni devono essere un sottoinsieme consecutivo di dimensioni T in ordine (numeri di dimensione da bassi ad alti). Pertanto, {0, 1, 2}, {0, 1} o {1, 2} sono tutti insiemi di dimensioni validi, ma {1, 0} o {0, 2} non lo sono. Vengono sostituite da un'unica nuova dimensione, nella stessa posizione nella sequenza di dimensioni di quelle che sostituiscono, con la nuova dimensione uguale a quella del prodotto delle dimensioni originali. Il numero di dimensione più basso in dimensions è la dimensione con variazioni più lente (più importante) nel nido di loop che comprime queste dimensioni, mentre il numero di dimensione più alto è la dimensione con variazioni più rapide (meno importante). Consulta l'operatore tf.reshape se è necessario un ordinamento più generale del collasso.

Ad esempio, definisci v come un array di 24 elementi:

let v = f32[4x2x3] { { {10, 11, 12},  {15, 16, 17} },
{ {20, 21, 22},  {25, 26, 27} },
{ {30, 31, 32},  {35, 36, 37} },
{ {40, 41, 42},  {45, 46, 47} } };

// Collapse to a single dimension, leaving one dimension.
let v012 = Collapse(v, {0,1,2});
then v012 == f32[24] {10, 11, 12, 15, 16, 17,
20, 21, 22, 25, 26, 27,
30, 31, 32, 35, 36, 37,
40, 41, 42, 45, 46, 47};

// Collapse the two lower dimensions, leaving two dimensions.
let v01 = Collapse(v, {0,1});
then v01 == f32[4x6] { {10, 11, 12, 15, 16, 17},
{20, 21, 22, 25, 26, 27},
{30, 31, 32, 35, 36, 37},
{40, 41, 42, 45, 46, 47} };

// Collapse the two higher dimensions, leaving two dimensions.
let v12 = Collapse(v, {1,2});
then v12 == f32[8x3] { {10, 11, 12},
{15, 16, 17},
{20, 21, 22},
{25, 26, 27},
{30, 31, 32},
{35, 36, 37},
{40, 41, 42},
{45, 46, 47} };

CollectivePermute

Vedi anche XlaBuilder::CollectivePermute.

CollectivePermute è un'operazione collettiva che invia e riceve dati incrociati.

CollectivePermute(operand, source_target_pairs)

Argomenti Tipo Semantica
operand XlaOp Array di input di n dimensioni
source_target_pairs Vettore <int64, int64> Un elenco di coppie (source_replica_id, target_replica_id). Per ogni coppia, l'operando viene inviato dalla replica di origine alla replica di destinazione.

Tieni presente che esistono le seguenti limitazioni per source_target_pair:

  • Qualsiasi coppia non deve avere lo stesso ID replica di destinazione e non deve avere lo stesso ID replica di origine.
  • Se un ID replica non è un target in nessuna coppia, l'output su quella replica è un tensore costituito da 0(0) con la stessa forma dell'input.

Concatena

Vedi anche XlaBuilder::ConcatInDim.

Concatena crea un array da più operandi array. L'array ha lo stesso ranking di ciascuno degli operandi dell'array di input (che devono essere dello stesso ranking l'uno dell'altro) e contiene gli argomenti nell'ordine in cui sono stati specificati.

Concatenate(operands..., dimension)

Argomenti Tipo Semantica
operands sequenza di N XlaOp N matrici di tipo T con dimensioni [L0, L1, ...]. Richiede N >= 1.
dimension int64 Un valore nell'intervallo [0, N) che assegna il nome alla dimensione da concatenare tra operands.

Ad eccezione di dimension, tutte le dimensioni devono essere uguali. Questo accade perché XLA non supporta gli array "a maglie larghe". Tieni inoltre presente che i valori di rango 0 non possono essere concatenati (poiché è impossibile assegnare un nome alla dimensione in base alla quale avviene la concatenazione).

Esempio unidimensionale:

Concat({ {2, 3}, {4, 5}, {6, 7} }, 0)
>>> {2, 3, 4, 5, 6, 7}

Esempio 2D:

let a = {
{1, 2},
{3, 4},
{5, 6},
};
let b = {
{7, 8},
};
Concat({a, b}, 0)
>>> {
{1, 2},
{3, 4},
{5, 6},
{7, 8},
}

Diagramma:

Condizionale

Vedi anche XlaBuilder::Conditional.

Conditional(pred, true_operand, true_computation, false_operand, false_computation)

Argomenti Tipo Semantica
pred XlaOp Scalare di tipo PRED
true_operand XlaOp Argomento di tipo \(T_0\)
true_computation XlaComputation XlaComputation di tipo \(T_0 \to S\)
false_operand XlaOp Argomento di tipo \(T_1\)
false_computation XlaComputation Calcolo Xla di tipo \(T_1 \to S\)

Esegue true_computation se pred è true, false_computation se pred è false e restituisce il risultato.

true_computation deve accettare un singolo argomento di tipo \(T_0\) e verrà invocato con true_operand, che deve essere dello stesso tipo. false_computation deve accettare un singolo argomento di tipo \(T_1\) e verrà invocato con false_operand, che deve essere dello stesso tipo. Il tipo del valore restituito di true_computation e false_computation deve essere lo stesso.

Tieni presente che verrà eseguito solo uno tra true_computation e false_computation a seconda del valore di pred.

Conditional(branch_index, branch_computations, branch_operands)

Argomenti Tipo Semantica
branch_index XlaOp Scalare di tipo S32
branch_computations sequenza di N XlaComputation XlaComputations di tipo \(T_0 \to S , T_1 \to S , ..., T_{N-1} \to S\)
branch_operands sequenza di N XlaOp Argomenti di tipo \(T_0 , T_1 , ..., T_{N-1}\)

Esegue branch_computations[branch_index] e restituisce il risultato. Se branch_index è un S32 minore di 0 o maggiore o uguale a N, allora branch_computations[N-1] viene eseguito come ramo predefinito.

Ogni branch_computations[b] deve accettare un singolo argomento di tipo \(T_b\) e verrà invocato con branch_operands[b], che deve essere dello stesso tipo. Il tipo del valore restituito di ogni branch_computations[b] deve essere lo stesso.

Tieni presente che verrà eseguita una sola delle branch_computations a seconda del valore di branch_index.

Conv (convoluzione)

Vedi anche XlaBuilder::Conv.

Come ConvWithGeneralPadding, ma il padding è specificato in modo abbreviato come SAME o VALID. Il padding UGUALE aggiunge zeri all'input (lhs) in modo che l'output abbia la stessa forma dell'input se non si tiene conto del passaggio. Il valore VALID padding indica semplicemente che non è presente alcun padding.

ConvWithGeneralPadding (convoluzione)

Vedi anche XlaBuilder::ConvWithGeneralPadding.

Calcola una convezione del tipo utilizzato nelle reti neurali. Qui, una convoluzione può essere considerata come una finestra n-dimensionale che si sposta in un'area di base n-dimensionale e viene eseguito un calcolo per ogni posizione possibile della finestra.

Argomenti Tipo Semantica
lhs XlaOp array di input di rango n+2
rhs XlaOp array di rango n+2 di pesi del kernel
window_strides ArraySlice<int64> array n-d di passi del kernel
padding ArraySlice< pair<int64,int64>> Array n-dimensionale di spaziature interne (basso, alto)
lhs_dilation ArraySlice<int64> Matrice del fattore di dilatazione n-d lhs
rhs_dilation ArraySlice<int64> array del fattore di dilatazione n-d rhs
feature_group_count int64 il numero di gruppi di funzionalità
batch_group_count int64 il numero di gruppi batch

Sia n il numero di dimensioni spaziali. L'argomento lhs è un array di rango n+2 che descrive l'area di base. Questo è chiamato input, anche se ovviamente anche il lato destro è un input. In una rete neurale, queste sono le attivazioni di input. Le dimensioni n+2 sono, in questo ordine:

  • batch: ogni coordinata in questa dimensione rappresenta un input indipendente per il quale viene eseguita la convergenza.
  • z/depth/features: a ogni posizione (y, x) nell'area di base è associato un vettore che viene inserito in questa dimensione.
  • spatial_dims: descrive le dimensioni spaziali n che definiscono l'area di base su cui si sposta la finestra.

L'argomento rhs è un array di rango n+2 che descrive il filtro/nucleo/finestra di convoluzione. Le dimensioni sono, in questo ordine:

  • output-z: la dimensione z dell'output.
  • input-z: il valore di questa dimensione moltiplicato per feature_group_count deve essere uguale al valore della dimensione z in lhs.
  • spatial_dims: descrive le dimensioni spaziali n che definiscono la finestra n-d che si sposta sull'area di base.

L'argomento window_strides specifica lo stride della finestra di convolvezione nelle dimensioni spaziali. Ad esempio, se lo stride nella prima dimensione spaziale è 3, la finestra può essere posizionata solo nelle coordinate in cui il primo indice spaziale è divisibile per 3.

L'argomento padding specifica la quantità di spaziatura interna pari a zero da applicare all'area di base. La quantità di spaziatura interna può essere negativa: il valore assoluto di spaziatura interna negativa indica il numero di elementi da rimuovere dalla dimensione specificata prima di eseguire la convoluzione. padding[0] specifica il padding per la dimensione y e padding[1] specifica il padding per la dimensione x. Ogni coppia ha il padding basso come primo elemento e il padding alto come secondo elemento. La spaziatura interna bassa viene applicata nella direzione degli indici inferiori, mentre la spaziatura interna alta viene applicata nella direzione degli indici più alti. Ad esempio, se padding[1] è (2,3), nella seconda dimensione spaziale verrà inserito un riempimento di 2 zeri a sinistra e di 3 zeri a destra. L'utilizzo della spaziatura interna equivale a inserire gli stessi valori zero nell'input (lhs) prima di eseguire la convoluzione.

Gli argomenti lhs_dilation e rhs_dilation specificano il fattore di dilatazione da applicare rispettivamente a lhs e rhs in ogni dimensione spaziale. Se il coefficiente di dilatazione in una dimensione spaziale è d, tra ciascuna delle voci di quella dimensione vengono inseriti implicitamente d-1 spazi, aumentando le dimensioni dell' آرایه. I buchi vengono riempiti con un valore no-op, che per la convergenza significa zero.

La dilatazione del membro destro è chiamata anche convolvezione atrous. Per maggiori dettagli, consulta tf.nn.atrous_conv2d. La dilatazione degli lhs è anche detta convoluzione trasposta. Per maggiori dettagli, vedi tf.nn.conv2d_transpose.

L'argomento feature_group_count (valore predefinito 1) può essere utilizzato per le convoluzioni raggruppate. feature_group_count deve essere un divisore sia della dimensione della funzionalità di input sia di quella di output. Se feature_group_count è maggiore di 1, significa che concettualmente la dimensione delle funzionalità di input e di output e la dimensione delle funzionalità di output rhs sono suddivise equamente in molti gruppi feature_group_count, ciascuno costituito da una sottosezione consecutiva di funzionalità. La dimensione della funzionalità di input rhs deve essere uguale alla dimensione della funzionalità di input lhs divisa per feature_group_count (quindi ha già le dimensioni di un gruppo di funzionalità di input). I gruppi i vengono utilizzati insieme per calcolare feature_group_count per molte convoluzioni separate. I risultati di queste convolute vengono concatenati nella dimensione della funzionalità di output.

Per la convoluzione in profondità, l'argomento feature_group_count verrebbe impostato sulla dimensione della caratteristica di input e il filtro verrebbe rimodellato da [filter_height, filter_width, in_channels, channel_multiplier] a [filter_height, filter_width, 1, in_channels * channel_multiplier]. Per maggiori dettagli, vedi tf.nn.depthwise_conv2d.

L'argomento batch_group_count (valore predefinito 1) può essere utilizzato per i filtri raggruppati durante la retropropagazione. batch_group_count deve essere un divisore della dimensione del batch lhs (input). Se batch_group_count è maggiore di 1, significa che la dimensione del batch di output deve avere dimensioni input batch / batch_group_count. batch_group_count deve essere un divisore della dimensione della funzionalità di output.

La forma di output ha le seguenti dimensioni, in questo ordine:

  • batch: il valore di questa dimensione moltiplicato per batch_group_count deve essere uguale alle dimensioni della dimensione batch in lhs.
  • z: stessa dimensione di output-z nel kernel (rhs).
  • spatial_dims: un valore per ogni posizionamento valido della finestra convolzionale.

La figura sopra mostra come funziona il campo batch_group_count. In pratica, dividiamo ogni batch lhs in gruppi batch_group_count e facciamo lo stesso per le funzionalità di output. Poi, per ciascuno di questi gruppi, eseguiamo convezioni a coppie e concatenamo l'output lungo la dimensione della funzionalità di output. La semantica operativa di tutte le altre dimensioni (caratteristica e spaziale) rimane la stessa.

I posizionamenti validi della finestra convoluzionale sono determinati dalle strisce e dalle dimensioni della superficie di base dopo la spaziatura interna.

Per descrivere il funzionamento di una convoluzione, considera una convoluzione 2D e scegli alcune coordinate batch, z, y e x fisse nell'output. (y,x) rappresenta la posizione di un angolo della finestra all'interno della superficie di base (ad esempio, l'angolo in alto a sinistra, a seconda di come interpreti le dimensioni spaziali). Ora abbiamo una finestra 2D, presa dall'area di base, in cui ogni punto 2D è associato a un vettore 1D, quindi otteniamo una casella 3D. Dal kernel convoluzionale, dopo aver corretto la coordinata di output z, abbiamo anche un riquadro 3D. Le due scatole hanno le stesse dimensioni, quindi possiamo prendere la somma dei prodotti elementari tra le due scatole (in modo simile a un prodotto scalare). Questo è il valore di output.

Tieni presente che se output-z è ad esempio 5, ogni posizione della finestra produce 5 valori nell'output nella dimensione z dell'output. Questi valori differiscono in base alla parte del kernel di convoluzione utilizzata: esiste una casella 3D separata di valori utilizzata per ogni coordinata output-z. Quindi puoi considerarlo come 5 convolute separate con un filtro diverso per ciascuna.

Ecco lo pseudo-codice per una convoluzione 2D con spaziatura interna e passo:

for (b, oz, oy, ox) {  // output coordinates
  value = 0;
  for (iz, ky, kx) {  // kernel coordinates and input z
    iy = oy*stride_y + ky - pad_low_y;
    ix = ox*stride_x + kx - pad_low_x;
    if ((iy, ix) inside the base area considered without padding) {
      value += input(b, iz, iy, ix) * kernel(oz, iz, ky, kx);
    }
  }
  output(b, oz, oy, ox) = value;
}

ConvertElementType

Vedi anche XlaBuilder::ConvertElementType.

Analogamente a un static_cast elemento per elemento in C++, esegue un'operazione di conversione elemento per elemento da una forma di dati a una forma di destinazione. Le dimensioni devono essere corrispondenti e la conversione deve essere basata sugli elementi; ad esempio, gli elementi s32 diventano elementi f32 tramite una routine di conversione da s32 a f32.

ConvertElementType(operand, new_element_type)

Argomenti Tipo Semantica
operand XlaOp array di tipo T con dimensioni D
new_element_type PrimitiveType digita U

Le dimensioni dell'operando e della forma di destinazione devono corrispondere. I tipi di elementi di origine e di destinazione non devono essere tuple.

Una conversione da T=s32 a U=f32 eseguirà una routine di conversione da intero a virgola mobile come arrotondamento al più vicino alla situazione uguale.

let a: s32[3] = {0, 1, 2};
let b: f32[3] = convert(a, f32);
then b == f32[3]{0.0, 1.0, 2.0}

CrossReplicaSum

Esegue AllReduce con il calcolo della somma.

CustomCall

Vedi anche XlaBuilder::CustomCall.

Chiama una funzione fornita dall'utente all'interno di un calcolo.

CustomCall(target_name, args..., shape)

Argomenti Tipo Semantica
target_name string Nome della funzione. Verrà emessa un'istruzione di chiamata che ha come target il nome del simbolo.
args sequenza di N XlaOp N argomenti di tipo arbitrario, che verranno passati alla funzione.
shape Shape Forma di output della funzione

La firma della funzione è la stessa, indipendentemente dall'arità o dal tipo di argomenti:

extern "C" void target_name(void* out, void** in);

Ad esempio, se CustomCall viene utilizzata come segue:

let x = f32[2] {1,2};
let y = f32[2x3] { {10, 20, 30}, {40, 50, 60} };

CustomCall("myfunc", {x, y}, f32[3x3])

Ecco un esempio di implementazione di myfunc:

extern "C" void myfunc(void* out, void** in) {
  float (&x)[2] = *static_cast<float(*)[2]>(in[0]);
  float (&y)[2][3] = *static_cast<float(*)[2][3]>(in[1]);
  EXPECT_EQ(1, x[0]);
  EXPECT_EQ(2, x[1]);
  EXPECT_EQ(10, y[0][0]);
  EXPECT_EQ(20, y[0][1]);
  EXPECT_EQ(30, y[0][2]);
  EXPECT_EQ(40, y[1][0]);
  EXPECT_EQ(50, y[1][1]);
  EXPECT_EQ(60, y[1][2]);
  float (&z)[3][3] = *static_cast<float(*)[3][3]>(out);
  z[0][0] = x[1] + y[1][0];
  // ...
}

La funzione fornita dall'utente non deve avere effetti collaterali e la sua esecuzione deve essere idempotente.

Punto

Vedi anche XlaBuilder::Dot.

Dot(lhs, rhs)

Argomenti Tipo Semantica
lhs XlaOp array di tipo T
rhs XlaOp array di tipo T

La semantica esatta di questa operazione dipende dal rango degli operandi:

Input Output Semantica
vettore [n] dot vettore [n] scalare prodotto scalare vettoriale
matrice [m x k] dot vettore [k] vettore [m] moltiplicazione matrice-vettore
matrice [m x k] dot matrice [k x n] matrice [m x n] moltiplicazione matrice-matrice

L'operazione esegue la somma dei prodotti sulla seconda dimensione di lhs (o sulla prima se ha rango 1) e sulla prima dimensione di rhs. Si tratta delle dimensioni "contrattate". Le dimensioni contrattate di lhs e rhs devono essere delle stesse dimensioni. In pratica, può essere utilizzato per eseguire prodotti scalari tra vettori, moltiplicazioni vettore/matrice o moltiplicazioni matrice/matrice.

DotGeneral

Vedi anche XlaBuilder::DotGeneral.

DotGeneral(lhs, rhs, dimension_numbers)

Argomenti Tipo Semantica
lhs XlaOp array di tipo T
rhs XlaOp array di tipo T
dimension_numbers DotDimensionNumbers numeri di dimensione in batch e in contratti

Simile a Dot, ma consente di specificare i numeri delle dimensioni di aggregazione e batch sia per lhs che per rhs.

Campi DotDimensioneNumbers Tipo Semantica
lhs_contracting_dimensions int64 ripetuto lhs numeri di dimensione contrattante
rhs_contracting_dimensions int64 ripetuto rhs numeri di dimensione contrattante
lhs_batch_dimensions int64 ripetuto lhs numeri di dimensioni batch
rhs_batch_dimensions int64 ripetuto rhs numeri delle dimensioni del lotto

DotGeneral esegue la somma dei prodotti per le dimensioni di contratto specificate in dimension_numbers.

I numeri di dimensione contraente associati da lhs e rhs non devono essere uguali, ma devono avere le stesse dimensioni.

Esempio di numeri di dimensione contratta:

lhs = { {1.0, 2.0, 3.0},
{4.0, 5.0, 6.0} }

rhs = { {1.0, 1.0, 1.0},
{2.0, 2.0, 2.0} }

DotDimensionNumbers dnums;
dnums.add_lhs_contracting_dimensions(1);
dnums.add_rhs_contracting_dimensions(1);

DotGeneral(lhs, rhs, dnums) -> { {6.0, 12.0},
{15.0, 30.0} }

I numeri delle dimensioni del batch associate di lhs e rhs devono avere le stesse dimensioni.

Esempio con numeri di dimensione del lotto (dimensione del lotto 2, matrici 2x2):

lhs = { { {1.0, 2.0},
{3.0, 4.0} },
{ {5.0, 6.0},
{7.0, 8.0} } }

rhs = { { {1.0, 0.0},
{0.0, 1.0} },
{ {1.0, 0.0},
{0.0, 1.0} } }

DotDimensionNumbers dnums;
dnums.add_lhs_contracting_dimensions(2);
dnums.add_rhs_contracting_dimensions(1);
dnums.add_lhs_batch_dimensions(0);
dnums.add_rhs_batch_dimensions(0);

DotGeneral(lhs, rhs, dnums) -> { { {1.0, 2.0},
{3.0, 4.0} },
{ {5.0, 6.0},
{7.0, 8.0} } }
Input Output Semantica
[b0, m, k] dot [b0, k, n] [b0, m, n] matmul batch
[b0, b1, m, k] dot [b0, b1, k, n] [b0, b1, m, n] batch batch

Ne consegue che il numero della dimensione risultante inizia con la dimensione batch, poi con la dimensione lhs non contrattuale/non batch e infine con la dimensione rhs non contrattuale/non batch.

DynamicSlice

Vedi anche XlaBuilder::DynamicSlice.

DynamicSlice estrae un sottoarray dall'array di input in corrispondenza di start_indices dinamico. La dimensione del segmento in ogni dimensione viene passata in size_indices, che specifica il punto di fine degli intervalli di segmenti esclusivi in ogni dimensione: [start, start + size). La forma di start_indices deve avere un rango == 1, con una dimensione pari al rango di operand.

DynamicSlice(operand, start_indices, size_indices)

Argomenti Tipo Semantica
operand XlaOp Matrice dimensionale N di tipo T
start_indices sequenza di N XlaOp Elenco di N numeri interi scalari contenenti gli indici iniziali del sezionamento per ogni dimensione. Il valore deve essere maggiore o uguale a zero.
size_indices ArraySlice<int64> Elenco di N numeri interi contenenti le dimensioni della sezione per ogni dimensione. Ogni valore deve essere strettamente maggiore di zero e start + size deve essere inferiore o uguale alla dimensione della dimensione per evitare il wrapping modulo la dimensione della dimensione.

Gli indici di taglio effettivi vengono calcolati applicando la seguente trasformazione per ogni indice i in [1, N) prima di eseguire il taglio:

start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - size_indices[i])

In questo modo, il segmento estratto è sempre in ambito rispetto all'array di operandi. Se il segmento è in-bounds prima dell'applicazione della trasformazione, la trasformazione non ha alcun effetto.

Esempio unidimensionale:

let a = {0.0, 1.0, 2.0, 3.0, 4.0}
let s = {2}

DynamicSlice(a, s, {2}) produces:
{2.0, 3.0}

Esempio 2D:

let b =
{ {0.0,  1.0,  2.0},
{3.0,  4.0,  5.0},
{6.0,  7.0,  8.0},
{9.0, 10.0, 11.0} }
let s = {2, 1}

DynamicSlice(b, s, {2, 2}) produces:
{ { 7.0,  8.0},
{10.0, 11.0} }

DynamicUpdateSlice

Vedi anche XlaBuilder::DynamicUpdateSlice.

DynamicUpdateSlice genera un risultato che corrisponde al valore dell'array di input operand, con una sezione update sovrascritta in start_indices. La forma di update determina la forma del sottoarray del risultato che viene aggiornato. La forma di start_indices deve avere ranking == 1, con dimensioni della dimensione uguali al ranking di operand.

DynamicUpdateSlice(operand, update, start_indices)

Argomenti Tipo Semantica
operand XlaOp Matrice dimensionale N di tipo T
update XlaOp Array N dimensionale di tipo T contenente l'aggiornamento della frazione. Ogni dimensione della forma dell'aggiornamento deve essere strettamente maggiore di zero e inizio + aggiornamento deve essere minore o uguale alla dimensione dell'operando per ogni dimensione per evitare di generare indici di aggiornamento fuori intervallo.
start_indices sequenza di N XlaOp Elenco di N numeri interi scalari contenenti gli indici iniziali della sezione per ogni dimensione. Il valore deve essere maggiore o uguale a zero.

Gli indici di taglio effettivi vengono calcolati applicando la seguente trasformazione per ogni indice i in [1, N) prima di eseguire il taglio:

start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - update.dimension_size[i])

In questo modo, il segmento aggiornato è sempre in ambito rispetto all'array di operandi. Se il segmento è in-bounds prima dell'applicazione della trasformazione, la trasformazione non ha alcun effetto.

Esempio unidimensionale:

let a = {0.0, 1.0, 2.0, 3.0, 4.0}
let u = {5.0, 6.0}
let s = {2}

DynamicUpdateSlice(a, u, s) produces:
{0.0, 1.0, 5.0, 6.0, 4.0}

Esempio bidimensionale:

let b =
{ {0.0,  1.0,  2.0},
{3.0,  4.0,  5.0},
{6.0,  7.0,  8.0},
{9.0, 10.0, 11.0} }
let u =
{ {12.0,  13.0},
{14.0,  15.0},
{16.0,  17.0} }

let s = {1, 1}

DynamicUpdateSlice(b, u, s) produces:
{ {0.0,  1.0,  2.0},
{3.0, 12.0, 13.0},
{6.0, 14.0, 15.0},
{9.0, 16.0, 17.0} }

Operazioni aritmetiche binarie a livello di elementi

Vedi anche XlaBuilder::Add.

È supportato un insieme di operazioni aritmetiche binarie a livello di elementi.

Op(lhs, rhs)

dove Op è uno dei seguenti valori: Add (addizione), Sub (sottrazione), Mul (moltiplicazione), Div (divisione), Pow (potenza), Rem (residuo), Max (massimo), Min (minimo), And (AND logico), Or (OR logico), Xor (XOR logico), ShiftLeft (Spostamento a sinistra), ShiftRightArithmetic (Spostamento a destra aritmetico), ShiftRightLogical (Spostamento a destra logico), Atan2 (arcotangente a 2 argomenti) o Complex (combina le parti reali e imaginary in un numero complesso)

Argomenti Tipo Semantica
lhs XlaOp operando sinistro: array di tipo T
rhs XlaOp operando sul lato destro: array di tipo T

Le forme degli argomenti devono essere simili o compatibili. Consulta la documentazione relativa alla trasmissione per sapere cosa significa che le forme siano compatibili. Il risultato di un'operazione ha una forma che è il risultato della trasmissione dei due array di input. In questa variante, le operazioni tra array di livelli diversi non sono supportate, a meno che uno degli operandi non sia uno scalare.

Quando Op è Rem, il segno del risultato viene dedotto dal dividendo e il valore assoluto del risultato è sempre inferiore al valore assoluto del divisore.

Il overflow della divisione di interi (divisione/resto con segno/senza segno per zero o divisione/resto con segno di INT_SMIN con -1) produce un valore definito dall'implementazione.

Esiste una variante alternativa con supporto della trasmissione di ranking diverso per queste operazioni:

Op(lhs, rhs, broadcast_dimensions)

dove Op è lo stesso valore indicato sopra. Questa variante dell'operazione deve essere utilizzata per le operazioni aritmetiche tra array di ranghi diversi (ad esempio l'aggiunta di una matrice a un vettore).

L'operando broadcast_dimensions aggiuntivo è un segmento di numeri interi utilizzato per espandere il rango dell'operando di rango inferiore fino al rango dell'operando di rango superiore. broadcast_dimensions mappa le dimensioni della forma di rango inferiore alle dimensioni della forma di rango superiore. Le dimensioni non mappate della forma espansa vengono completate con dimensioni di dimensione 1. La trasmissione delle dimensioni degenerate quindi trasmette le forme lungo queste dimensioni degenerate per equalizzare le forme di entrambi gli operandi. La semantica è descritta in dettaglio nella pagina relativa alla trasmissione.

Operazioni di confronto elemento per elemento

Vedi anche XlaBuilder::Eq.

È supportato un insieme di operazioni di confronto binario elementari standard. Tieni presente che la semantica del confronto in virgola mobile standard IEEE 754 viene applicata quando si confrontano i tipi con virgola mobile.

Op(lhs, rhs)

Dove Op è uno di Eq (uguale a), Ne (non uguale a), Ge (maggiore o uguale a), Gt (maggiore di), Le (minore o uguale a), Lt (minore di). Un altro insieme di operatori, EqTotalOrder, NeTotalOrder, GeTotalOrder, GtTotalOrder, LeTotalOrder e LtTotalOrder, fornisce le stesse funzionalità, ma supportano in aggiunta un ordine totale rispetto ai numeri in virgola mobile, applicando -NaN < -Inf < -Finite < -0 < +0 < +Finite < +NaNInf.

Argomenti Tipo Semantica
lhs XlaOp operando sinistro: array di tipo T
rhs XlaOp operando sul lato destro: array di tipo T

Le forme degli argomenti devono essere simili o compatibili. Consulta la documentazione sulla trasmissione per scoprire cosa significa la compatibilità delle forme. Il risultato di un'operazione ha una forma che è il risultato della trasmissione dei due array di input con il tipo di elemento PRED. In questa variante, le operazioni tra array di ranghi diversi non sono supportate, a meno che uno degli operandi non sia uno scalare.

Esiste una variante alternativa con supporto della trasmissione di ranking diverso per queste operazioni:

Op(lhs, rhs, broadcast_dimensions)

dove Op è lo stesso valore indicato sopra. Questa variante dell'operazione deve essere utilizzata per le operazioni di confronto tra array di ranghi diversi (ad esempio l'aggiunta di una matrice a un vettore).

L'operando broadcast_dimensions aggiuntivo è un segmento di numeri interi che specifica le dimensioni da utilizzare per la trasmissione degli operandi. La semantica è descritta in dettaglio nella pagina relativa alla trasmissione.

Funzioni unarie elementari

XlaBuilder supporta le seguenti funzioni unarie elementari:

Abs(operand) Assoluto elemento per elemento x -> |x|.

Cbrt(operand) Operazione di radice cubica elemento per elemento x -> cbrt(x).

Ceil(operand) Ceilo elemento per elemento x -> ⌈x⌉.

Clz(operand) Conta i zeri iniziali elemento per elemento.

Cos(operand) Coseno elemento per elemento x -> cos(x).

Erf(operand) Funzione di errore elemento per elemento x -> erf(x) dove

\(\text{erf}(x) = \frac{2}{\sqrt{\pi} }\int_0^x e^{-t^2} \, dt\).

Exp(operand) Esponenziale naturale elemento per elemento x -> e^x.

Expm1(operand) Esponenziale naturale elemento per elemento meno uno x -> e^x - 1.

Floor(operand) Valore minimo per elemento x -> ⌊x⌋.

Imag(operand) Parte immaginaria basata sugli elementi di una forma complessa (o reale). x -> imag(x). Se l'operando è di tipo a virgola mobile, restituisce 0.

IsFinite(operand) Verifica se ogni elemento di operand è finito, ovvero non è infinito positivo o negativo e non è NaN. Restituisce un array di valori PRED con la stessa forma dell'input, dove ogni elemento è true se e solo se l'elemento di input corrispondente è finito.

Log(operand) Logaritmo naturale elemento per elemento x -> ln(x).

Log1p(operand) Logaritmo naturale spostato per elemento x -> ln(1+x).

Logistic(operand) Calcolo della funzione logistica elemento per elemento x -> logistic(x).

Neg(operand) Negazione elemento per elemento x -> -x.

Not(operand) Logica a livello di elemento non x -> !(x).

PopulationCount(operand) Calcola il numero di bit impostati in ogni elemento di operand.

Real(operand) Parte reale elemento per elemento di una forma complessa (o reale). x -> real(x). Se l'operando è di tipo a virgola mobile, restituisce lo stesso valore.

Round(operand) Arrotondamento elemento per elemento, i valori uguali a zero vengono ignorati.

RoundNearestEven(operand) Arrotondamento elemento per elemento, con arrotondamento al numero pari più vicino.

Rsqrt(operand) Reciproco elemento per elemento dell'operazione di radice quadrata x -> 1.0 / sqrt(x).

Sign(operand) Operazione di firma a livello di elemento x -> sgn(x) in cui

\[\text{sgn}(x) = \begin{cases} -1 & x < 0\\ -0 & x = -0\\ NaN & x = NaN\\ +0 & x = +0\\ 1 & x > 0 \end{cases}\]

utilizzando l'operatore di confronto del tipo di elemento operand.

Sin(operand) Seno elemento per elemento x -> sin(x).

Sqrt(operand) Operazione di radice quadrata elemento per elemento x -> sqrt(x).

Tan(operand) Tangente x -> tan(x) rispetto all'elemento.

Tanh(operand) Tangente iperbolica basata sull'elemento x -> tanh(x).

Argomenti Tipo Semantica
operand XlaOp L'operando alla funzione

La funzione viene applicata a ciascun elemento nell'array operand, generando un array con la stessa forma. operand può essere uno scalare (ranking 0) consentito.

Fft

L'operazione FFT XLA implementa le trasformazioni di Fourier dirette e inverse per input/output reali e complessi. Sono supportati gli FFT multidimensionali su un massimo di 3 assi.

Vedi anche XlaBuilder::Fft.

Argomenti Tipo Semantica
operand XlaOp L'array di cui stiamo calcolando la trasformata di Fourier.
fft_type FftType Consulta la tabella riportata di seguito.
fft_length ArraySlice<int64> Le lunghezze nel dominio del tempo degli assi in fase di trasformazione. Questo è necessario in particolare per l'IRFFT per ridimensionare correttamente l'asse più interno, poiché RFFT(fft_length=[16]) ha la stessa forma di output di RFFT(fft_length=[17]).
FftType Semantica
FFT Inoltra FFT da complessi a complessi. La forma non cambia.
IFFT FFT inverso da complesso a complesso. La forma non cambia.
RFFT Inoltra FFT da reale a complesso. La forma dell'asse più interno viene ridotta a fft_length[-1] // 2 + 1 se fft_length[-1] è un valore diverso da zero, omettendo la parte coniugata inversa del segnale trasformato oltre la frequenza di Nyquist.
IRFFT FFT inverso da reale a complesso (ovvero richiede un insieme di dati complessi e restituisce il valore reale). La forma dell'asse più interno viene espansa a fft_length[-1] se fft_length[-1] è un valore diverso da zero, deducendo la parte del segnale trasformato oltre la frequenza di Nyquist dalla coniugata inversa delle voci 1 a fft_length[-1] // 2 + 1.

FFT multidimensionale

Se viene fornito più di un fft_length, è equivalente all'applicazione di una cascata di operazioni FFT a ciascuno degli assi più interni. Tieni presente che per i casi reale->complesso e complesso->reale, la trasformazione dell'asse più interno viene eseguita (in modo efficace) per prima (RFFT; per ultima per IRFFT), motivo per cui è l'asse più interno a cambiare dimensione. Le altre trasformazioni degli assi saranno quindi complesso->complesso.

Dettagli di implementazione

L'FFT della CPU è supportato da TensorFFT di Eigen. FFT GPU utilizza cuFFT.

Raccogliere

La XLA riunisce le operazioni di stitching di diverse sezioni (ciascuna sezione con un offset di runtime potenzialmente diverso) di un array di input.

Semantica generale

Vedi anche XlaBuilder::Gather. Per una descrizione più intuitiva, consulta la sezione "Descrizione informale" di seguito.

gather(operand, start_indices, offset_dims, collapsed_slice_dims, slice_sizes, start_index_map)

Argomenti Tipo Semantica
operand XlaOp L'array da cui stiamo raccogliendo i dati.
start_indices XlaOp Array contenente gli indici iniziali delle sezioni raccolte.
index_vector_dim int64 La dimensione in start_indices che "contiene" gli indici iniziali. Di seguito è riportata una descrizione dettagliata.
offset_dims ArraySlice<int64> L'insieme di dimensioni nella forma di output che eseguono l'offset in una matrice suddivisa in base all'operando.
slice_sizes ArraySlice<int64> slice_sizes[i] sono i limiti della sezione nella dimensione i.
collapsed_slice_dims ArraySlice<int64> L'insieme di dimensioni in ogni frazione che vengono compresse. Queste dimensioni devono avere dimensioni pari a 1.
start_index_map ArraySlice<int64> Una mappa che descrive come mappare gli indici in start_indices agli indici legali nell'operando.
indices_are_sorted bool Indica se è garantito che gli indici siano ordinati dal chiamante.

Per praticità, etichettiamo le dimensioni nell'array di output non in offset_dims come batch_dims.

L'output è un array di ranking batch_dims.size + offset_dims.size.

operand.rank deve essere uguale alla somma di offset_dims.size e collapsed_slice_dims.size. Inoltre, slice_sizes.size deve essere uguale a operand.rank.

Se index_vector_dim è uguale a start_indices.rank, consideriamo implicitamente start_indices una dimensione 1 finale (ad es. se start_indices era di forma [6,7] e index_vector_dim è 2, consideriamo implicitamente la forma di start_indices come [6,7,1]).

I limiti per l'array di output lungo la dimensione i vengono calcolati come segue:

  1. Se i è presente in batch_dims (ovvero è uguale a batch_dims[k] per alcuni k), scegliamo i limiti delle dimensioni corrispondenti tra start_indices.shape, saltando index_vector_dim (ad es. scegli start_indices.shape.dims[k] se k < index_vector_dim e start_indices.shape.dims[k+1] in caso contrario).

  2. Se i è presente in offset_dims (ovvero è uguale a offset_dims[k] per qualche k), selezioniamo il limite corrispondente da slice_sizes dopo aver tenuto conto di collapsed_slice_dims (ovvero selezioniamo adjusted_slice_sizes[k] dove adjusted_slice_sizes è slice_sizes con i limiti agli indici collapsed_slice_dims rimossi).

Formalmente, l'indice dell'operando In corrispondente a un determinato indice di output Out viene calcolato come segue:

  1. Sia G = { Out[k] per k in batch_dims }. Utilizza G per suddividere un vettore S in modo tale che S[i] = start_indices[Combina(G, i)] dove Combina(A, b) inserisce b nella posizione index_vector_dim in A. Tieni presente che questo è ben definito anche se G è vuoto: se G è vuoto, S = start_indices.

  2. Crea un indice iniziale, Sin, in operand utilizzando S distribuendo S tramite start_index_map. Più precisamente:

    1. Sin[start_index_map[k]] = S[k] se k < start_index_map.size.

    2. Sin[_] = 0 in caso contrario.

  3. Crea un indice Oin in operand distribuendo gli indici nelle dimensioni di offset in Out in base all'insieme collapsed_slice_dims. Più nello specifico:

    1. Oin[remapped_offset_dims(k)] = Out[offset_dims[k]] se k < offset_dims.size (remapped_offset_dims è definito di seguito).

    2. Oin[_] = 0 in caso contrario.

  4. In è Oin + Sin, dove + rappresenta l'aggiunta di elementi.

remapped_offset_dims è una funzione monotona con dominio [0, offset_dims.size) e intervallo [0, operand.rank) \ collapsed_slice_dims. Pertanto, se ad esempio offset_dims.size è 4, operand.rank è 6 e collapsed_slice_dims è {0, 2} quindi remapped_offset_dims è {01, 13, 24, 35}.

Se indices_are_sorted è impostato su true, XLA può assumere che start_indices siano ordinati (in ordine crescente, dopo aver sparso i relativi valori in base a start_index_map) dall'utente. In caso contrario, la semantica viene definita per l'implementazione.

Descrizione informale ed esempi

In modo informale, ogni indice Out nell'array di output corrisponde a un elemento E nell'array dell'operando, calcolato come segue:

  • Utilizziamo le dimensioni batch in Out per cercare un indice iniziale a partire da start_indices.

  • Utilizziamo start_index_map per mappare l'indice iniziale (la cui dimensione può essere inferiore a operand.rank) a un indice iniziale "completo" in operand.

  • Dividiamo dinamicamente un segmento con dimensioni slice_sizes utilizzando l'indice iniziale completo.

  • Rimodelniamo la sezione comprimindo le dimensioni collapsed_slice_dims. Poiché tutte le dimensioni degli slice compressi devono avere un limite di 1, questa trasformazione è sempre valida.

  • Utilizziamo le dimensioni di offset in Out per indicizzare questa sezione in modo da ottenere l'elemento di input, E, corrispondente all'indice di output Out.

Il valore index_vector_dim è impostato su start_indices.rank - 1 in tutti gli esempi che seguono. I valori più interessanti per index_vector_dim non modificano sostanzialmente l'operazione, ma rendono la rappresentazione visiva più complicata.

Per capire meglio come si combinano tutti gli elementi precedenti, diamo un'occhiata a un esempio che raccoglie 5 slice di forma [8,6] da un array [16,11]. La posizione di un segmento nell'array [16,11] può essere rappresentata come un vettore di indice di forma S64[2], pertanto l'insieme di 5 posizioni può essere rappresentato come un S64[5,2] array.

Il comportamento dell'operazione di collect può quindi essere rappresentato come una trasformazione dell'indice che prende [G,O0,O1], un indice nella forma di output, e lo mappa a un elemento nell'array di input nel seguente modo:

Per prima cosa selezioniamo un vettore (X,Y) dall'array di indici di raccolta utilizzando G. L'elemento nell'array di output all'indice [G,O0,O1] è quindi l'elemento nell'array di input all'indice [X+O0,Y+O1].

slice_sizes è [8,6], che stabilisce l'intervallo di O0 e O1 e questo determina a sua volta i limiti della sezione.

Questa operazione di raccolta agisce come una sezione dinamica batch con G come dimensione batch.

Gli indici di aggregazione possono essere multidimensionali. Ad esempio, una versione più generale dell'esempio precedente che utilizza un array di "indici di aggregazione" di forma [4,5,2] tradurrebbe gli indici come segue:

Anche in questo caso, si tratta di un segmento dinamico del batch G0 e G1 come dimensioni del batch. La dimensione del segmento rimane [8,6].

L'operazione di aggregazione in XLA generalizza la semantica informale descritta sopra nel seguente modo:

  1. Possiamo configurare le dimensioni della forma di output che sono le dimensioni di offset (le dimensioni contenenti O0, O1 nell'ultimo esempio). Le dimensioni del batch di output (le dimensioni contenenti G0, G1 nell'ultimo esempio) sono definite come le dimensioni di output che non sono dimensioni di offset.

  2. Il numero di dimensioni di offset dell'output esplicitamente presenti nella forma di output potrebbe essere inferiore al rango dell'input. Queste dimensioni "mancanti", elencate esplicitamente come collapsed_slice_dims, devono avere una dimensione del segmento di 1. Poiché hanno una dimensione di slice pari a 1, l'unico indice valido per loro è 0 e l'elisione non introduce ambiguità.

  3. Il segmento estratto dall'array "Indici di aggregazione" ((X, Y) nell'ultimo esempio) può avere meno elementi rispetto al ranking dell'array di input e una mappatura esplicita determina come deve essere espanso l'indice per avere lo stesso ranking dell'input.

Come ultimo esempio, utilizziamo i numeri (2) e (3) per implementare tf.gather_nd:

G0 e G1 vengono utilizzati per estrarre un indice iniziale dall'array di indici di aggregazione come di consueto, tranne per il fatto che l'indice iniziale ha un solo elemento, X. Analogamente, esiste un solo indice di offset di output con il valore O0. start_index_mapremapped_offset_dimsXX000000000OOOOGGGG11GatherIndicestf.gather_nd

slice_sizes per questa richiesta è [1,11]. Intuitivamente, questo significa che ogni indice X nell'array di indici di aggregazione seleziona un'intera riga e il risultato è la concatenazione di tutte queste righe.

GetDimensionSize

Vedi anche XlaBuilder::GetDimensionSize.

Restituisce la dimensione della dimensione specificata dell'operando. L'operando deve avere la forma di un array.

GetDimensionSize(operand, dimension)

Argomenti Tipo Semantica
operand XlaOp Array di input di n dimensioni
dimension int64 Un valore nell'intervallo [0, n) che specifica la dimensione

SetDimensionSize

Vedi anche XlaBuilder::SetDimensionSize.

Imposta la dimensione dinamica della dimensione specificata di XlaOp. L'operando deve avere la forma di un array.

SetDimensionSize(operand, size, dimension)

Argomenti Tipo Semantica
operand XlaOp Array di input n-dimensionale.
size XlaOp int32 che rappresenta la dimensione dinamica di runtime.
dimension int64 Un valore nell'intervallo [0, n) che specifica la dimensione.

Passa attraverso l'operando come risultato, con la dimensione dinamica monitorata dal compilatore.

I valori con spaziatura interna verranno ignorati dalle operazioni di riduzione a valle.

let v: f32[10] = f32[10]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
let five: s32 = 5;
let six: s32 = 6;

// Setting dynamic dimension size doesn't change the upper bound of the static
// shape.
let padded_v_five: f32[10] = set_dimension_size(v, five, /*dimension=*/0);
let padded_v_six: f32[10] = set_dimension_size(v, six, /*dimension=*/0);

// sum == 1 + 2 + 3 + 4 + 5
let sum:f32[] = reduce_sum(padded_v_five);
// product == 1 * 2 * 3 * 4 * 5
let product:f32[] = reduce_product(padded_v_five);

// Changing padding size will yield different result.
// sum == 1 + 2 + 3 + 4 + 5 + 6
let sum:f32[] = reduce_sum(padded_v_six);

GetTupleElement

Vedi anche XlaBuilder::GetTupleElement.

Indici in una tupla con un valore costante di tempo di compilazione.

Il valore deve essere una costante di compilazione affinché l'inferenza della forma possa determinare il tipo del valore risultante.

È analogo a std::get<int N>(t) in C++. Concettualmente:

let v: f32[10] = f32[10]{0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
let s: s32 = 5;
let t: (f32[10], s32) = tuple(v, s);
let element_1: s32 = gettupleelement(t, 1);  // Inferred shape matches s32.

Vedi anche tf.tuple.

Infeed

Vedi anche XlaBuilder::Infeed.

Infeed(shape)

Argomento Tipo Semantica
shape Shape Forma dei dati letti dall'interfaccia Infeed. Il campo del layout della forma deve essere impostato in modo da corrispondere al layout dei dati inviati al dispositivo; in caso contrario, il comportamento non è definito.

Legge un singolo elemento di dati dall'interfaccia di streaming infeed implicita del dispositivo, interpretando i dati come la forma e il layout specificati e restituisce un XlaOp di dati. In un calcolo sono consentite più operazioni infeed, ma deve essere presente un ordine totale tra le operazioni infeed. Ad esempio, due infeed nel codice di seguito hanno un ordine totale poiché esiste una dipendenza tra i loop while.

result1 = while (condition, init = init_value) {
  Infeed(shape)
}

result2 = while (condition, init = result1) {
  Infeed(shape)
}

Le forme di tuple nidificate non sono supportate. Per una forma di tupla vuota, l'operazione Infeed è in pratica un'operazione no-op e procede senza leggere alcun dato dall'infeed del dispositivo.

Iota

Vedi anche XlaBuilder::Iota.

Iota(shape, iota_dimension)

Consente di creare una costante letterale sul dispositivo anziché un trasferimento dell'host potenzialmente di grandi dimensioni. Crea un array con una forma specificata e contiene valori che iniziano da zero e aumentano di uno lungo la dimensione specificata. Per i tipi con virgola mobile, l'array generato è equivalente a ConvertElementType(Iota(...)), dove Iota è di tipo integrale e la conversione è di tipo in virgola mobile.

Argomenti Tipo Semantica
shape Shape Forma dell'array creata da Iota()
iota_dimension int64 La dimensione da incrementare.

Ad esempio, Iota(s32[4, 8], 0) restituisce

  [[0, 0, 0, 0, 0, 0, 0, 0 ],
   [1, 1, 1, 1, 1, 1, 1, 1 ],
   [2, 2, 2, 2, 2, 2, 2, 2 ],
   [3, 3, 3, 3, 3, 3, 3, 3 ]]

Resi a Iota(s32[4, 8], 1)

  [[0, 1, 2, 3, 4, 5, 6, 7 ],
   [0, 1, 2, 3, 4, 5, 6, 7 ],
   [0, 1, 2, 3, 4, 5, 6, 7 ],
   [0, 1, 2, 3, 4, 5, 6, 7 ]]

Mappa

Vedi anche XlaBuilder::Map.

Map(operands..., computation)

Argomenti Tipo Semantica
operands sequenza di N XlaOp s N array di tipo T0..T{N-1}
computation XlaComputation calcolo di tipo T_0, T_1, .., T_{N + M -1} -> S con N parametri di tipo T e M di tipo arbitrario
dimensions Array int64 array di dimensioni della mappa

Applica una funzione scalare sugli array operands specificati, producendo un array delle stesse dimensioni in cui ogni elemento è il risultato della funzione mappata applicata agli elementi corrispondenti negli array di input.

La funzione mappata è un calcolo arbitrario con la limitazione di avere N input di tipo scalare T e un singolo output di tipo S. L'output ha le stesse dimensioni degli operandi, tranne per il fatto che il tipo di elemento T viene sostituito con S.

Ad esempio: Map(op1, op2, op3, computation, par1) mappa elem_out <- computation(elem1, elem2, elem3, par1) a ciascun indice (multidimensionale) negli array di input per produrre l'array di output.

OptimizationBarrier

Impedisce a qualsiasi passaggio di ottimizzazione di spostare i calcoli oltre la barriera.

Garantisce che tutti gli input vengano valutati prima di qualsiasi operatore che dipende dagli output della barriera.

Cuscinetto

Vedi anche XlaBuilder::Pad.

Pad(operand, padding_value, padding_config)

Argomenti Tipo Semantica
operand XlaOp array di tipo T
padding_value XlaOp scalare di tipo T per riempire la spaziatura interna aggiunta
padding_config PaddingConfig La quantità di spaziatura interna su entrambi i lati (basso, alto) e tra gli elementi di ogni dimensione

Espande l'array operand specificato inserendo spaziatura interna intorno all'array e tra gli elementi dell'array con il valore padding_value specificato. padding_config specifica la quantità di spaziatura interna e interna per ogni dimensione.

PaddingConfig è un campo ripetuto di PaddingConfigDimension, che contiene tre campi per ogni dimensione: edge_padding_low, edge_padding_high e interior_padding.

edge_padding_low e edge_padding_high specificano la quantità di spaziatura interna aggiunta rispettivamente all'estremità inferiore (accanto all'indice 0) e all'estremità superiore (accanto all'indice più alto) di ogni dimensione. La quantità di spaziatura interna può essere negativa: il valore assoluto della spaziatura interna negativa indica il numero di elementi da rimuovere dalla dimensione specificata.

interior_padding specifica la quantità di spaziatura interna aggiunta tra due elementi in ogni dimensione. Non può essere negativo. Il padding interno avviene logicamente prima del padding dei bordi, pertanto, nel caso di padding dei bordi negativo, gli elementi vengono rimossi dall'operando con padding interno.

Questa operazione non esegue alcuna operazione se tutte le coppie di spaziatura interna sono (0, 0) e tutti i valori di spaziatura interna sono 0. La figura seguente mostra esempi di valori diversi di edge_padding e interior_padding per un array bidimensionale.

Recv

Vedi anche XlaBuilder::Recv.

Recv(shape, channel_handle)

Argomenti Tipo Semantica
shape Shape la forma dei dati per ricevere
channel_handle ChannelHandle identificatore univoco per ogni coppia di invio/ricezione

Riceve i dati della forma specificata da un'istruzione Send in un altro calcolo che condivide lo stesso handle di canale. Restituisce un XlaOp per i dati ricevuti.

L'API client dell'operazione Recv rappresenta la comunicazione sincrona. Tuttavia, l'istruzione viene decomposta internamente in due istruzioni HLO (Recv e RecvDone) per abilitare i trasferimenti di dati asincroni. Vedi anche HloInstruction::CreateRecv e HloInstruction::CreateRecvDone.

Recv(const Shape& shape, int64 channel_id)

Alloca le risorse necessarie per ricevere i dati da un'istruzione Send con lo stesso channel_id. Restituisce un contesto per le risorse allocate, che viene utilizzato da una successiva istruzione RecvDone per attendere il completamento del trasferimento dei dati. Il contesto è una tupla di {receive buffer (shape), request identifier (U32)} e può essere utilizzato solo da un'istruzione RecvDone.

RecvDone(HloInstruction context)

Dato un contesto creato da un'istruzione Recv, attende il completamento del trasferimento dei dati e restituisce i dati ricevuti.

Riduci

Vedi anche XlaBuilder::Reduce.

Applica una funzione di riduzione a uno o più array in parallelo.

Reduce(operands..., init_values..., computation, dimensions)

Argomenti Tipo Semantica
operands Sequenza di N XlaOp N array di tipo T_0, ..., T_{N-1}.
init_values Sequenza di N XlaOp N scalari di tipi T_0, ..., T_{N-1}.
computation XlaComputation calcolo di tipo T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}).
dimensions Array int64 Array non ordinato di dimensioni da ridurre.

Dove:

  • N deve essere maggiore o uguale a 1.
  • Il calcolo deve essere "approssimativamente" associativo (vedi sotto).
  • Tutti gli array di input devono avere le stesse dimensioni.
  • Tutti i valori iniziali devono formare un'identità in computation.
  • Se N = 1, Collate(T) è T.
  • Se N > 1, Collate(T_0, ..., T_{N-1}) è una tupla di elementi N di tipo T.

Questa operazione riduce una o più dimensioni di ogni array di input in scalari. Il ranking di ogni matrice restituita è rank(operand) - len(dimensions). L'output dell'operazione è Collate(Q_0, ..., Q_N), dove Q_i è un array di tipo T_i, le cui dimensioni sono descritte di seguito.

È consentito associare nuovamente il calcolo della riduzione a backend diversi. Ciò può portare a differenze numeriche, poiché alcune funzioni di riduzione come l'addizione non sono associative per i valori float. Tuttavia, se l'intervallo dei dati è limitato, l'addizione a virgola mobile è sufficientemente simile all'associazione per la maggior parte degli utilizzi pratici.

Esempi

Quando si riduce una dimensione in un singolo array 1D con valori [10, 11, 12, 13], con la funzione di riduzione f (ovvero computation), il valore può essere calcolato come

f(10, f(11, f(12, f(init_value, 13)))

ma ci sono anche molte altre possibilità, ad es.

f(init_value, f(f(10, f(init_value, 11)), f(f(init_value, 12), f(init_value, 13))))

Di seguito è riportato un esempio approssimativo di pseudo-codice di come potrebbe essere implementata la riduzione, utilizzando la somma come calcolo della riduzione con un valore iniziale di 0.

result_shape <- remove all dims in dimensions from operand_shape

# Iterate over all elements in result_shape. The number of r's here is equal
# to the rank of the result
for r0 in range(result_shape[0]), r1 in range(result_shape[1]), ...:
  # Initialize this result element
  result[r0, r1...] <- 0

  # Iterate over all the reduction dimensions
  for d0 in range(dimensions[0]), d1 in range(dimensions[1]), ...:
    # Increment the result element with the value of the operand's element.
    # The index of the operand's element is constructed from all ri's and di's
    # in the right order (by construction ri's and di's together index over the
    # whole operand shape).
    result[r0, r1...] += operand[ri... di]

Ecco un esempio di riduzione di un array 2D (matrice). La forma ha rango 2, dimensione 0 di dimensione 2 e dimensione 1 di dimensione 3:

Risultati della riduzione delle dimensioni 0 o 1 con una funzione "add":

Tieni presente che entrambi i risultati della riduzione sono array 1D. Il diagramma mostra una colonna e un'altra come riga per comodità visiva.

Per un esempio più complesso, ecco un array 3D. Il suo rango è 3, la dimensione 0 di dimensione 4, la dimensione 1 di dimensione 2 e la dimensione 2 di dimensione 3. Per semplicità, i valori da 1 a 6 vengono replicati nella dimensione 0.

Analogamente all'esempio 2D, possiamo ridurre una sola dimensione. Se riduciamo la dimensione 0, ad esempio, otteniamo un array di rango 2 in cui tutti i valori della dimensione 0 sono stati combinati in un valore scalare:

|  4   8  12 |
| 16  20  24 |

Se riduciamo la dimensione 2, otteniamo anche un array di rango 2 in cui tutti i valori della dimensione 2 sono stati ridotti in uno scalare:

| 6  15 |
| 6  15 |
| 6  15 |
| 6  15 |

Tieni presente che l'ordine relativo tra le dimensioni rimanenti nell'input viene mantenuto nell'output, ma ad alcune dimensioni potrebbero essere assegnati nuovi numeri (poiché il ranking cambia).

Possiamo anche ridurre più dimensioni. Le dimensioni addizionali 0 e 1 producono l'array 1D [20, 28, 36].

Riducendo l'array 3D su tutte le sue dimensioni si ottiene lo 84 scalare.

Riduzione variabile

Quando N > 1, l'applicazione della funzione di riduzione è leggermente più complessa, in quanto viene applicata contemporaneamente a tutti gli input. Gli operandi vengono forniti al calcolo nel seguente ordine:

  • Valore ridotto in esecuzione per il primo operando
  • Valore ridotto dell'operando n-esimo
  • Valore di input per il primo operando
  • Valore di input per l'operando n-esimo

Ad esempio, prendi in considerazione la seguente funzione di riduzione, che può essere utilizzata per calcolare in parallelo il valore massimo e l'argmax di una matrice 1D:

f: (Float, Int, Float, Int) -> Float, Int
f(max, argmax, value, index):
  if value >= max:
    return (value, index)
  else:
    return (max, argmax)

Per gli array di input 1-D V = Float[N], K = Int[N] e i valori di inizializzazione I_V = Float, I_K = Int, il risultato f_(N-1) della riduzione nell'unica dimensione di input è equivalente alla seguente applicazione ricorsiva:

f_0 = f(I_V, I_K, V_0, K_0)
f_1 = f(f_0.first, f_0.second, V_1, K_1)
...
f_(N-1) = f(f_(N-2).first, f_(N-2).second, V_(N-1), K_(N-1))

Applicando questa riduzione a un array di valori e a un array di indici sequenziali (ovvero iota), viene eseguita la co-iterazione sugli array e viene restituita una tupla contenente il valore massimo e l'indice corrispondente.

RiduciPrecisione

Vedi anche XlaBuilder::ReducePrecision.

Modella l'effetto della conversione dei valori in virgola mobile in un formato di minore precisione (ad esempio IEEE-FP16) e di nuovo nel formato originale. Il numero di bit dell'esponente e della mantissa nel formato a precisione inferiore può essere specificato arbitrariamente, anche se tutte le dimensioni in bit potrebbero non essere supportate in tutte le implementazioni hardware.

ReducePrecision(operand, mantissa_bits, exponent_bits)

Argomenti Tipo Semantica
operand XlaOp array di tipo a virgola mobile T.
exponent_bits int32 numero di bit dell'esponente in formato a precisione inferiore
mantissa_bits int32 numero di bit della frazione in formato a precisione inferiore

Il risultato è un array di tipo T. I valori di input vengono arrotondati al valore representabile più vicino con il numero specificato di bit della frazione (utilizzando la semantica "pari") e tutti i valori che superano l'intervallo specificato dal numero di bit dell'esponente vengono vincolati a infinito positivo o negativo. I valori NaN vengono conservati, anche se possono essere convertiti in valori NaN canonici.

Il formato a precisione inferiore deve avere almeno un bit di esponente (per distinguere un valore zero da un infinito, poiché entrambi hanno una mantissa pari a zero) e deve avere un numero non negativo di bit di mantissa. Il numero di bit dell'esponente o della mantissa può superare il valore corrispondente per il tipo T; la parte corrispondente della conversione è quindi semplicemente un'operazione senza effetti.

ReduceScatter

Vedi anche XlaBuilder::ReduceScatter.

ReduceScatter è un'operazione collettiva che esegue efficacemente un AllReduce e poi sparge il risultato dividendolo in blocchi shard_count lungo scatter_dimension e la replica i nel gruppo di repliche riceve lo shard ith.

ReduceScatter(operand, computation, scatter_dim, shard_count, replica_group_ids, channel_id)

Argomenti Tipo Semantica
operand XlaOp Matrice o una tupla di matrici non vuota da ridurre nelle repliche.
computation XlaComputation Calcolo della riduzione
scatter_dimension int64 La dimensione da visualizzare in un grafico a dispersione.
shard_count int64 Numero di blocchi da suddividere scatter_dimension
replica_groups vettore di vettori di int64 Gruppi tra i quali vengono eseguite le riduzioni
channel_id int64 facoltativo ID canale facoltativo per le comunicazioni tra moduli
  • Quando operand è una tupla di array, la riduzione-sparsa viene eseguita su ogni elemento della tupla.
  • replica_groups è un elenco di gruppi di repliche tra cui viene eseguita la riduzione (l'ID replica per la replica corrente può essere recuperato utilizzando ReplicaId). L'ordine delle repliche in ogni gruppo determina l'ordine in cui verrà distribuito il risultato di all-reduce. replica_groups deve essere vuoto (in questo caso tutte le repliche appartengono a un singolo gruppo) o contenere lo stesso numero di elementi delle repliche. Se esistono più gruppi di repliche, devono avere tutti le stesse dimensioni. Ad esempio, replica_groups = {0, 2}, {1, 3} esegue la riduzione tra le repliche 0 e 2 e 1 e 3 e poi sparge il risultato.
  • shard_count è la dimensione di ogni gruppo di repliche. Ne abbiamo bisogno nei casi in cui replica_groups sono vuoti. Se replica_groups non è vuoto, shard_count deve essere uguale alle dimensioni di ciascun gruppo di repliche.
  • channel_id viene utilizzato per la comunicazione tra moduli: solo le operazioni di reduce-scatter con lo stesso channel_id possono comunicare tra loro.

La forma di output è la forma di input con scatter_dimension ridotta shard_count volte. Ad esempio, se ci sono due repliche e l'operando ha rispettivamente i valori [1.0, 2.25] e [3.0, 5.25] nelle due repliche, il valore di output di questa operazione in cui scatter_dim è 0 sarà [4.0] per la prima replica e [7.5] per la seconda replica.

ReduceWindow

Vedi anche XlaBuilder::ReduceWindow.

Applica una funzione di riduzione a tutti gli elementi di ogni finestra di una sequenza di N array multi-dimensionali, producendo come output un singolo array o una tupla di N array multi-dimensionali. Ogni array di output ha lo stesso numero di elementi del numero di posizioni valide della finestra. Un livello di pooling può essere espresso come ReduceWindow. Come per Reduce, il valore computation applicato viene sempre passato a init_values sul lato sinistro.

ReduceWindow(operands..., init_values..., computation, window_dimensions, window_strides, padding)

Argomenti Tipo Semantica
operands N XlaOps Una sequenza di N array multidimensionali di tipi T_0,..., T_{N-1}, ciascuno dei quali rappresenta l'area di base su cui è posizionata la finestra.
init_values N XlaOps Gli N valori iniziali per la riduzione, uno per ciascuno degli N operandi. Per maggiori dettagli, vedi Ridurre.
computation XlaComputation Funzione di riduzione di tipo T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}), da applicare agli elementi in ogni finestra di tutti gli operandi di input.
window_dimensions ArraySlice<int64> array di numeri interi per i valori delle dimensioni della finestra
window_strides ArraySlice<int64> array di numeri interi per i valori del passo della finestra
base_dilations ArraySlice<int64> array di interi per i valori di dilatazione di base
window_dilations ArraySlice<int64> array di numeri interi per i valori di dilatazione della finestra
padding Padding tipo di spaziatura interna per la finestra (Padding::kSame, che pad in modo da avere la stessa forma di output dell'input se il passo è 1, o Padding::kValid, che non utilizza spaziatura interna e "arresta" la finestra quando non si adatta più)

Dove:

  • N deve essere maggiore o uguale a 1.
  • Tutti gli array di input devono avere le stesse dimensioni.
  • Se N = 1, Collate(T) è T.
  • Se N > 1, Collate(T_0, ..., T_{N-1}) è una tupla di elementi N di tipo (T0,...T{N-1}).

Di seguito il codice e la figura mostrano un esempio di utilizzo di ReduceWindow. L'input è una matrica di dimensioni [4x6] e sia window_dimensions sia window_stride_dimensions sono [2x3].

// Create a computation for the reduction (maximum).
XlaComputation max;
{
  XlaBuilder builder(client_, "max");
  auto y = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "y");
  auto x = builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "x");
  builder.Max(y, x);
  max = builder.Build().value();
}

// Create a ReduceWindow computation with the max reduction computation.
XlaBuilder builder(client_, "reduce_window_2x3");
auto shape = ShapeUtil::MakeShape(F32, {4, 6});
auto input = builder.Parameter(0, shape, "input");
builder.ReduceWindow(
    input,
    /*init_val=*/builder.ConstantLiteral(LiteralUtil::MinValue(F32)),
    *max,
    /*window_dimensions=*/{2, 3},
    /*window_stride_dimensions=*/{2, 3},
    Padding::kValid);

L'intervallo di 1 in una dimensione specifica che la posizione di una finestra nella dimensione è a 1 elemento di distanza dalla finestra adiacente. Per specificare che nessuna finestra si sovrapponga tra loro, window_stride_dimensions deve essere uguale a window_dimensions. La figura seguente illustra l'uso di due valori di passo diversi. Il riempimento viene applicato a ogni dimensione dell'input e i calcoli sono gli stessi come se l'input fosse stato inserito con le dimensioni che ha dopo il riempimento.

Per un esempio di riempimento non banale, considera il calcolo del minimo della finestra di riduzione (valore iniziale MAX_FLOAT) con la dimensione 3 e lo stride 2 sull'array di input[10000, 1000, 100, 10, 1]. Il riempimento di kValid calcola il minimo su due finestre valide: [10000, 1000, 100] e [100, 10, 1], generando un output [100, 1]. Se si completa kSame, l'array viene riempito per prima cosa in modo che la forma dopo la finestra di riduzione sia la stessa dell'input per la fase uno, aggiungendo elementi iniziali su entrambi i lati, ottenendo [MAX_VALUE, 10000, 1000, 100, 10, 1, MAX_VALUE]. L'esecuzione di reduce-window sull'array con spaziatura aggiuntiva viene eseguita su tre finestre [MAX_VALUE, 10000, 1000], [1000, 100, 10], [10, 1, MAX_VALUE] e genera [1000, 10, 1].

L'ordine di valutazione della funzione di riduzione è arbitrario e può essere non deterministico. Di conseguenza, la funzione di riduzione non deve essere eccessivamente sensibile alla riassociazione. Per ulteriori dettagli, consulta la discussione sull'associatività nel contesto di Reduce.

ReplicaId

Vedi anche XlaBuilder::ReplicaId.

Restituisce l'ID univoco (scalare U32) della replica.

ReplicaId()

L'ID univoco di ogni replica è un numero intero non firmato nell'intervallo [0, N), dove N è il numero di repliche. Poiché tutte le repliche eseguono lo stesso programma, una chiamata ReplicaId() nel programma restituirà un valore diverso su ogni replica.

Rimodella

Vedi anche XlaBuilder::Reshape e l'operazione Collapse.

Rimodella le dimensioni di un array in una nuova configurazione.

Reshape(operand, new_sizes) Reshape(operand, dimensions, new_sizes)

Argomenti Tipo Semantica
operand XlaOp array di tipo T
dimensions Vettore int64 ordine di compressione delle dimensioni
new_sizes Vettore int64 vettore delle dimensioni delle nuove dimensioni

In termini concettuali, la funzione reshape appiattisce prima una matrice in un vettore unidimensionale di valori di dati, quindi perfeziona questo vettore in una nuova forma. Gli argomenti di input sono un array arbitrario di tipo T, un vettore costante in fase di compilazione di indici di dimensione e un vettore costante in fase di compilazione delle dimensioni per il risultato. I valori nel vettore dimension, se specificati, devono essere una permutazione di tutte le dimensioni di T. Il valore predefinito, se non specificato, è {0, ..., rank - 1}. L'ordine delle dimensioni in dimensions va dalla dimensione con variazioni più lente (più importante) alla dimensione con variazioni più rapide (meno importante) nel nido di loop che comprime l'array di input in una singola dimensione. Il vettore new_sizes determina la dimensione dell'array di output. Il valore all'indice 0 in new_sizes è la dimensione della dimensione 0, il valore all'indice 1 è la dimensione della dimensione 1 e così via. Il prodotto delle dimensioni new_size deve essere uguale al prodotto delle dimensioni delle dimensioni dell'operando. Quando perfezioni l'array compresso nell'array multidimensionale definito da new_sizes, le dimensioni in new_sizes vengono ordinate da una minore variazione (maggiore maggiore) a una più rapida (maggiormente minore).

Ad esempio, definisci v come un array di 24 elementi:

let v = f32[4x2x3] { { {10, 11, 12}, {15, 16, 17} },
                    { {20, 21, 22}, {25, 26, 27} },
                    { {30, 31, 32}, {35, 36, 37} },
                    { {40, 41, 42}, {45, 46, 47} } };

In-order collapse:
let v012_24 = Reshape(v, {0,1,2}, {24});
then v012_24 == f32[24] {10, 11, 12, 15, 16, 17, 20, 21, 22, 25, 26, 27,
                         30, 31, 32, 35, 36, 37, 40, 41, 42, 45, 46, 47};

let v012_83 = Reshape(v, {0,1,2}, {8,3});
then v012_83 == f32[8x3] { {10, 11, 12}, {15, 16, 17},
                          {20, 21, 22}, {25, 26, 27},
                          {30, 31, 32}, {35, 36, 37},
                          {40, 41, 42}, {45, 46, 47} };

Out-of-order collapse:
let v021_24 = Reshape(v, {1,2,0}, {24});
then v012_24 == f32[24]  {10, 20, 30, 40, 11, 21, 31, 41, 12, 22, 32, 42,
                          15, 25, 35, 45, 16, 26, 36, 46, 17, 27, 37, 47};

let v021_83 = Reshape(v, {1,2,0}, {8,3});
then v021_83 == f32[8x3] { {10, 20, 30}, {40, 11, 21},
                          {31, 41, 12}, {22, 32, 42},
                          {15, 25, 35}, {45, 16, 26},
                          {36, 46, 17}, {27, 37, 47} };


let v021_262 = Reshape(v, {1,2,0}, {2,6,2});
then v021_262 == f32[2x6x2] { { {10, 20}, {30, 40},
                              {11, 21}, {31, 41},
                              {12, 22}, {32, 42} },
                             { {15, 25}, {35, 45},
                              {16, 26}, {36, 46},
                              {17, 27}, {37, 47} } };

Come caso speciale, reshape può trasformare un array con singolo elemento in uno scalare e viceversa. Ad esempio,

Reshape(f32[1x1] { {5} }, {0,1}, {}) == 5;
Reshape(5, {}, {1,1}) == f32[1x1] { {5} };

Indietro

Vedi anche XlaBuilder::Rev.

Rev(operand, dimensions)

Argomenti Tipo Semantica
operand XlaOp array di tipo T
dimensions ArraySlice<int64> dimensioni da invertire

Inverte l'ordine degli elementi nell'array operand in base al valore specificato dimensions, generando un array di output della stessa forma. Ogni elemento dell'array operando in un indice multidimensionale viene memorizzato nell'array di output in un indice trasformato. L'indice multidimensionale viene trasformato invertendo l'indice in ogni dimensione da invertire (ad es. se una dimensione di dimensione N è una delle dimensioni di inversione, il suo indice i viene trasformato in N - 1 - i).

Un utilizzo dell'operazione Rev è invertire l'array dei pesi di convoluzione lungo le due dimensioni della finestra durante il calcolo del gradiente nelle reti neurali.

RngNormal

Vedi anche XlaBuilder::RngNormal.

Costruisce un output di una determinata forma con numeri casuali generati seguendo la \(N(\mu, \sigma)\) distribuzione normale. I parametri \(\mu\) e \(\sigma\)e la forma di output devono avere un tipo di elemento a virgola mobile. Inoltre, i parametri devono avere un valore scalare.

RngNormal(mu, sigma, shape)

Argomenti Tipo Semantica
mu XlaOp Scalare di tipo T che specifica la media dei numeri generati
sigma XlaOp Scalare di tipo T che specifica la deviazione standard di quanto generato
shape Shape Forma di output di tipo T

RngUniform

Vedi anche XlaBuilder::RngUniform.

Crea un output di una determinata forma con numeri casuali generati seguendo la distribuzione uniforme nell'intervallo \([a,b)\). I parametri e il tipo di elemento di output devono essere di tipo booleano, di tipo integrale o di tipo a virgola mobile e i tipi devono essere coerenti. Al momento i backend di CPU e GPU supportano solo F64, F32, F16, BF16, S64, U64, S32 e U32. Inoltre, i parametri devono avere un valore scalare. Se \(b <= a\) il risultato è definito dall'implementazione.

RngUniform(a, b, shape)

Argomenti Tipo Semantica
a XlaOp Scalare di tipo T che specifica il limite inferiore dell'intervallo
b XlaOp Scalare di tipo T che specifica il limite superiore dell'intervallo
shape Shape Forma di output di tipo T

RngBitGenerator

Genera un output con una forma specifica riempita con bit casuali uniformi utilizzando l'algoritmo specificato (o quello predefinito del backend) e restituisce uno stato aggiornato (con la stessa forma dello stato iniziale) e i dati casuali generati.

Lo stato iniziale è lo stato iniziale dell'attuale generazione di numeri casuali. La forma e i valori validi richiesti dipendono dall'algoritmo utilizzato.

L'output è garantito come funzione deterministica dello stato iniziale, ma non è garantito come deterministico tra backend e versioni diverse del compilatore.

RngBitGenerator(algorithm, key, shape)

Argomenti Tipo Semantica
algorithm RandomAlgorithm dall'algoritmo PRNG.
initial_state XlaOp Stato iniziale per l'algoritmo PRNG.
shape Shape Forma di output per i dati generati.

Valori disponibili per algorithm:

A dispersione

L'operazione di dispersione XLA genera una sequenza di risultati che sono i valori dell'array di input operands, con diversi slice (agli indici specificati da scatter_indices) aggiornati con la sequenza di valori in updates utilizzando update_computation.

Vedi anche XlaBuilder::Scatter.

scatter(operands..., scatter_indices, updates..., update_computation, index_vector_dim, update_window_dims, inserted_window_dims, scatter_dims_to_operand_dims)

Argomenti Tipo Semantica
operands Sequenza di N XlaOp N array di tipi T_0, ..., T_N in cui eseguire la dispersione.
scatter_indices XlaOp Matrice contenente gli indici iniziali delle sezioni su cui devono essere distribuite.
updates Sequenza di N XlaOp N array di tipi T_0, ..., T_N. updates[i] contiene i valori che devono essere utilizzati per lo scattering di operands[i].
update_computation XlaComputation Calcolo da utilizzare per combinare i valori esistenti nell'array di input e gli aggiornamenti durante la dispersione. Questo calcolo deve essere di tipo T_0, ..., T_N, T_0, ..., T_N -> Collate(T_0, ..., T_N).
index_vector_dim int64 La dimensione in scatter_indices che contiene gli indici iniziali.
update_window_dims ArraySlice<int64> L'insieme di dimensioni nella forma updates che sono dimensioni della finestra.
inserted_window_dims ArraySlice<int64> L'insieme di dimensioni della finestra che deve essere inserito nella forma updates.
scatter_dims_to_operand_dims ArraySlice<int64> Una mappa delle dimensioni dagli indici di dispersione allo spazio dell'indice dell'operando. Questo array è interpretato come la mappatura di i a scatter_dims_to_operand_dims[i] . Deve essere uno a uno e totale.
indices_are_sorted bool Indica se è garantito che gli indici siano ordinati dal chiamante.
unique_indices bool Indica se il chiamante garantisce che gli indici siano univoci.

Dove:

  • N deve essere maggiore o uguale a 1.
  • operands[0], ..., operands[N-1] devono avere tutte le stesse dimensioni.
  • updates[0], ..., updates[N-1] devono avere tutti le stesse dimensioni.
  • Se N = 1, Collate(T) è T.
  • Se N > 1, Collate(T_0, ..., T_N) è una tupla di elementi N di tipo T.

Se index_vector_dim è uguale a scatter_indices.rank, consideriamo implicitamente scatter_indices come contenente una dimensione 1 finale.

Definiamo update_scatter_dims di tipo ArraySlice<int64> come l'insieme di le dimensioni nel shape updates che non sono in update_window_dims, in ordine crescente.

Gli argomenti di dispersione devono seguire questi vincoli:

  • Ogni array updates deve essere di rango update_window_dims.size + scatter_indices.rank - 1.

  • I limiti della dimensione i in ogni array updates devono essere conformi a quanto segue:

    • Se i è presente in update_window_dims (ovvero è uguale a update_window_dims[k] per alcuni k), il limite della dimensione i in updates non deve superare il limite corrispondente di operand dopo aver tenuto conto di inserted_window_dims (ovvero adjusted_window_bounds[k], dove adjusted_window_bounds contiene i limiti di operand con i limiti agli indici inserted_window_dims rimossi).
    • Se i è presente in update_scatter_dims (ovvero uguale a update_scatter_dims[k] per alcuni k), il limite della dimensione i in updates deve essere uguale al limite corrispondente di scatter_indices, saltando index_vector_dim (ad es. scatter_indices.shape.dims[k], se k < index_vector_dim e scatter_indices.shape.dims[k+1] in caso contrario).
  • update_window_dims deve essere in ordine crescente, non deve avere numeri di dimensione ripetuti ed essere compreso nell'intervallo [0, updates.rank).

  • inserted_window_dims deve essere in ordine crescente, non deve avere numeri di dimensioni ripetuti ed essere nell'intervallo [0, operand.rank).

  • operand.rank deve essere uguale alla somma di update_window_dims.size e inserted_window_dims.size.

  • scatter_dims_to_operand_dims.size deve essere uguale a scatter_indices.shape.dims[index_vector_dim] e i suoi valori devono essere nell'intervallo [0, operand.rank).

Per un determinato indice U in ogni array updates, l'indice I corrispondente nell'array operands a cui deve essere applicato questo aggiornamento viene calcolato come segue:

  1. Sia G = { U[k] per k in update_scatter_dims }. Utilizza G per cercare un vettore di indice S nell'array scatter_indices in modo tale che S[i] = scatter_indices[Combina(G, i)] dove Combina(A, b) inserisce b nelle posizioni index_vector_dim in A.
  2. Crea un indice Sin in operand utilizzando S distribuendo S con la mappa scatter_dims_to_operand_dims. In modo più formale:
    1. Sin[scatter_dims_to_operand_dims[k]] = S[k] se k < scatter_dims_to_operand_dims.size.
    2. Sin[_] = 0 in caso contrario.
  3. Crea un indice Win in ogni array operands disperdendo gli indici in corrispondenza di update_window_dims in U in base a inserted_window_dims. In modo più formale:
    1. Win[window_dims_to_operand_dims(k)] = U[k] se k si trova in update_window_dims, dove window_dims_to_operand_dims è la funzione monotonica con dominio [0, update_window_dims.size) e intervallo [0, operand.rank) \ inserted_window_dims. Ad esempio, se update_window_dims.size è 4, operand.rank è 6 e inserted_window_dims è {0, 2}, allora window_dims_to_operand_dims è {01, 13, 24, 35}.
    2. Win[_] = 0 in caso contrario.
  4. I è Win + Sin, dove + è l'addizione elemento per elemento.

In sintesi, l'operazione di dispersione può essere definita come segue.

  • Inizializza output con operands, ovvero per tutti gli indici J, per tutti gli indici O nell'array operands[J]:
    output[J][O] = operands[J][O]
  • Per ogni indice U nell'array updates[J] e l'indice corrispondente O nell'array operand[J], se O è un indice valido per output:
    (output[0][O], ..., output[N-1][O]) =update_computation(output[0][O], ..., ,output[N-1][O],updates[0][U], ...,updates[N-1][U])

L'ordine in cui vengono applicati gli aggiornamenti non è deterministico. Pertanto, quando più indice in updates fanno riferimento allo stesso indice in operands, il valore corrispondente in output sarà non deterministico.

Tieni presente che il primo parametro passato in update_computation sarà sempre il valore corrente dell'array output e il secondo parametro sarà sempre il valore dell'array updates. Questo è importante in particolare nei casi in cui il valore update_computation non è commutativo.

Se indices_are_sorted è impostato su true, XLA può assumere che scatter_indices siano ordinati (in ordine crescente, dopo aver sparso i relativi valori in base a scatter_dims_to_operand_dims) dall'utente. In caso contrario, la semantica viene definita l'implementazione.

Se unique_indices è impostato su true, XLA può presumere che tutti gli elementi tra cui è stato eseguito lo scattering siano univoci. Pertanto, XLA potrebbe utilizzare operazioni non atomiche. Se unique_indices è impostato su true e gli indici a cui viene eseguita la distribuzione non sono unici, la semantica è definita dall'implementazione.

In modo informale, l'operazione scatter può essere considerata come l'inverso dell'operazione gather, ovvero l'operazione scatter aggiorna gli elementi nell'input estratti dall'operazione gather corrispondente.

Per una descrizione informale dettagliata ed esempi, consulta la sezione "Descrizione informale" in Gather.

Seleziona

Vedi anche XlaBuilder::Select.

Costruisce un array di output dagli elementi di due array di input, in base ai valori di un array di predicati.

Select(pred, on_true, on_false)

Argomenti Tipo Semantica
pred XlaOp array di tipo PRED
on_true XlaOp array di tipo T
on_false XlaOp array di tipo T

Gli array on_true e on_false devono avere la stessa forma. Questa è anche la forma dell'array di output. L'array pred deve avere la stessa dimensione di on_true e on_false, con il tipo di elemento PRED.

Per ogni elemento P di pred, l'elemento corrispondente dell'array di output viene preso da on_true se il valore di P è true e da on_false se il valore di P è false. Come forma limitata di trasmissione, pred può essere uno scalare di tipo PRED. In questo caso, l'array di output viene preso interamente da on_true se pred è true e da on_false se pred è false.

Esempio con pred non scalare:

let pred: PRED[4] = {true, false, false, true};
let v1: s32[4] = {1, 2, 3, 4};
let v2: s32[4] = {100, 200, 300, 400};
==>
Select(pred, v1, v2) = s32[4]{1, 200, 300, 4};

Esempio con valore scalare pred:

let pred: PRED = true;
let v1: s32[4] = {1, 2, 3, 4};
let v2: s32[4] = {100, 200, 300, 400};
==>
Select(pred, v1, v2) = s32[4]{1, 2, 3, 4};

Sono supportate le selezioni tra tuple. A questo scopo, le tuple sono considerate tipi scalari. Se on_true e on_false sono tuple (che devono avere la stessa forma), pred deve essere uno scalare di tipo PRED.

SelectAndScatter

Vedi anche XlaBuilder::SelectAndScatter.

Questa operazione può essere considerata un'operazione composta che prima calcola ReduceWindow sull'array operand per selezionare un elemento da ogni finestra, quindi distribuisce l'array source agli indici degli elementi selezionati per costruire un array di output con la stessa forma dell'array di operandi. La funzione select binaria viene utilizzata per selezionare un elemento da ogni finestra applicandola a ogni finestra e viene chiamata con la proprietà che il vettore di indici del primo parametro è lessicograficamente inferiore al vettore di indici del secondo parametro. La funzione select restituisce true se viene selezionato il primo parametro e false se viene selezionato il secondo parametro. Inoltre, la funzione deve rispettare la transitività (ad es. se select(a, b) e select(b, c) sono true, anche select(a, c) è true) in modo che l'elemento selezionato non dipenda dall'ordine degli elementi attraversati per una determinata finestra.

La funzione scatter viene applicata a ogni indice selezionato nell'array di output. Sono necessari due parametri scalari:

  1. Valore corrente nell'indice selezionato nell'array di output
  2. Il valore di dispersione di source che si applica all'indice selezionato

Combina i due parametri e restituisce un valore scalare utilizzato per aggiornare il valore nell'indice selezionato nell'array di output. Inizialmente, tutti gli indici dell'array di output sono impostati su init_value.

L'array di output ha la stessa forma dell'array operand e l'array source deve avere la stessa forma del risultato dell'applicazione di un'operazione ReduceWindow sull'array operand. SelectAndScatter può essere utilizzato per eseguire la propagazione a ritroso dei valori del gradiente per un livello di pooling in una rete neurale.

SelectAndScatter(operand, select, window_dimensions, window_strides, padding, source, init_value, scatter)

Argomenti Tipo Semantica
operand XlaOp array di tipo T su cui scorrono le finestre
select XlaComputation Calcolo binario di tipo T, T -> PRED, da applicare a tutti gli elementi di ogni finestra; restituisce true se è selezionato il primo parametro e false se è selezionato il secondo parametro
window_dimensions ArraySlice<int64> array di numeri interi per i valori delle dimensioni della finestra
window_strides ArraySlice<int64> array di numeri interi per i valori del passo della finestra
padding Padding tipo di spaziatura interna per la finestra (Padding::kSame o Padding::kValid)
source XlaOp array di tipo T con i valori da spargere
init_value XlaOp Valore scalare di tipo T per il valore iniziale dell'array di output
scatter XlaComputation calcolo binario di tipo T, T -> T, per applicare ogni elemento di origine della dispersione con il relativo elemento di destinazione

La figura seguente mostra esempi di utilizzo di SelectAndScatter, in cui la funzione select calcola il valore massimo tra i suoi parametri. Tieni presente che quando le finestre si sovrappongono, come nella figura (2) di seguito, un indice dell'array operand può essere selezionato più volte da finestre diverse. Nella figura, l'elemento di valore 9 è selezionato da entrambe le finestre superiori (blu e rossa) e la funzione di addizione binaria scatter produce l'elemento di output di valore 8 (2 + 6).

L'ordine di valutazione della funzione scatter è arbitrario e può essere non deterministico. Pertanto, la funzione scatter non deve essere eccessivamente sensibile alla riaccoppiamento. Per ulteriori dettagli, consulta la discussione sull'associatività nel contesto di Reduce.

Invia

Vedi anche XlaBuilder::Send.

Send(operand, channel_handle)

Argomenti Tipo Semantica
operand XlaOp dati da inviare (array di tipo T)
channel_handle ChannelHandle identificatore univoco per ogni coppia di invii/ricezioni

Invia i dati dell'operando specificato a un'istruzione Recv in un altro calcolo che condivide lo stesso handle del canale. Non restituisce dati.

Analogamente all'operazione Recv, l'API client dell'operazione Send rappresenta la comunicazione sincrona ed è scomposta internamente in 2 istruzioni HLO (Send e SendDone) per consentire i trasferimenti di dati asincroni. Vedi anche HloInstruction::CreateSend e HloInstruction::CreateSendDone.

Send(HloInstruction operand, int64 channel_id)

Avvia un trasferimento asincrono dell'operando nelle risorse allocate dall'istruzione Recv con lo stesso ID canale. Restituisce un contesto, che viene utilizzato da un'istruzione SendDone successiva per attendere il completamento del trasferimento dei dati. Il contesto è una tupla di {operando (forma), identificatore richiesta (U32)} e può essere utilizzato solo da un'istruzione SendDone.

SendDone(HloInstruction context)

Dato un contesto creato da un'istruzione Send, attende il completamento del trasferimento dei dati. L'istruzione non restituisce dati.

Programmazione delle istruzioni del canale

L'ordine di esecuzione delle 4 istruzioni per ciascun canale (Recv, RecvDone, Send, SendDone) è il seguente.

  • Recv si verifica prima del giorno Send
  • Send si verifica prima del giorno RecvDone
  • Recv si verifica prima del giorno RecvDone
  • Send si verifica prima del giorno SendDone

Quando i compilatori backend generano una pianificazione lineare per ogni calcolo che comunica tramite le istruzioni del canale, non devono esserci cicli tra i calcoli. Ad esempio, le programmazioni riportate di seguito portano a dei deadlock.

Sezione

Vedi anche XlaBuilder::Slice.

La suddivisione estrae un sottoarray dall'array di input. Il sottoarray ha lo stesso ranking dell'input e contiene i valori all'interno di un riquadro di delimitazione all'interno dell'array di input, in cui le dimensioni e gli indici del riquadro di delimitazione sono forniti come argomenti per l'operazione della sezione.

Slice(operand, start_indices, limit_indices, strides)

Argomenti Tipo Semantica
operand XlaOp Array N-dimensionale di tipo T
start_indices ArraySlice<int64> Elenco di N numeri interi contenenti gli indici iniziali del segmento per ogni dimensione. I valori devono essere maggiori o uguali a zero.
limit_indices ArraySlice<int64> Elenco di N numeri interi contenenti gli indici di fine (esclusivi) per il segmento per ogni dimensione. Ogni valore deve essere maggiore o uguale al rispettivo valore start_indices per la dimensione e minore o uguale alla dimensione della dimensione.
strides ArraySlice<int64> Elenco di N numeri interi che determinano lo stride di input della sezione. La sezione sceglie ogni elemento strides[d] nella dimensione d.

Esempio unidimensionale:

let a = {0.0, 1.0, 2.0, 3.0, 4.0}
Slice(a, {2}, {4}) produces:
  {2.0, 3.0}

Esempio bidimensionale:

let b =
 { {0.0,  1.0,  2.0},
   {3.0,  4.0,  5.0},
   {6.0,  7.0,  8.0},
   {9.0, 10.0, 11.0} }

Slice(b, {2, 1}, {4, 3}) produces:
  { { 7.0,  8.0},
    {10.0, 11.0} }

Ordina

Vedi anche XlaBuilder::Sort.

Sort(operands, comparator, dimension, is_stable)

Argomenti Tipo Semantica
operands ArraySlice<XlaOp> Gli operandi da ordinare.
comparator XlaComputation Il calcolo del comparatore da utilizzare.
dimension int64 La dimensione in base alla quale ordinare.
is_stable bool Indica se deve essere utilizzato l'ordinamento stabile.

Se viene fornito un solo operando:

  • Se l'operando è un tensore di rango 1 (un array), il risultato è un array ordinato. Se vuoi ordinare l'array in ordine crescente, il comparatore deve eseguire un confronto minore di. Formalmente, dopo l'ordinamento dell'array, vale per tutte le posizioni dell'indice i, j con i < j che è comparator(value[i], value[j]) = comparator(value[j], value[i]) = false o comparator(value[i], value[j]) = true.

  • Se l'operando ha un ranking più elevato, l'operando viene ordinato in base alla dimensione fornita. Ad esempio, per un tensore di rango 2 (una matrice), un valore della dimensione 0 ordina in modo indipendente ogni colonna, mentre un valore della dimensione 1 ordina in modo indipendente ogni riga. Se non viene fornito alcun numero di dimensione, per impostazione predefinita viene scelta l'ultima dimensione. Per la dimensione ordinata, viene applicato lo stesso ordine di ordinamento del caso con ranking 1.

Se vengono forniti operandi n > 1:

  • Tutti gli operandi n devono essere tensori con le stesse dimensioni. I tipi di elementi degli tensori possono essere diversi.

  • Tutti gli operandi vengono ordinati insieme, non singolarmente. Concettualmente, gli operandi vengono trattati come una tupla. Quando si verifica se è necessario scambiare gli elementi di ciascun operando nelle posizioni dell'indice i e j, il comparatore viene chiamato con parametri scalari 2 * n, dove il parametro 2 * k corrisponde al valore nella posizione i dell'operando k-th e il parametro 2 * k + 1 corrisponde al valore nella posizione j dell'operando k-th. In genere, il comparatore confronta i parametri 2 * k e 2 * k + 1 tra loro e potrebbe utilizzare altre coppie di parametri come criteri di parità.

  • Il risultato è una tupla costituita dagli operandi in ordine crescente (in base alla dimensione indicata sopra). L'operando i-th della tupla corrisponde all'operando i-th di Ordina.

Ad esempio, se sono presenti tre operandi operand0 = [3, 1], operand1 = [42, 50], operand2 = [-3.0, 1.1] e il comparatore confronta solo i valori di operand0 con un valore minore, l'output dell'ordinamento è la tupla ([1, 3], [50, 42], [1.1, -3.0]).

Se is_stable è impostato su true, l'ordinamento è garantito come stabile, ovvero se esistono elementi considerati uguali dal comparatore, l'ordine relativo dei valori uguali viene mantenuto. Due elementi e1 e e2 sono uguali se e solo se comparator(e1, e2) = comparator(e2, e1) = false. Per impostazione predefinita, is_stable è impostato su false.

Transpose

Vedi anche l'operazione tf.reshape.

Transpose(operand)

Argomenti Tipo Semantica
operand XlaOp L'operando da trasporre.
permutation ArraySlice<int64> Come permutare le dimensioni.

Attiva le dimensioni dell'operando con la permutazione specificata, quindi ∀ i . 0 ≤ i < rank ⇒ input_dimensions[permutation[i]] = output_dimensions[i].

È lo stesso di Reshape(operand, permutation, Permute(permutation, operand.shape.dimensions)).

TriangularSolve

Vedi anche XlaBuilder::TriangularSolve.

Risolve i sistemi di equazioni lineari con matrici di coefficienti inferiori o superiori tramite sostituzione diretta o inversa. Eseguendo la trasmissione lungo le dimensioni principali, questa routine risolve uno dei sistemi di matrici op(a) * x = b o x * op(a) = b per la variabile x, dati a e b, dove op(a) è op(a) = a, op(a) = Transpose(a) o op(a) = Conj(Transpose(a)).

TriangularSolve(a, b, left_side, lower, unit_diagonal, transpose_a)

Argomenti Tipo Semantica
a XlaOp un array con ranking > 2 di tipo complesso o con rappresentazione in virgola mobile, con forma [..., M, M].
b XlaOp un array di rango > 2 dello stesso tipo con forma [..., M, K] se left_side è vero, [..., K, M] altrimenti.
left_side bool indica se risolvere un sistema del tipo op(a) * x = b (true) o x * op(a) = b (false).
lower bool se utilizzare il triangolo superiore o inferiore di a.
unit_diagonal bool Se true, si presume che gli elementi diagonali di a siano 1 e non vengano acceduti.
transpose_a Transpose se utilizzare a così com'è, trasporlo o prendere la sua trasposta coniugata.

I dati di input vengono letti solo dal triangolo inferiore/superiore di a, a seconda del valore di lower. I valori dell'altro triangolo vengono ignorati. I dati di output vengono restituiti nello stesso triangolo; i valori nell'altro triangolo sono definiti dall'implementazione e possono essere qualsiasi cosa.

Se il ranking di a e b è maggiore di 2, vengono trattati come batch di matrici, dove tutte le dimensioni tranne le 2 secondarie sono dimensioni batch. a e b devono avere le stesse dimensioni del batch.

Tupla

Vedi anche XlaBuilder::Tuple.

Una tupla contenente un numero variabile di handle di dati, ognuno dei quali ha una propria forma.

È analogo a std::tuple in C++. Concettualmente:

let v: f32[10] = f32[10]{0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
let s: s32 = 5;
let t: (f32[10], s32) = tuple(v, s);

È possibile scomporre le tuple (con accesso) tramite l'operazione GetTupleElement.

La funzione

Vedi anche XlaBuilder::While.

While(condition, body, init)

Argomenti Tipo Semantica
condition XlaComputation XlaComputation di tipo T -> PRED che definisce la condizione di terminazione del loop.
body XlaComputation XlaComputation di tipo T -> T che definisce il corpo del loop.
init T Valore iniziale per il parametro condition e body.

Esegue in sequenza il body finché il condition non fallisce. È simile a un tipico while loop in molti altri linguaggi, ad eccezione delle differenze e delle limitazioni elencate di seguito.

  • Un nodo While restituisce un valore di tipo T, che è il risultato dell'ultima esecuzione di body.
  • La forma del tipo T viene determinata in modo statico e deve essere la stessa in tutte le iterazioni.

I parametri T dei calcoli vengono inizializzati con il valore init nella prima iterazione e vengono aggiornati automaticamente al nuovo risultato da body in ogni iterazione successiva.

Un caso d'uso principale del nodo While è implementare l'esecuzione ripetuta dell'addestramento nelle reti neurali. Di seguito è riportato lo pseudocodice semplificato con un grafico che rappresenta il calcolo. Il codice è disponibile in while_test.cc. Il tipo T in questo esempio è un Tuple costituito da un int32 per il conteggio delle iterazioni e un vector[10] per l'accumulatore. Per 1000 iterazioni, il loop continua ad aggiungere un vettore costante all'accumulatore.

// Pseudocode for the computation.
init = {0, zero_vector[10]} // Tuple of int32 and float[10].
result = init;
while (result(0) < 1000) {
  iteration = result(0) + 1;
  new_vector = result(1) + constant_vector[10];
  result = {iteration, new_vector};
}