Vorgangssemantik

Im Folgenden wird die Semantik von Vorgängen beschrieben, die in der XlaBuilder-Schnittstelle definiert sind. In der Regel werden diese Vorgänge 1:1-Vorgängen zugeordnet, die in der RPC-Schnittstelle in xla_data.proto definiert sind.

Hinweis zur Nomenklatur: Der generalisierte Datentyp XLA ist ein n-dimensionales Array, das Elemente eines einheitlichen Typs enthält (z. B. 32-Bit-Gleitkommazahl). In der gesamten Dokumentation wird array verwendet, um ein Array mit beliebiger Dimension anzugeben. Sonderfälle haben spezifischere und bekanntere Namen. Ein Vektor ist beispielsweise ein eindimensionales Array und eine Matrix ein zweidimensionales Array.

AfterAll

Weitere Informationen finden Sie unter XlaBuilder::AfterAll.

„AfterAll“ verwendet eine unterschiedliche Anzahl von Tokens und erzeugt ein einzelnes Token. Tokens sind primitive Typen, in die sich Nebenwirkungsvorgängen miteinander verbinden können, um eine Sortierung zu erzwingen. AfterAll kann als Join von Tokens verwendet werden, um einen Vorgang nach einer festgelegten Anzahl von Vorgängen zu bestellen.

AfterAll(operands)

Argumente Typ Semantik
operands XlaOp unterschiedlich viele Tokens

AllGather

Weitere Informationen finden Sie unter XlaBuilder::AllGather.

Führt eine Verkettung über Replikate hinweg durch.

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

Argumente Typ Semantik
operand XlaOp Array zum Verketten zwischen Replikaten
all_gather_dim int64 Verkettungsdimension
replica_groups Vektor der Vektoren von int64 Gruppen, zwischen denen die Verkettung durchgeführt wird,
channel_id int64 (optional) Optionale Kanal-ID für die modulübergreifende Kommunikation
  • replica_groups ist eine Liste von Replikatgruppen, zwischen denen die Verkettung durchgeführt wird (die Replikat-ID für das aktuelle Replikat kann mit ReplicaId abgerufen werden). Die Reihenfolge der Replikate in jeder Gruppe bestimmt die Reihenfolge, in der sich die Eingaben im Ergebnis befinden. replica_groups muss entweder leer sein (in diesem Fall gehören alle Replikate zu einer einzigen Gruppe in der Reihenfolge von 0 bis N - 1) oder muss dieselbe Anzahl von Elementen wie die Anzahl der Replikate enthalten. replica_groups = {0, 2}, {1, 3} führt beispielsweise eine Verkettung zwischen den Replikaten 0 und 2 sowie 1 und 3 durch.
  • shard_count ist die Größe der einzelnen Replikatgruppen. Wir benötigen dies, wenn replica_groups leer ist.
  • channel_id wird für die modulübergreifende Kommunikation verwendet: Nur all-gather-Vorgänge mit derselben channel_id können miteinander kommunizieren.

Die Ausgabeform ist die Eingabeform mit dem all_gather_dim shard_count-mal größer. Wenn beispielsweise zwei Replikate vorhanden sind und der Operand den Wert [1.0, 2.5] bzw. [3.0, 5.25] auf den beiden Replikaten hat, ist der Ausgabewert dieses Vorgangs mit all_gather_dim 0 für beide Replikate [1.0, 2.5, 3.0, 5.25].

AllReduce

Weitere Informationen finden Sie unter XlaBuilder::AllReduce.

Führt eine benutzerdefinierte Berechnung für alle Replikate durch.

AllReduce(operand, computation, replica_group_ids, channel_id)

Argumente Typ Semantik
operand XlaOp Array oder ein nicht leeres Tupel von Arrays, um die Anzahl der Replikate zu reduzieren
computation XlaComputation Reduktionsberechnung
replica_groups Vektor der Vektoren von int64 Gruppen, zwischen denen die Reduzierungen
channel_id int64 (optional) Optionale Kanal-ID für die modulübergreifende Kommunikation
  • Wenn operand ein Tupel von Arrays ist, wird die All-Reduce-Funktion für jedes Element des Tupels ausgeführt.
  • replica_groups ist eine Liste von Replikatgruppen, zwischen denen die Reduzierung durchgeführt wird (die Replikat-ID für das aktuelle Replikat kann mit ReplicaId abgerufen werden). replica_groups muss entweder leer sein (in diesem Fall gehören alle Replikate zu einer einzigen Gruppe) oder dieselbe Anzahl von Elementen wie die Anzahl der Replikate enthalten. Beispiel: replica_groups = {0, 2}, {1, 3} reduziert die Replikate 0 und 2 sowie 1 und 3.
  • channel_id wird für die modulübergreifende Kommunikation verwendet: Nur all-reduce-Vorgänge mit derselben channel_id können miteinander kommunizieren.

Die Ausgabeform ist mit der Eingabeform identisch. Wenn beispielsweise zwei Replikate vorhanden sind und der Operand in den beiden Replikaten den Wert [1.0, 2.5] bzw. [3.0, 5.25] hat, ist der Ausgabewert dieses Vorgangs und der Summenberechnung für beide Replikate [4.0, 7.75]. Wenn die Eingabe ein Tupel ist, ist auch die Ausgabe ein Tupel.

Für die Berechnung des Ergebnisses von AllReduce ist eine Eingabe von jedem Replikat erforderlich. Wenn also ein Replikat einen AllReduce-Knoten häufiger als ein anderes ausführt, wartet das frühere Replikat unbegrenzt. Da die Replikate alle dasselbe Programm ausführen, gibt es nicht viele Möglichkeiten, dies zu tun. Es ist jedoch möglich, wenn die Bedingung einer Während-Schleife von Daten aus der Einspeisung abhängt und die eingespeisten Daten die Während-Schleife dazu veranlassen, ein Replikat öfter als ein anderes zu durchlaufen.

AllToAll

Weitere Informationen finden Sie unter XlaBuilder::AllToAll.

AllToAll ist ein kollektiver Vorgang, bei dem Daten von allen Kernen an alle Kerne gesendet werden. Sie besteht aus zwei Phasen:

  1. Streuphase. Auf jedem Kern wird der Operand in eine split_count-Anzahl von Blöcken entlang des split_dimensions aufgeteilt. Die Blöcke werden auf alle Kerne verteilt, z. B. wird der i-te-Block zum i-ten Kern gesendet.
  2. Die Erfassungsphase. Jeder Kern verkettet die empfangenen Blöcke entlang der concat_dimension.

Die teilnehmenden Kerne können so konfiguriert werden:

  • replica_groups: Jede Replikatgruppe enthält eine Liste der an der Berechnung beteiligten Replikat-IDs (die Replikat-ID für das aktuelle Replikat kann mit ReplicaId abgerufen werden). „AllToAll“ wird in Untergruppen in der angegebenen Reihenfolge angewendet. replica_groups = { {1,2,3}, {4,5,0} } bedeutet beispielsweise, dass ein „AllToAll“ in den Replikaten {1, 2, 3} und in der Erfassungsphase angewendet wird und die empfangenen Blöcke in der Reihenfolge von 1, 2, 3 verkettet werden. Dann wird ein weiteres AllToAll innerhalb der Replikate 4, 5, 0 angewendet und die Verkettungsreihenfolge lautet ebenfalls 4, 5, 0. Wenn replica_groups leer ist, gehören alle Replikate in der Verkettungsreihenfolge ihrer Darstellung zu einer Gruppe.

Voraussetzungen:

  • Die Dimensionsgröße des Operanden auf split_dimension ist durch split_count teilbar.
  • Die Form des Operanden ist kein Tupel.

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

Argumente Typ Semantik
operand XlaOp n-dimensionales Eingabearray
split_dimension int64 Ein Wert im Intervall [0, n), der die Dimension bezeichnet, anhand der der Operand aufgeteilt wird
concat_dimension int64 Ein Wert im Intervall [0, n), der die Dimension angibt, mit der die Aufteilungsblöcke verkettet sind
split_count int64 Die Anzahl der Kerne, die an diesem Vorgang beteiligt sind. Wenn replica_groups leer ist, sollte dies die Anzahl der Replikate sein. Andernfalls sollte dies der Anzahl der Replikate in jeder Gruppe entsprechen.
replica_groups Vektor: ReplicaGroup Jede Gruppe enthält eine Liste von Replikat-IDs.

Unten sehen Sie ein Beispiel für 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 diesem Beispiel nehmen 4 Kerne am Alltoall teil. Auf jedem Kern wird der Operand in 4 Teile entlang der Dimension 0 aufgeteilt, sodass jeder Teil die Form f32[4,4] hat. Die vier Teile sind auf alle Kerne verteilt. Dann verkettet jeder Kern die erhaltenen Teile entlang der Dimension 1 in der Reihenfolge von Kern 0–4. Die Ausgabe auf jedem Kern hat also die Form f32[16,4].

BatchNormGrad

Eine ausführliche Beschreibung des Algorithmus finden Sie unter XlaBuilder::BatchNormGrad und im ursprünglichen Batch-Normalisierungsdokument.

Berechnet Gradienten der Batchnorm.

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

Argumente Typ Semantik
operand XlaOp n-dimensionales Array, das normalisiert werden soll (x)
scale XlaOp 1-dimensionales Array (\(\gamma\))
mean XlaOp 1-dimensionales Array (\(\mu\))
variance XlaOp 1-dimensionales Array (\(\sigma^2\))
grad_output XlaOp Farbverläufe an BatchNormTraining (\(\nabla y\)) weitergegeben
epsilon float Epsilon-Wert (\(\epsilon\))
feature_index int64 Index für Featuredimension in operand

Für jedes Element in der Featuredimension (feature_index ist der Index für die Featuredimension in operand) berechnet der Vorgang die Gradienten in Bezug auf operand, offset und scale für alle anderen Dimensionen. feature_index muss ein gültiger Index für die Featuredimension in operand sein.

Die drei Gradienten werden durch die folgenden Formeln definiert (angenommen ein vierdimensionales Array als operand und mit dem Featuredimensionsindex l, der Batchgröße m und den räumlichen Größen w und 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} \]

Die Eingaben mean und variance stellen Momentwerte aus Batch- und räumlichen Dimensionen dar.

Der Ausgabetyp ist ein Tupel aus drei Handles:

Ausgaben Typ Semantik
grad_operand XlaOp Gradienten in Bezug auf die Eingabe operand ($\nabla x$)
grad_scale XlaOp Gradienten in Bezug auf die Eingabe scale ($\nabla \gamma$)
grad_offset XlaOp Gradienten in Bezug auf die Eingabe offset($\nabla \beta$)

BatchNormInference

Eine ausführliche Beschreibung des Algorithmus finden Sie unter XlaBuilder::BatchNormInference und im ursprünglichen Batch-Normalisierungsdokument.

Normalisiert ein Array über Batch- und räumliche Dimensionen hinweg.

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

Argumente Typ Semantik
operand XlaOp n-dimensionales Array, das normalisiert werden soll
scale XlaOp 1-dimensionales Array
offset XlaOp 1-dimensionales Array
mean XlaOp 1-dimensionales Array
variance XlaOp 1-dimensionales Array
epsilon float Epsilon-Wert
feature_index int64 Index für Featuredimension in operand

Für jedes Element in der Featuredimension (feature_index ist der Index für die Featuredimension in operand) berechnet der Vorgang den Mittelwert und die Varianz über alle anderen Dimensionen hinweg und verwendet den Mittelwert und die Varianz, um jedes Element in operand zu normalisieren. feature_index muss ein gültiger Index für die Featuredimension in operand sein.

BatchNormInference entspricht dem Aufrufen von BatchNormTraining, ohne mean und variance für jeden Batch zu berechnen. Die Eingabewerte mean und variance werden stattdessen als geschätzte Werte verwendet. Der Zweck dieser Operation besteht darin, die Latenz in der Inferenz zu verringern, daher der Name BatchNormInference.

Die Ausgabe ist ein n-dimensionales, normalisiertes Array mit derselben Form wie die Eingabe-operand.

BatchNormTraining

Eine ausführliche Beschreibung des Algorithmus finden Sie unter XlaBuilder::BatchNormTraining und the original batch normalization paper.

Normalisiert ein Array über Batch- und räumliche Dimensionen hinweg.

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

Argumente Typ Semantik
operand XlaOp n-dimensionales Array, das normalisiert werden soll (x)
scale XlaOp 1-dimensionales Array (\(\gamma\))
offset XlaOp 1-dimensionales Array (\(\beta\))
epsilon float Epsilon-Wert (\(\epsilon\))
feature_index int64 Index für Featuredimension in operand

Für jedes Element in der Featuredimension (feature_index ist der Index für die Featuredimension in operand) berechnet der Vorgang den Mittelwert und die Varianz über alle anderen Dimensionen hinweg und verwendet den Mittelwert und die Varianz, um jedes Element in operand zu normalisieren. feature_index muss ein gültiger Index für die Featuredimension in operand sein.

Der Algorithmus geht für jeden Batch in operand \(x\) , der m-Elemente mit w und h als Größe der räumlichen Dimensionen enthält (vorausgesetzt, operand ist ein vierdimensionales Array):

  • Berechnet den Batchmittelwert \(\mu_l\) für jede Feature-l in der Feature-Dimension: \(\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}\)

  • Berechnet die Batchabweichung \(\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$

  • Normalisiert, skaliert und verschoben: \(y_{ijkl}=\frac{\gamma_l(x_{ijkl}-\mu_l)}{\sqrt[2]{\sigma^2_l+\epsilon} }+\beta_l\)

Der Epsilon-Wert, in der Regel eine kleine Zahl, wird addiert, um Fehler durch eine Division durch Null zu vermeiden.

Der Ausgabetyp ist ein Tupel aus drei XlaOps:

Ausgaben Typ Semantik
output XlaOp n-dimensionales Array mit derselben Form wie die Eingabeoperand (y)
batch_mean XlaOp 1-dimensionales Array (\(\mu\))
batch_var XlaOp 1-dimensionales Array (\(\sigma^2\))

batch_mean und batch_var sind Momente, die über die Batch- und räumlichen Dimensionen hinweg mithilfe der obigen Formeln berechnet werden.

BitcastConvertType

Weitere Informationen finden Sie unter XlaBuilder::BitcastConvertType.

Führt ähnlich wie bei einem tf.bitcast in TensorFlow einen elementweisen Bitcast-Vorgang von einer Datenform zu einer Zielform aus. Die Eingabe- und Ausgabegröße muss übereinstimmen: s32-Elemente werden z.B. über die Bitcast-Routine zu f32-Elementen und ein s32-Element wird zu vier s8-Elementen. Bitcast wird als Low-Level-Umwandlung implementiert, sodass Maschinen mit unterschiedlichen Gleitkommadarstellungen unterschiedliche Ergebnisse liefern.

BitcastConvertType(operand, new_element_type)

Argumente Typ Semantik
operand XlaOp Array vom Typ T mit Dimmungen D
new_element_type PrimitiveType Typ U

Die Abmessungen des Operanden und der Zielform müssen übereinstimmen, mit Ausnahme der letzten Dimension, die sich durch das Verhältnis der primitiven Größe vor und nach der Konvertierung ändert.

Die Quell- und Zielelementtypen dürfen keine Tupel sein.

Bitcast-Konvertierung in einen primitiven Typ mit unterschiedlicher Breite

Die HLO-Anweisung BitcastConvert unterstützt den Fall, dass die Größe des Ausgabeelementtyps T' nicht mit der Größe des Eingabeelements T übereinstimmt. Da der gesamte Vorgang konzeptionell ein Bitcast ist und die zugrunde liegenden Byte nicht geändert werden, muss sich die Form des Ausgabeelements ändern. Für B = sizeof(T), B' = sizeof(T') gibt es zwei mögliche Fälle.

Beim B > B' erhält die Ausgabeform zuerst eine neue Nebendimension der Größe B/B'. Beispiel:

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

Die Regel für effektive Skalare bleibt dieselbe:

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

Alternativ dazu muss für den Befehl B' > B für die Anweisung die letzte logische Dimension der Eingabeform gleich B'/B sein. Diese Dimension wird dann bei der Konvertierung entfernt:

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

Beachten Sie, dass Konvertierungen zwischen verschiedenen Bitbreiten nicht elementweise sind.

Nachricht an alle

Weitere Informationen finden Sie unter XlaBuilder::Broadcast.

Fügt einem Array Dimensionen hinzu, indem die Daten im Array dupliziert werden.

Broadcast(operand, broadcast_sizes)

Argumente Typ Semantik
operand XlaOp Das zu duplizierende Array
broadcast_sizes ArraySlice<int64> Die Größen der neuen Dimensionen

Die neuen Dimensionen werden links eingefügt. Wenn broadcast_sizes also die Werte {a0, ..., aN} und die Operandenform die Abmessungen {b0, ..., bM} hat, hat die Form der Ausgabe die Abmessungen {a0, ..., aN, b0, ..., bM}.

den neuen Dimensionsindex in Kopien des Operanden, d.h.

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

Wenn operand beispielsweise ein skalarer f32 mit dem Wert 2.0f ist und broadcast_sizes den Wert {2, 3} hat, ist das Ergebnis ein Array mit der Form f32[2, 3] und alle Werte im Ergebnis sind 2.0f.

BroadcastInDim

Weitere Informationen finden Sie unter XlaBuilder::BroadcastInDim.

Erweitert die Größe und den Rang eines Arrays, indem die Daten im Array dupliziert werden.

BroadcastInDim(operand, out_dim_size, broadcast_dimensions)

Argumente Typ Semantik
operand XlaOp Das zu duplizierende Array
out_dim_size ArraySlice<int64> Die Größen der Abmessungen der Zielform
broadcast_dimensions ArraySlice<int64> welcher Dimension in der Zielform die einzelnen Dimensionen der Operandenform entsprechen

Ähnlich wie bei Broadcast, ermöglicht aber das Hinzufügen von Dimensionen an beliebiger Stelle und das Erweitern vorhandener Dimensionen auf Größe 1.

Das operand wird an die von out_dim_size beschriebene Form übertragen. broadcast_dimensions ordnet die Abmessungen von operand den Abmessungen der Zielform zu, d.h., die i-te Dimension des Operanden wird der „broadcast_dimension[i]“-Dimension der Ausgabeform zugeordnet. Die Abmessungen von operand müssen Größe 1 oder dieselbe Größe wie die Dimension in der Ausgabeform haben, der sie zugeordnet sind. Die übrigen Abmessungen werden mit Abmessungen der Größe 1 gefüllt. Beim Broadcasting ohne Dimensionierung wird dann ein Broadcast an diese degenerierten Dimensionen gesendet, um die Ausgabeform zu erreichen. Die Semantik wird auf der Broadcasting-Seite ausführlich beschrieben.

Call

Weitere Informationen finden Sie unter XlaBuilder::Call.

Löst eine Berechnung mit den angegebenen Argumenten aus.

Call(computation, args...)

Argumente Typ Semantik
computation XlaComputation Berechnung des Typs T_0, T_1, ..., T_{N-1} -> S mit N Parametern beliebigen Typs
args Folge von N XlaOp s N-Argumente beliebigen Typs

Die ARity und die Typen von args müssen mit den Parametern von computation übereinstimmen. args darf nicht vorhanden sein.

Cholesky

Weitere Informationen finden Sie unter XlaBuilder::Cholesky.

Berechnet die Cholesky-Zerlegung eines Batches symmetrischer (Hermitianischer) positiver Definite Matrizen.

Cholesky(a, lower)

Argumente Typ Semantik
a XlaOp ein Array des Typs Rang > 2 eines komplexen oder Gleitkommatyps.
lower bool ob das obere oder untere Dreieck von a verwendet werden soll.

Wenn lower true ist, werden Matrizen mit niedriger Dreiecksform l so berechnet, dass $a = l . l^T$. Wenn lower false ist, werden Matrizen des oberen Dreiecks u so berechnet, dass\(a = u^T . u\).

Eingabedaten werden je nach Wert von lower nur aus dem unteren/oberen Dreieck von a gelesen. Werte aus dem anderen Dreieck werden ignoriert. Die Ausgabedaten werden im gleichen Dreieck zurückgegeben. Die Werte im anderen Dreieck sind implementierungsdefiniert und können beliebig sein.

Wenn der Rang von a größer als 2 ist, wird a als Batch von Matrizen behandelt, wobei alle außer den zwei untergeordneten Dimensionen Batchdimensionen sind.

Wenn a nicht symmetrisch (Hermitian) positiv-definite ist, ist das Ergebnis implementierungsdefiniert.

Klemme

Weitere Informationen finden Sie unter XlaBuilder::Clamp.

Setzt einen Operanden auf den Bereich zwischen einem Mindest- und Höchstwert.

Clamp(min, operand, max)

Argumente Typ Semantik
min XlaOp Array vom Typ T
operand XlaOp Array vom Typ T
max XlaOp Array vom Typ T

Bei einem Operanden sowie einem Mindest- und Höchstwert wird der Operand zurückgegeben, wenn er im Bereich zwischen dem Mindest- und Höchstwert liegt. Andernfalls wird der Mindestwert zurückgegeben, wenn der Operand unter diesem Bereich liegt, oder den Maximalwert, wenn der Operand über diesem Bereich liegt. Der Wert ist clamp(a, x, b) = min(max(a, x), b).

Alle drei Arrays müssen die gleiche Form haben. Als eingeschränkte Form der Broadcasting-Funktion können min und/oder max ein Skalar vom Typ T sein.

Beispiel mit dem skalaren min und 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};

Minimieren

Weitere Informationen finden Sie unter XlaBuilder::Collapse und tf.reshape-Vorgang.

Minimiert die Dimensionen eines Arrays in eine Dimension.

Collapse(operand, dimensions)

Argumente Typ Semantik
operand XlaOp Array vom Typ T
dimensions Vektor: int64 eine zusammenhängende Teilmenge der T-Dimensionen.

Bei der Minimierung wird die angegebene Teilmenge der Abmessungen des Operanden durch eine einzelne Dimension ersetzt. Die Eingabeargumente sind ein beliebiges Array vom Typ T und ein Vektor der kompilierten Zeitkonstanten von Dimensionsindexen. Die Dimensionsindexe müssen eine fortlaufende Teilmenge der T-Dimensionen sein (aufsteigend). Daher sind {0, 1, 2}, {0, 1} oder {1, 2} gültige Dimensionssätze, {1, 0} oder {0, 2} jedoch nicht. Sie werden durch eine einzelne neue Dimension ersetzt, die sich an derselben Position in der Abfolge der Dimensionen befindet wie die, die sie ersetzen, wobei die neue Dimensionsgröße dem Produkt der ursprünglichen Dimensionsgrößen entspricht. Die niedrigste Dimensionsnummer in dimensions ist die langsamste variierende Dimension (größte) in der Schleifenverschachtelung, die diese Dimensionen minimiert, und die höchste Dimensionsnummer ist die am schnellsten variierende (größte) Dimension. Wenn eine allgemeinere Minimierungsreihenfolge erforderlich ist, verwenden Sie den Operator tf.reshape.

Ist v beispielsweise ein Array mit 24 Elementen:

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

Weitere Informationen finden Sie unter XlaBuilder::CollectivePermute.

CollectivePermute ist ein kollektiver Vorgang, der Daten replikatübergreifend sendet und empfängt.

CollectivePermute(operand, source_target_pairs)

Argumente Typ Semantik
operand XlaOp n-dimensionales Eingabearray
source_target_pairs Vektor: <int64, int64> Eine Liste der Paare (source_Replikat_id, target_repli_id). Für jedes Paar wird der Operand vom Quellreplikat zum Zielreplikat gesendet.

Beachten Sie, dass für source_target_pair die folgenden Einschränkungen gelten:

  • Zwei Paare sollten nicht dieselbe Zielreplikat-ID und dieselbe Quellreplikat-ID haben.
  • Wenn eine Replikat-ID kein Ziel in einem Paar ist, ist die Ausgabe für dieses Replikat ein Tensor, der aus 0(s) mit derselben Form wie die Eingabe besteht.

Verketten

Weitere Informationen finden Sie unter XlaBuilder::ConcatInDim.

Durch Verketten wird ein Array aus mehreren Array-Operanden erstellt. Das Array hat den gleichen Rang wie die einzelnen Eingabearray-Operanden (die denselben Rang wie die anderen haben müssen) und enthält die Argumente in der Reihenfolge, in der sie angegeben wurden.

Concatenate(operands..., dimension)

Argumente Typ Semantik
operands Sequenz von N XlaOp N Arrays vom Typ T mit den Dimensionen [L0, L1, ...]. Erfordert N >= 1.
dimension int64 Ein Wert im Intervall [0, N), der die Dimension angibt, die zwischen operands verkettet werden soll.

Mit Ausnahme von dimension müssen alle Dimensionen gleich sein. Das liegt daran, dass XLA keine „unregelmäßigen“ Arrays unterstützt. Rang-0-Werte können nicht verkettet werden, da es nicht möglich ist, die Dimension zu benennen, mit der die Verkettung erfolgt.

Beispiel für eindimensionales Element:

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

Beispiel für zweidimensionale Darstellung:

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

Diagramm:

Bedingt

Weitere Informationen finden Sie unter XlaBuilder::Conditional.

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

Argumente Typ Semantik
pred XlaOp Skalar des Typs PRED
true_operand XlaOp Argument vom Typ \(T_0\)
true_computation XlaComputation XlaComputation vom Typ \(T_0 \to S\)
false_operand XlaOp Argument vom Typ \(T_1\)
false_computation XlaComputation XlaComputation vom Typ \(T_1 \to S\)

Führt true_computation aus, wenn pred gleich true ist, false_computation, wenn pred gleich false ist, und gibt das Ergebnis zurück.

true_computation muss ein einzelnes Argument vom Typ \(T_0\) annehmen und wird mit true_operand aufgerufen, das denselben Typ haben muss. false_computation muss ein einzelnes Argument vom Typ \(T_1\) annehmen und wird mit false_operand aufgerufen, das denselben Typ haben muss. Der Typ des zurückgegebenen Werts von true_computation und false_computation muss identisch sein.

Je nach Wert von pred wird nur entweder true_computation oder false_computation ausgeführt.

Conditional(branch_index, branch_computations, branch_operands)

Argumente Typ Semantik
branch_index XlaOp Skalar des Typs S32
branch_computations Sequenz von N XlaComputation Xla-Berechnungen vom Typ \(T_0 \to S , T_1 \to S , ..., T_{N-1} \to S\)
branch_operands Sequenz von N XlaOp Argumente des Typs \(T_0 , T_1 , ..., T_{N-1}\)

Führt branch_computations[branch_index] aus und gibt das Ergebnis zurück. Wenn branch_index ein S32 ist, der < 0 oder >= N ist, wird branch_computations[N-1] als Standardzweig ausgeführt.

Jeder branch_computations[b] muss ein einzelnes Argument vom Typ \(T_b\) annehmen und wird mit branch_operands[b] aufgerufen, die denselben Typ haben müssen. Der Typ des zurückgegebenen Werts für jede branch_computations[b] muss gleich sein.

Abhängig vom Wert von branch_index wird nur eine der branch_computations ausgeführt.

Conv. (Faltung)

Weitere Informationen finden Sie unter XlaBuilder::Conv.

Wie ConvWithGeneralPadding, aber das Padding wird kurz als SAME oder VALID angegeben. Mit dem gleichen Padding wird die Eingabe (lhs) mit Nullen aufgefüllt, sodass die Ausgabe dieselbe Form wie die Eingabe hat, wenn keine Streifen berücksichtigt werden. GÜLTIGES Padding bedeutet einfach kein Auffüllen.

ConvWithGeneralPadding (Faltung)

Weitere Informationen finden Sie unter XlaBuilder::ConvWithGeneralPadding.

Berechnet eine Faltung des in neuronalen Netzwerken verwendeten Typs. Hier kann man sich eine Faltung als n-dimensionales Fenster vorstellen, das sich über eine n-dimensionale Grundfläche bewegt. Für jede mögliche Position des Fensters wird eine Berechnung durchgeführt.

Argumente Typ Semantik
lhs XlaOp n+2-Array von Eingaben einstufen
rhs XlaOp n+2-Array der Kernel-Gewichtungen einstufen
window_strides ArraySlice<int64> n-d-Array von Kernelschritten
padding ArraySlice< pair<int64,int64>> n-d-Array von (niedrig, hoch) Abstand
lhs_dilation ArraySlice<int64> n-d-lhs-Dilationsfaktor-Array
rhs_dilation ArraySlice<int64> n-d-rhs-Dilationsfaktor-Array
feature_group_count int64 Anzahl der Featuregruppen
batch_group_count int64 Anzahl der Batchgruppen

Lass n die Anzahl der räumlichen Dimensionen sein. Das Argument lhs ist ein Rang-n+2-Array, das die Basisfläche beschreibt. Diese wird als Eingabe bezeichnet, auch wenn „rhs“ eine Eingabe ist. In einem neuronalen Netzwerk sind das die Eingabeaktivierungen. Die n+2-Dimensionen sind in dieser Reihenfolge:

  • batch: Jede Koordinate in dieser Dimension stellt eine unabhängige Eingabe dar, für die Faltungen ausgeführt werden.
  • z/depth/features: Jeder (y,x)-Position im Grundbereich ist ein Vektor zugeordnet, der in diese Dimension einfließt.
  • spatial_dims: Beschreibt die räumlichen n-Dimensionen, die die Grundfläche definieren, über die sich das Fenster bewegt.

Das Argument rhs ist ein Rang-n+2-Array, das den Faltfilter/Kernel/Fenster beschreibt. Die Dimensionen werden in dieser Reihenfolge angezeigt:

  • output-z: Die Dimension z der Ausgabe.
  • input-z: Die Größe dieser Dimension multipliziert mit feature_group_count sollte der Größe der Dimension z in lhs entsprechen.
  • spatial_dims: Beschreibt die räumlichen n-Dimensionen, die das n-d-Fenster definieren, das sich über die Grundfläche bewegt.

Das Argument window_strides gibt den Schritt des Faltfensters in den räumlichen Dimensionen an. Wenn der Schritt in der ersten räumlichen Dimension beispielsweise 3 ist, kann das Fenster nur auf Koordinaten platziert werden, bei denen der erste räumliche Index durch 3 teilbar ist.

Das Argument padding gibt den Wert null an, der auf die Basisfläche angewendet werden soll. Der Umfang des Paddings kann negativ sein. Der absolute Wert des negativen Paddings gibt die Anzahl der Elemente an, die vor der Faltung aus der angegebenen Dimension entfernt werden sollen. padding[0] gibt den Abstand für die Dimension y an und padding[1] den Abstand der Dimension x. Jedes Paar hat als erstes Element den niedrigen Abstand und als zweites Element den hohen Abstand. Der niedrige Abstand wird in Richtung der niedrigeren Indizes angewendet, während der hohe Abstand in Richtung der höheren Indizes angewendet wird. Wenn padding[1] beispielsweise (2,3) ist, wird links eine Auffüllung von 2 Nullen und rechts in der zweiten räumlichen Dimension ein Abstand von 3 Nullen erstellt. Das Auffüllen entspricht dem Einfügen derselben Nullwerte in die Eingabe (lhs) vor der Faltung.

Die Argumente lhs_dilation und rhs_dilation geben den Dilatationsfaktor an, der in jeder räumlichen Dimension auf die LHs bzw. RHs angewendet werden soll. Wenn der Dilatationsfaktor in einer räumlichen Dimension d ist, werden d-1 Löcher implizit zwischen den einzelnen Einträgen in dieser Dimension platziert, wodurch das Array vergrößert wird. Die Löcher werden mit einem Nullwert gefüllt, der für die Faltung Nullen bedeutet.

Die Dilatation der RHS wird auch als „verbrechende Faltung“ bezeichnet. Weitere Informationen findest du unter tf.nn.atrous_conv2d. Die Dilatation des LHs wird auch als transponierte Faltung bezeichnet. Weitere Informationen finden Sie unter tf.nn.conv2d_transpose.

Das Argument feature_group_count (Standardwert 1) kann für gruppierte Faltungen verwendet werden. feature_group_count muss ein Divisor der Eingabe- und der Ausgabefeature-Dimension sein. Wenn feature_group_count größer als 1 ist, bedeutet dies, dass die Eingabe- und Ausgabefeature-Dimension sowie die rhs-Ausgabefeature-Dimension konzeptionell gleichmäßig in viele feature_group_count-Gruppen aufgeteilt werden, wobei jede Gruppe aus einer aufeinanderfolgenden Untersequenz von Features besteht. Die Eingabefeature-Dimension von rhs muss der Eingabefeature-Dimension lhs geteilt durch feature_group_count entsprechen (sie hat also bereits die Größe einer Gruppe von Eingabefeatures). Die i-ten Gruppen werden zusammen verwendet, um feature_group_count für viele separate Faltungen zu berechnen. Die Ergebnisse dieser Faltungen werden in der Ausgabefeature-Dimension verkettet.

Für die tiefenfaltbare Faltung wird das Argument feature_group_count auf die Dimension des Eingabefeatures festgelegt und der Filter wird von [filter_height, filter_width, in_channels, channel_multiplier] in [filter_height, filter_width, 1, in_channels * channel_multiplier] umgeformt. Weitere Informationen findest du unter tf.nn.depthwise_conv2d.

Das Argument batch_group_count (Standardwert 1) kann während der Rückpropagierung für gruppierte Filter verwendet werden. batch_group_count muss ein Divisor der Größe der Batchdimension lhs (Eingabe) sein. Wenn batch_group_count größer als 1 ist, bedeutet dies, dass die Batchdimension der Ausgabe die Größe input batch / batch_group_count haben sollte. batch_group_count muss ein Divisor der Größe des Ausgabefeatures sein.

Die Ausgabeform hat folgende Abmessungen in dieser Reihenfolge:

  • batch: Die Größe dieser Dimension multipliziert mit batch_group_count sollte der Größe der batch-Dimension in lhs entsprechen.
  • z: Gleiche Größe wie output-z auf dem Kernel (rhs).
  • spatial_dims: Ein Wert für jede gültige Platzierung des Faltfensters.

Die Abbildung oben zeigt, wie das Feld batch_group_count funktioniert. Wir teilen jeden LH-Batch in batch_group_count-Gruppen auf und machen dasselbe mit den Ausgabefeatures. Dann führen wir für jede dieser Gruppen paarweise Faltungen durch und verketten die Ausgabe entlang der Ausgabefeature-Dimension. Die Betriebssemantik aller anderen Dimensionen (Feature und räumlich) bleibt gleich.

Die gültige Platzierung des Faltfensters wird durch die Schritte und die Größe der Grundfläche nach der Auffüllung bestimmt.

Betrachten Sie eine 2D-Faltung zur Beschreibung der Funktionsweise einer Faltung und wählen Sie in der Ausgabe einige feste batch-, z-, y- und x-Koordinaten aus. Dann ist (y,x) die Position einer Ecke des Fensters innerhalb des Grundbereichs (z.B. die obere linke Ecke, je nachdem, wie Sie die räumlichen Dimensionen interpretieren). Wir haben jetzt ein 2D-Fenster aus dem Grundbereich, in dem jeder 2D-Punkt einem 1D-Vektor zugeordnet ist, sodass wir eine 3D-Box erhalten. Da wir im Convolutional-Kernel die Ausgabekoordinate z korrigiert haben, gibt es auch eine 3D-Box. Die beiden Boxen haben dieselben Abmessungen, sodass wir die Summe der elementweisen Produkte zwischen den beiden Feldern (ähnlich einem Punktprodukt) ermitteln können. Das ist der Ausgabewert.

Wenn output-z z.B. 5 an, dann erzeugt jede Position des Fensters 5 Werte in der Ausgabe für die Dimension z der Ausgabe. Diese Werte unterscheiden sich darin, welcher Teil des Convolutional Kernel verwendet wird. Für jede output-z-Koordinate wird ein separates 3D-Feld mit Werten verwendet. Sie können sich das als fünf separate Faltungen mit jeweils einem anderen Filter vorstellen.

Hier ist Pseudocode für eine 2D-Faltung mit Auffüllung und Abstand:

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

Weitere Informationen finden Sie unter XlaBuilder::ConvertElementType.

Führt ähnlich wie ein elementweises static_cast in C++ einen elementweisen Konvertierungsvorgang von einer Datenform in eine Zielform aus. Die Dimensionen müssen übereinstimmen und die Conversion erfolgt elementweise. So werden s32-Elemente über eine s32-zu-f32-Conversion-Routine zu f32-Elementen.

ConvertElementType(operand, new_element_type)

Argumente Typ Semantik
operand XlaOp Array vom Typ T mit Dimmungen D
new_element_type PrimitiveType Typ U

Die Abmessungen des Operanden und der Zielform müssen übereinstimmen. Die Quell- und Zielelementtypen dürfen keine Tupel sein.

Eine Konvertierung wie T=s32 in U=f32 führt eine normalisierende Int-to-Float-Konvertierungsroutine aus, z. B. „Round-to-Nearest-Even“.

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

Führt AllReduce mit einer Summenberechnung aus.

CustomCall

Weitere Informationen finden Sie unter XlaBuilder::CustomCall.

Rufen Sie eine vom Nutzer bereitgestellte Funktion innerhalb einer Berechnung auf.

CustomCall(target_name, args..., shape)

Argumente Typ Semantik
target_name string Name der Funktion. Es wird eine Aufrufanweisung ausgegeben, die auf diesen Symbolnamen ausgerichtet ist.
args Folge von N XlaOp s N-Argumente beliebigen Typs, die an die Funktion übergeben werden.
shape Shape Ausgabeform der Funktion

Die Funktionssignatur ist dieselbe, unabhängig von der Arity oder dem Typ der Argumente:

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

Angenommen, „CustomCall“ wird wie folgt verwendet:

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

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

Hier ist ein Beispiel für eine Implementierung von 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];
  // ...
}

Die vom Nutzer bereitgestellte Funktion darf keine Nebeneffekte haben und ihre Ausführung muss idempotent sein.

Dot

Weitere Informationen finden Sie unter XlaBuilder::Dot.

Dot(lhs, rhs)

Argumente Typ Semantik
lhs XlaOp Array vom Typ T
rhs XlaOp Array vom Typ T

Die genaue Semantik dieser Operation hängt von den Rängen der Operanden ab:

Eingabe Ausgabe Semantik
Vektor [n] dot Vektor [n] Skalar Vektor-Punktprodukt
Matrix [m x k] dot-Vektor [k] Vektor [m] Matrixvektormultiplikation
Matrix [m x k] dot Matrix [k x n] Matrix [m x n] Matrix-Matrixmultiplikation

Der Vorgang führt die Summe der Produkte über die zweite Dimension von lhs (oder die erste Dimension, wenn sie Rang 1 hat) und die erste Dimension von rhs aus. Dies sind die „vertraglich vereinbarten“ Dimensionen. Die vertraglich vereinbarten Abmessungen von lhs und rhs müssen dieselbe Größe haben. In der Praxis können damit Punktprodukte zwischen Vektoren, Vektor-/Matrixmultiplikationen oder Matrix-/Matrixmultiplikationen durchgeführt werden.

DotGeneral

Weitere Informationen finden Sie unter XlaBuilder::DotGeneral.

DotGeneral(lhs, rhs, dimension_numbers)

Argumente Typ Semantik
lhs XlaOp Array vom Typ T
rhs XlaOp Array vom Typ T
dimension_numbers DotDimensionNumbers Vertrags- und Batchdimensionsnummern

Ähnlich wie bei Dot, allerdings können Vertrags- und Batch-Dimensionsnummern sowohl für lhs als auch für rhs angegeben werden.

DotDimensionNumbers-Felder Typ Semantik
lhs_contracting_dimensions wiederholter int64-Wert lhs kontrastierende Dimensionsnummern
rhs_contracting_dimensions wiederholter int64-Wert rhs kontrastierende Dimensionsnummern
lhs_batch_dimensions wiederholter int64-Wert lhs Batchdimensionsnummern
rhs_batch_dimensions wiederholter int64-Wert rhs Batchdimensionsnummern

DotGeneral führt die Summe der Produkte aus, die die in dimension_numbers angegebenen Abmessungen enthalten.

Verknüpfte vertraglich vereinbarte Dimensionsnummern aus lhs und rhs müssen nicht identisch sein, müssen aber dieselben Dimensionsgrößen haben.

Beispiel mit vertraglich vereinbarten Dimensionszahlen:

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

Verknüpfte Batch-Dimensionsnummern aus lhs und rhs müssen dieselben Dimensionsgrößen haben.

Beispiel mit Batch-Dimensionsnummern (Batchgröße 2, 2 x 2-Matrizen):

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} } }
Eingabe Ausgabe Semantik
[b0, m, k] dot [b0, k, n] [b0, m, n] Batch MatMul
[b0, b1, m, k] dot [b0, b1, k, n] [b0, b1, m, n] Batch MatMul

Daraus folgt, dass die resultierende Dimensionsnummer mit der Batchdimension beginnt, dann mit der Dimension lhs ohne Vertrags-/Batch-Dimension und schließlich mit der Dimension rhs ohne Vertrags- oder Batchdimension.

DynamicSlice

Weitere Informationen finden Sie unter XlaBuilder::DynamicSlice.

DynamicSlice extrahiert ein Teilarray aus dem Eingabearray bei dynamischem start_indices. Die Größe des Kreissegments in jeder Dimension wird in size_indices übergeben, die den Endpunkt exklusiver Segmentintervalle in jeder Dimension angeben: [Start, Start + Größe). Die Form von start_indices muss ein Rang == 1 sein, wobei die Dimensionsgröße dem Rang von operand entspricht.

DynamicSlice(operand, start_indices, size_indices)

Argumente Typ Semantik
operand XlaOp N-dimensionales Array vom Typ T
start_indices Sequenz von N XlaOp Liste von n skalaren Ganzzahlen, die die Startindizes des Segments für jede Dimension enthalten. Wert muss größer oder gleich null sein.
size_indices ArraySlice<int64> Liste von N Ganzzahlen, die die Segmentgröße für jede Dimension enthalten. Jeder Wert muss größer als null sein und „start“ und „size“ müssen kleiner oder gleich der Größe der Dimension sein, damit die Größe der Modulo-Dimension nicht umbrochen wird.

Die effektiven Slice-Indexe werden berechnet, indem die folgende Transformation auf jeden Index-i in [1, N) angewendet wird, bevor das Slice ausgeführt wird:

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

Dadurch wird sichergestellt, dass das extrahierte Segment in Bezug auf das Operandenarray immer innerhalb der Grenzen liegt. Wenn das Slice innerhalb des Bereichs liegt, bevor die Transformation angewendet wird, hat die Transformation keine Auswirkungen.

Beispiel für eindimensionales Element:

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

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

Beispiel für zweidimensionale Darstellung:

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

Weitere Informationen finden Sie unter XlaBuilder::DynamicUpdateSlice.

DynamicUpdateSlice generiert ein Ergebnis, das dem Wert des Eingabearrays operand entspricht, wobei das Segment update bei start_indices überschrieben wird. Die Form von update bestimmt die Form des Teilarrays des Ergebnisses, das aktualisiert wird. Die Form von start_indices muss ein Rang == 1 sein, wobei die Dimensionsgröße dem Rang von operand entspricht.

DynamicUpdateSlice(operand, update, start_indices)

Argumente Typ Semantik
operand XlaOp N-dimensionales Array vom Typ T
update XlaOp N-dimensionales Array vom Typ T, das die Segmentaktualisierung enthält Jede Dimension der Aktualisierungsform muss größer als null sein und Start + Aktualisierung muss kleiner oder gleich der Operandengröße für jede Dimension sein, um zu vermeiden, dass Aktualisierungsindexe außerhalb des gültigen Bereichs liegen.
start_indices Sequenz von N XlaOp Liste von n skalaren Ganzzahlen, die die Startindizes des Segments für jede Dimension enthalten. Wert muss größer oder gleich null sein.

Die effektiven Slice-Indexe werden berechnet, indem die folgende Transformation auf jeden Index-i in [1, N) angewendet wird, bevor das Slice ausgeführt wird:

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

Dadurch wird sichergestellt, dass das aktualisierte Slice immer innerhalb des Operandenarrays liegt. Wenn das Slice innerhalb des Bereichs liegt, bevor die Transformation angewendet wird, hat die Transformation keine Auswirkungen.

Beispiel für eindimensionales Element:

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}

Beispiel für zweidimensionale Darstellung:

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

Elementweise binäre Arithmetische Operationen

Weitere Informationen finden Sie unter XlaBuilder::Add.

Es wird eine Reihe von elementweisen binären arithmetischen Operationen unterstützt.

Op(lhs, rhs)

Dabei ist Op entweder Add (Addition), Sub (Subtraktion), Mul (Multiplikation), Div (Division), Rem (Rest), Max (maximal), Min (Minimum), LogicalAnd (logisches UND) oder LogicalOr (logisches ODER).

Argumente Typ Semantik
lhs XlaOp Linker Operand auf der linken Seite: Array vom Typ T
rhs XlaOp right-hand-side operand: Array vom Typ T

Die Formen der Argumente müssen entweder ähnlich oder kompatibel sein. Weitere Informationen zur Kompatibilität von Formen finden Sie in der Broadcasting-Dokumentation. Das Ergebnis einer Operation hat eine Form, die das Ergebnis der Übertragung der beiden Eingabearrays ist. Bei dieser Variante werden Operationen zwischen Arrays unterschiedlicher Ränge nicht unterstützt, es sei denn, einer der Operanden ist ein Skalar.

Wenn Op Rem ist, wird das Vorzeichen des Ergebnisses vom Dividenden entnommen, und der absolute Wert des Ergebnisses ist immer kleiner als der absolute Wert des Divisors.

Überlauf der Ganzzahl-Division (vorzeichenbehaftete/unsignierte Division/Rest durch null oder vorzeichenbehaftete Division/Rest von INT_SMIN mit -1) führt zu einem für die Implementierung definierten Wert.

Für diese Vorgänge gibt es eine alternative Variante, die Übertragungen unterschiedlicher Rangfolge unterstützt:

Op(lhs, rhs, broadcast_dimensions)

Dabei ist Op mit dem oben identischen Wert identisch. Diese Variante der Operation sollte für arithmetische Operationen zwischen Arrays unterschiedlicher Ränge verwendet werden, z. B. für das Hinzufügen einer Matrix zu einem Vektor.

Der zusätzliche Operand broadcast_dimensions ist ein Segment von Ganzzahlen, mit dem der Rang des Operanden des niedrigeren Rangs auf den Rang des Operanden des höheren Rangs erweitert wird. broadcast_dimensions ordnet die Dimensionen der Form mit niedrigerem Rang den Dimensionen der Form mit höherem Rang zu. Die nicht zugeordneten Maße der maximierten Form werden mit Maßen der Größe 1 gefüllt. Beim Broadcasting nach Dimensionierung werden dann die Formen entlang dieser Degenerierungsdimensionen gesendet, um die Formen beider Operanden anzugleichen. Die Semantik wird auf der Broadcasting-Seite ausführlich beschrieben.

Elementweise Vergleichsvorgänge

Weitere Informationen finden Sie unter XlaBuilder::Eq.

Es wird eine Reihe standardmäßiger elementweiser binärer Vergleichsvorgänge unterstützt. Beachten Sie, dass beim Vergleich von Gleitkommatypen die standardmäßige IEEE 754-Gleitkommasemantik für den Vergleich gilt.

Op(lhs, rhs)

Dabei ist Op einer der folgenden Werte: Eq (gleich), Ne (nicht gleich), Ge (größer oder gleich), Gt (größer als), Le (kleiner oder gleich) oder Lt (kleiner als). Eine weitere Gruppe von Operatoren, EqTotalOrder, NeTotalOrder, GeTotalOrder, GtTotalOrder, LeTotalOrder und LtTotalOrder und LtTotalOrder, bieten dieselben Funktionen, mit der Ausnahme, dass sie zusätzlich eine Gesamtreihenfolge über den Gleitkommazahlen unterstützen, indem sie -NaN < -Inf < -Finite < -0 < +0 < +Finite <N +Inf.

Argumente Typ Semantik
lhs XlaOp Linker Operand auf der linken Seite: Array vom Typ T
rhs XlaOp right-hand-side operand: Array vom Typ T

Die Formen der Argumente müssen entweder ähnlich oder kompatibel sein. Weitere Informationen zur Kompatibilität von Formen finden Sie in der Broadcasting-Dokumentation. Das Ergebnis eines Vorgangs hat eine Form. Dies ist das Ergebnis der Übertragung der beiden Eingabearrays mit dem Elementtyp PRED. Bei dieser Variante werden Vorgänge zwischen Arrays unterschiedlicher Ränge nicht unterstützt, es sei denn, einer der Operanden ist ein Skalar.

Für diese Vorgänge gibt es eine alternative Variante, die Übertragungen unterschiedlicher Rangfolge unterstützt:

Op(lhs, rhs, broadcast_dimensions)

Dabei ist Op mit dem oben identischen Wert identisch. Diese Variante des Vorgangs sollte für Vergleichsvorgänge zwischen Arrays unterschiedlicher Ränge verwendet werden, z. B. zum Hinzufügen einer Matrix zu einem Vektor.

Der zusätzliche Operand broadcast_dimensions ist ein Segment von Ganzzahlen, mit denen die Dimensionen festgelegt werden, die für die Übertragung der Operanden verwendet werden sollen. Die Semantik wird auf der Broadcasting-Seite ausführlich beschrieben.

Elementweise unäre Funktionen

XlaBuilder unterstützt die folgenden elementweisen unären Funktionen:

Abs(operand) Elementweises ABS x -> |x|.

Ceil(operand) Elementweises Maximum x -> ⌈x⌉.

Cos(operand) Elementweiser Kosinus x -> cos(x).

Exp(operand) Elementweises natürliches exponentielles x -> e^x.

Floor(operand) Elementweiser Mindestbetrag: x -> ⌊x⌋.

Imag(operand) Elementweiser imaginärer Teil einer komplexen (oder reellen) Form. x -> imag(x). Wenn der Operand ein Gleitkommatyp ist, wird 0 zurückgegeben.

IsFinite(operand) Prüft, ob jedes Element von operand endlich, also nicht positiv oder negativ unendlich und nicht NaN ist. Gibt ein Array von PRED-Werten mit derselben Form wie die Eingabe zurück, wobei jedes Element nur dann true ist, wenn das entsprechende Eingabeelement endlich ist.

Log(operand) Elementweiser natürlicher Logarithmus x -> ln(x).

LogicalNot(operand) Elementweise logisch, nicht x -> !(x).

Logistic(operand) Berechnung der elementweisen logistischen Funktion x -> logistic(x).

PopulationCount(operand) Berechnet die Anzahl der Bits, die in den einzelnen Elementen von operand festgelegt sind.

Neg(operand) Elementweise Negation x -> -x.

Real(operand) Elementweiser reeller Teil einer komplexen (oder reellen) Form. x -> real(x). Wenn der Operand ein Gleitkommatyp ist, wird derselbe Wert zurückgegeben.

Rsqrt(operand) Elementweise Reziproz der Quadratwurzeloperation x -> 1.0 / sqrt(x).

Sign(operand) Vorgang für elementweises Zeichen x -> sgn(x), wobei

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

mit dem Vergleichsoperator des Elementtyps operand.

Sqrt(operand) Elementweise Quadratwurzel-Operation x -> sqrt(x).

Cbrt(operand) Elementweiser kubischer Stammvorgang x -> cbrt(x).

Tanh(operand) Elementweiser hyperbolischer Tangens x -> tanh(x).

Round(operand) Elementweise Rundung, gleicht sich von null ab.

RoundNearestEven(operand) Elementweise Rundung, wird auf die nächste Gerade ausgerichtet.

Argumente Typ Semantik
operand XlaOp Operand der Funktion

Die Funktion wird auf jedes Element im Array operand angewendet, sodass ein Array dieselbe Form hat. operand darf ein Skalar sein (Rang 0).

FFT

Der XLA-FFT-Vorgang implementiert die Vorwärts- und Inverse Fourier-Transformationen für reale und komplexe Ein-/Ausgaben. Mehrdimensionale FFTs auf bis zu drei Achsen werden unterstützt.

Weitere Informationen finden Sie unter XlaBuilder::Fft.

Argumente Typ Semantik
operand XlaOp Das Array, das mit der Fourier-Transformation transformiert wird.
fft_type FftType Siehe die folgende Tabelle.
fft_length ArraySlice<int64> Die Zeitdomainlängen der transformierten Achsen. Dies ist insbesondere erforderlich, damit IRFFT die Größe der innersten Achse richtig anpasst, da RFFT(fft_length=[16]) dieselbe Ausgabeform wie RFFT(fft_length=[17]) hat.
FftType Semantik
FFT Komplexe bis komplexe FFT weiterleiten. Die Form bleibt unverändert.
IFFT Umgekehrte komplex-komplexe FFT. Die Form bleibt unverändert.
RFFT Leiten Sie eine reale zu einer komplexen FFT weiter. Die Form der innersten Achse wird auf fft_length[-1] // 2 + 1 reduziert, wenn fft_length[-1] ein Wert ungleich null ist, wobei der umgekehrte konjugierte Teil des transformierten Signals über die Nyquist-Frequenz hinaus weggelassen wird.
IRFFT Umgekehrte reelle zu komplexe FFT (z.B. mit komplex, gibt reelle Zahlen zurück). Die Form der innersten Achse wird zu fft_length[-1] erweitert, wenn fft_length[-1] ein Wert ungleich null ist. Dabei wird der Teil des transformierten Signals jenseits der Nyquist-Frequenz aus der umgekehrten Konjugation der 1- zu fft_length[-1] // 2 + 1-Einträge abgeleitet.

Mehrdimensionale FFT

Wenn mehr als 1 fft_length angegeben ist, entspricht dies dem Anwenden einer Kaskade von FFT-Vorgängen auf jede der innersten Achsen. Beachten Sie, dass bei den reellen und komplexsten realen Fällen die Transformation der innersten Achse (effektiv) zuerst ausgeführt wird (RFFT; letzte für IRFFT). Deshalb ändert sich die Größe der innersten Achse. Andere Achsentransformationen werden dann komplex->komplex.

Implementierungsdetails

Die CPU-FFT wird vom Eigens TensorFFT unterstützt. GPU-FFT verwendet cuFFT.

Zusammentragen

Mit dem XLA-Vorgang für das Erfassen werden mehrere Segmente eines Eingabearrays zusammengefügt (jedes Segment mit einem möglicherweise anderen Laufzeitversatz).

Allgemeine Semantik

Weitere Informationen finden Sie unter XlaBuilder::Gather. Eine intuitivere Beschreibung finden Sie unten im Abschnitt „Informelle Beschreibung“.

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

Argumente Typ Semantik
operand XlaOp Das Array, aus dem wir Daten sammeln.
start_indices XlaOp Array mit den Startindexen der gesammelten Segmente
index_vector_dim int64 Die Dimension in start_indices, die die Startindexe „enthält“. Unten finden Sie eine detaillierte Beschreibung.
offset_dims ArraySlice<int64> Der Satz von Dimensionen in der Ausgabeform, die zu einem Array verschoben werden, unterteilt vom Operanden.
slice_sizes ArraySlice<int64> slice_sizes[i] sind die Grenzen für das Slice der Dimension i.
collapsed_slice_dims ArraySlice<int64> Die Maße in jedem Kreissegment, die minimiert werden. Diese Abmessungen müssen die Größe 1 haben.
start_index_map ArraySlice<int64> Eine Map, die beschreibt, wie Indizes in start_indices den rechtlichen Indizes in Operanden zugeordnet werden.
indices_are_sorted bool Gibt an, ob die Indizes garantiert vom Aufrufer sortiert werden.

Der Einfachheit halber kennzeichnen wir Dimensionen im Ausgabearray, nicht in offset_dims, als batch_dims.

Die Ausgabe ist ein Array mit dem Rang batch_dims.size + offset_dims.size.

operand.rank muss der Summe von offset_dims.size und collapsed_slice_dims.size entsprechen. Außerdem muss slice_sizes.size gleich operand.rank sein.

Wenn index_vector_dim gleich start_indices.rank ist, betrachten wir implizit start_indices als eine nachgestellte 1-Dimension. Wenn also start_indices die Form [6,7] hatte und index_vector_dim 2 ist, wird die Form von start_indices implizit als [6,7,1] betrachtet.

Die Grenzen für das Ausgabearray entlang der Dimension i werden so berechnet:

  1. Wenn i in batch_dims vorhanden ist (d. h. für einige k gleich batch_dims[k] ist), wählen wir die entsprechenden Dimensionsgrenzen aus start_indices.shape aus und überspringen index_vector_dim. Wählen Sie also start_indices.shape.dims[k] aus, wenn k < index_vector_dim und ansonsten start_indices.shape.dims[k+1].

  2. Wenn i in offset_dims vorhanden ist (d. h. für einige k gleich offset_dims[k]), wird die entsprechende Grenze von slice_sizes ausgewählt, nachdem collapsed_slice_dims berücksichtigt wurde. Das heißt, wir wählen adjusted_slice_sizes[k] aus, wobei adjusted_slice_sizes gleich slice_sizes ist und die Grenzen der Indexe collapsed_slice_dims entfernt wurden.

Formal wird der Operandenindex In, der einem bestimmten Ausgabeindex Out entspricht, so berechnet:

  1. Lass G = { Out[k] für k in batch_dims }. Verwende G, um einen Vektor-S so herauszuschneiden, dass S[i] = start_indices[Kombination(G, i)], wobei Kombinieren(A, b) b an Position index_vector_dim in A einfügt. Dies ist auch dann gut definiert, wenn G leer ist: Wenn G leer ist, dann S = start_indices.

  2. Erstellen Sie mit S einen Startindex Sin in operand. Dazu verteilen Sie S mit start_index_map. Genauer gesagt:

    1. Sin[start_index_map[k]] = S[k], wenn k < start_index_map.size.

    2. SAndernfalls in[_] = 0.

  3. Erstellen Sie einen Index Oin in operand, indem Sie die Indexe an den Offset-Dimensionen in Out gemäß dem Satz collapsed_slice_dims verteilen. Genauer gesagt:

    1. Oin[remapped_offset_dims(k)] = Out[offset_dims[k]], wenn k < offset_dims.size (remapped_offset_dims unten definiert).

    2. OAndernfalls in[_] = 0.

  4. In ist Oin + Sin, wobei + für das elementweise Addieren steht.

remapped_offset_dims ist eine monotone Funktion mit der Domain [0, offset_dims.size) und dem Bereich [0, operand.rank) \ collapsed_slice_dims. Wenn z.B. offset_dims.size ist 4, operand.rank ist 6 und collapsed_slice_dims ist {0, 2}, dann ist remapped_offset_dims {01, 13, 24, 35}.

Wenn indices_are_sorted auf „true“ gesetzt ist, kann XLA davon ausgehen, dass start_indices (in aufsteigender start_index_map-Reihenfolge) nach dem Nutzer sortiert werden. Andernfalls ist die Semantik Implementierung definiert.

Informelle Beschreibung und Beispiele

Informell entspricht jede Index-Out im Ausgabearray einem Element E im Operandenarray, das so berechnet wird:

  • Wir verwenden die Batchdimensionen in Out, um einen Startindex aus start_indices zu ermitteln.

  • Wir verwenden start_index_map, um den Startindex, dessen Größe kleiner als operand.rank sein kann, einem "vollständigen" Startindex in der operand zuzuordnen.

  • Ein Segment mit der Größe slice_sizes wird anhand des vollständigen Startindex dynamisch zerlegt.

  • Wir ändern die Form des Slice, indem wir die collapsed_slice_dims-Dimensionen minimieren. Da alle Abmessungen des minimierten Segments eine Grenze von 1 haben müssen, ist diese Umformung immer zulässig.

  • Wir verwenden die Offset-Dimensionen in Out, um in dieses Segment zu indexieren und das Eingabeelement E zu erhalten, das dem Ausgabeindex Out entspricht.

index_vector_dim ist in allen folgenden Beispielen auf start_indices.rank1 gesetzt. Interessantere Werte für index_vector_dim ändern den Vorgang nicht grundlegend, machen die visuelle Darstellung jedoch umständlicher.

Sehen wir uns ein Beispiel an, bei dem fünf Teile der Form [8,6] aus einem [16,11]-Array gesammelt werden, um eine Vorstellung davon zu bekommen, wie all dies zusammenpasst. Die Position eines Segments im [16,11]-Array kann als Indexvektor der Form S64[2] dargestellt werden, sodass die Gruppe von fünf Positionen als S64[5,2]-Array dargestellt werden kann.

Das Verhalten des collect-Vorgangs kann dann als Indextransformation dargestellt werden, die [G, O0, O1], einen Index in der Ausgabeform, annimmt und auf folgende Weise einem Element im Eingabearray zuordnet:

Zunächst wählen wir mit G einen (X,Y) Vektor aus dem Array zum Erfassen der Indexe aus. Das Element im Ausgabearray bei Index [G,O0,O1] ist dann das Element im Eingabearray an Index [X+O0,Y+O1].

slice_sizes ist [8,6], wodurch der Bereich von O0 und O1 sowie dies wiederum die Grenzen des Slice bestimmt.

Dieser Erfassungsvorgang fungiert als dynamisches Batchsegment mit G als Batchdimension.

Die Sammel-Indizes können mehrdimensional sein. Eine allgemeinere Version des obigen Beispiels, bei der ein Array vom Typ „Indexe sammeln“ der Form [4,5,2] verwendet wird, würde Indexe so übertragen:

Dies fungiert wiederum als dynamisches Batch-Slice G0 und G1 als Batch-Dimensionen. Die Segmentgröße beträgt immer noch [8,6].

Der Collect-Vorgang in XLA verallgemeinert die oben beschriebene informelle Semantik auf folgende Arten:

  1. Wir können konfigurieren, welche Abmessungen in der Ausgabeform die Versatzabmessungen sind (die im letzten Beispiel O0, O1 enthalten). Die Ausgabedimensionen im Batch (die im letzten Beispiel G0 und im letzten Beispiel G1 enthalten) sind als Ausgabedimensionen definiert, die keine Offsetdimensionen sind.

  2. Die Anzahl der Ausgabe-Offset-Dimensionen, die explizit in der Ausgabeform vorhanden sind, kann kleiner als der Eingaberang sein. Diese „fehlenden“ Dimensionen, die explizit als collapsed_slice_dims aufgeführt sind, müssen eine Segmentgröße von 1 haben. Da sie eine Slice-Größe von 1 haben, ist der einzige gültige Index für sie 0. Das Entfernen dieser Segmente führt nicht zu Mehrdeutigkeit.

  3. Das aus dem Array "collect Indices" extrahierte Segment (im letzten Beispiel X, Y) hat möglicherweise weniger Elemente als der Eingabearrayrang. Eine explizite Zuordnung gibt außerdem vor, wie der Index erweitert werden soll, damit er denselben Rang wie die Eingabe hat.

Als letztes Beispiel verwenden wir (2) und (3) zur Implementierung von tf.gather_nd:

G0 und G1 werden wie gewohnt verwendet, um einen Startindex aus dem Array für die Erfassung der Indexe wie gewohnt zu segmentieren. Der Startindex hat jedoch nur ein Element, X. Ebenso gibt es nur einen Ausgabe-Offset-Index mit dem Wert O0. Bevor sie jedoch als Indexe in das Eingabearray verwendet werden, werden diese gemäß „Collect Index Mapping“ (start_index_map in der formalen Beschreibung) und „Offset Mapping“ (remapped_offset_dims in der formalen Beschreibung) in [X,0] bzw. [0,O0] erweitert, was insgesamt [X,O0] ergibt. Mit anderen Worten: Der Ausgabeindex [O0] wird für den Index [O1], also den Index [O1, {1] und {1 den Ausgabeindex [1], {1 (, [0], und , .1, [G], [1], [], [1], [1], [], [G], [1], [], [1], [0], [1], [1], [], [1], {1,} [0] enthält.10GGGatherIndicestf.gather_nd

slice_sizes ist für diesen Fall [1,11]. Das bedeutet, dass jeder Index-X im Array zum Erfassen der Indexe eine ganze Zeile auswählt und das Ergebnis die Verkettung all dieser Zeilen ist.

GetDimensionSize

Weitere Informationen finden Sie unter XlaBuilder::GetDimensionSize.

Gibt die Größe der gegebenen Dimension des Operanden zurück. Der Operand muss Arrayform sein.

GetDimensionSize(operand, dimension)

Argumente Typ Semantik
operand XlaOp n-dimensionales Eingabearray
dimension int64 Ein Wert im Intervall [0, n), der die Dimension angibt

SetDimensionSize

Weitere Informationen finden Sie unter XlaBuilder::SetDimensionSize.

Legt die dynamische Größe der in XlaOp angegebenen Dimension fest. Der Operand muss Arrayform sein.

SetDimensionSize(operand, size, dimension)

Argumente Typ Semantik
operand XlaOp n-dimensionalen Eingabearray.
size XlaOp int32, die die dynamische Laufzeitgröße darstellt.
dimension int64 Ein Wert im Intervall [0, n), der die Dimension angibt.

Sie übergeben den Operanden als Ergebnis, wobei die dynamische Dimension vom Compiler erfasst wird.

Aufgefüllte Werte werden von nachgelagerten Reduktionsvorgängen ignoriert.

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

Weitere Informationen finden Sie unter XlaBuilder::GetTupleElement.

Indexiert ein Tupel mit einem Wert für die Compile-Zeitkonstante.

Der Wert muss eine Compile-Zeitkonstante sein, damit die Forminferenz den Typ des Ergebniswerts bestimmen kann.

Dies entspricht std::get<int N>(t) in C++. Konzeptionell:

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.

Weitere Informationen finden Sie unter tf.tuple.

Einspeisung

Weitere Informationen finden Sie unter XlaBuilder::Infeed.

Infeed(shape)

Argument Typ Semantik
shape Shape Form der aus der In-Feed-Oberfläche gelesenen Daten. Das Layoutfeld der Form muss so festgelegt werden, dass es dem Layout der an das Gerät gesendeten Daten entspricht. Andernfalls ist sein Verhalten nicht definiert.

Liest ein einzelnes Datenelement aus der impliziten Infeed-Streaming-Oberfläche des Geräts, interpretiert die Daten als gegebene Form und ihr Layout und gibt eine XlaOp der Daten zurück. In einer Berechnung sind mehrere Infeed-Vorgänge zulässig, unter den Infeed-Vorgängen muss aber eine Gesamtreihenfolge vorhanden sein. Zum Beispiel haben zwei Infeeds im folgenden Code eine Gesamtreihenfolge, da zwischen den while-Schleifen eine Abhängigkeit besteht.

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

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

Verschachtelte Tupelformen werden nicht unterstützt. Bei einer leeren Tupelform ist der Infeed-Vorgang praktisch ein Leerlauf und wird fortgesetzt, ohne dass Daten aus der Einspeisung des Geräts gelesen werden.

Iota

Weitere Informationen finden Sie unter XlaBuilder::Iota.

Iota(shape, iota_dimension)

Erstellt ein konstantes Literal auf dem Gerät anstelle einer potenziell großen Hostübertragung. Erstellt ein Array mit einer angegebenen Form und Werten, die bei null beginnen und entlang der angegebenen Dimension um eins erhöht werden. Bei Gleitkommatypen entspricht das erstellte Array ConvertElementType(Iota(...)), wobei Iota ein ganzzahliger Typ ist und die Konvertierung in den Gleitkommatyp erfolgt.

Argumente Typ Semantik
shape Shape Form des durch Iota() erstellten Arrays
iota_dimension int64 Dimension, mit der die Inkrementierung erfolgt.

Beispiel: Iota(s32[4, 8], 0) gibt Folgendes zurück:

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

Iota(s32[4, 8], 1) für Retouren

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

Karte

Weitere Informationen finden Sie unter XlaBuilder::Map.

Map(operands..., computation)

Argumente Typ Semantik
operands Folge von N XlaOp s N Arrays der Typen T0..T{N-1}
computation XlaComputation Berechnung des Typs T_0, T_1, .., T_{N + M -1} -> S mit N Parametern des Typs T und M eines beliebigen Typs
dimensions int64-Array Array von Kartenmaßen

Wendet eine Skalarfunktion über die gegebenen operands-Arrays an und ergibt ein Array derselben Dimensionen, wobei jedes Element das Ergebnis der zugeordneten Funktion ist, die auf die entsprechenden Elemente in den Eingabearrays angewendet wird.

Die zugeordnete Funktion ist eine beliebige Berechnung mit der Einschränkung, dass sie n Eingaben vom skalaren Typ T und eine einzelne Ausgabe vom Typ S hat. Die Ausgabe hat dieselben Maße wie die Operanden, mit der Ausnahme, dass der Elementtyp T durch S ersetzt wird.

Beispiel: Map(op1, op2, op3, computation, par1) ordnet elem_out <- computation(elem1, elem2, elem3, par1) bei jedem (mehrdimensionalen) Index in den Eingabearrays zu, um das Ausgabearray zu erzeugen.

OptimizationBarrier

Blockiert alle Optimierungsübergänge, sodass Berechnungen über die Barriere hinweg übergangen werden.

Sorgt dafür, dass alle Eingaben vor Operatoren ausgewertet werden, die von den Ausgaben der Barriere abhängig sind.

Pad

Weitere Informationen finden Sie unter XlaBuilder::Pad.

Pad(operand, padding_value, padding_config)

Argumente Typ Semantik
operand XlaOp Array vom Typ T
padding_value XlaOp Skalar vom Typ T, um den zusätzlichen Abstand zu füllen
padding_config PaddingConfig Abstand an beiden Rändern (niedrig, hoch) und zwischen den Elementen jeder Dimension

Maximiert das angegebene Array operand durch Auffüllung um das Array sowie zwischen den Elementen des Arrays mit dem angegebenen padding_value. padding_config gibt den Rand- und Innenabstand für jede Dimension an.

PaddingConfig ist ein wiederkehrendes Feld von PaddingConfigDimension, das drei Felder für jede Dimension enthält: edge_padding_low, edge_padding_high und interior_padding.

edge_padding_low und edge_padding_high geben den Abstand an, der am unteren Ende (neben Index 0) bzw. am oberen Rand (neben dem höchsten Index) jeder Dimension hinzugefügt wird. Der Umfang des Rands kann negativ sein – der absolute Wert gibt die Anzahl der Elemente an, die aus der angegebenen Dimension entfernt werden sollen.

interior_padding gibt den Abstand zwischen zwei Elementen in jeder Dimension an. Er darf nicht negativ sein. Der Innenrand geschieht logisch vor dem Randabstand, sodass bei einem negativen Randabstand die Elemente aus dem Operanden mit Auffüllung entfernt werden.

Diese Operation ist ein Nullbefehl, wenn die Paare für den Randabstand alle (0, 0) und die Werte für das innere Auffüllen alle 0 sind. Die folgende Abbildung zeigt Beispiele für verschiedene edge_padding- und interior_padding-Werte für ein zweidimensionales Array.

Empfang

Weitere Informationen finden Sie unter XlaBuilder::Recv.

Recv(shape, channel_handle)

Argumente Typ Semantik
shape Shape Form der zu empfangenden Daten
channel_handle ChannelHandle eindeutige Kennung für jedes Sende-/Empfangspaar

Empfängt Daten der angegebenen Form von einer Send-Anweisung in einer anderen Berechnung, die denselben Kanal-Handle verwendet. Gibt eine XlaOp für die empfangenen Daten zurück.

Die Client API des Recv-Vorgangs stellt die synchrone Kommunikation dar. Die Anweisung wird jedoch intern in zwei HLO-Anweisungen (Recv und RecvDone) aufgeteilt, um asynchrone Datenübertragungen zu ermöglichen. Weitere Informationen finden Sie unter HloInstruction::CreateRecv und HloInstruction::CreateRecvDone.

Recv(const Shape& shape, int64 channel_id)

Weist Ressourcen zu, die zum Empfangen von Daten aus einer Send-Anweisung mit derselben channel_id erforderlich sind. Gibt einen Kontext für die zugewiesenen Ressourcen zurück, der von der folgenden RecvDone-Anweisung verwendet wird, um auf den Abschluss der Datenübertragung zu warten. Der Kontext ist ein Tupel von {Empfangspuffer (Form), Anfragekennung (U32)} und kann nur von einer RecvDone-Anweisung verwendet werden.

RecvDone(HloInstruction context)

Bei einem von einer Recv-Anweisung erstellten Kontext wird auf den Abschluss der Datenübertragung gewartet und gibt die empfangenen Daten zurück.

Einschränken

Weitere Informationen finden Sie unter XlaBuilder::Reduce.

Wendet eine Reduktionsfunktion gleichzeitig auf ein oder mehrere Arrays an.

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

Argumente Typ Semantik
operands Sequenz von N XlaOp N Arrays vom Typ T_0, ..., T_{N-1}
init_values Sequenz von N XlaOp N Skalare vom Typ T_0, ..., T_{N-1}.
computation XlaComputation Berechnung des Typs T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}).
dimensions int64-Array ein ungeordnetes Array mit Dimensionen auf, das reduziert werden soll.

Wobei:

  • N muss größer oder gleich 1 sein.
  • Die Berechnung muss „ungefähr“ assoziativ sein (siehe unten).
  • Alle Eingabearrays müssen dieselben Abmessungen haben.
  • Alle Anfangswerte müssen unter computation eine Identität bilden.
  • Wenn N = 1, ist Collate(T) T.
  • Bei N > 1 ist Collate(T_0, ..., T_{N-1}) ein Tupel von N-Elementen vom Typ T.

Durch diesen Vorgang werden eine oder mehrere Dimensionen jedes Eingabearrays in Skalare reduziert. Der Rang jedes zurückgegebenen Arrays ist rank(operand) - len(dimensions). Die Ausgabe des Vorgangs ist Collate(Q_0, ..., Q_N), wobei Q_i ein Array vom Typ T_i ist, dessen Dimensionen unten beschrieben werden.

Die Berechnung der Reduzierung kann über verschiedene Back-Ends neu zugeordnet werden. Dies kann zu numerischen Unterschieden führen, da einige Reduktionsfunktionen wie das Addieren nicht assoziativ für Gleitkommazahlen sind. Wenn der Datenbereich jedoch begrenzt ist, reicht das Hinzufügen von Gleitkommazahlen für die meisten praktischen Zwecke aus, um assoziativ zu sein.

Beispiele

Wenn Sie eine Dimension in einem einzelnen 1D-Array mit den Werten [10, 11, 12, 13] und der Reduktionsfunktion f (hier computation) reduzieren, kann dies so berechnet werden:

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

aber es gibt noch viele andere Möglichkeiten, z.B.

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

Im Folgenden finden Sie ein grobes Pseudocode-Beispiel dafür, wie eine Reduktion implementiert werden könnte, wobei die Summierung als Reduktionsberechnung mit einem Anfangswert von 0 verwendet wird.

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]

Hier ein Beispiel für die Reduzierung eines 2D-Arrays (Matrix). Die Form hat den Rang 2, die Dimension 0 der Größe 2 und die Dimension 1 der Größe 3:

Ergebnisse der Reduzierung von Dimensionen 0 oder 1 mit einer „add“-Funktion:

Beide Reduktionsergebnisse sind 1D-Arrays. Das Diagramm zeigt eine als Spalte und eine als Zeile nur zur besseren Übersicht.

Ein komplexeres Beispiel ist hier ein 3D-Array. Ihr Rang ist 3, Dimension 0 von Größe 4, Dimension 1 von Größe 2 und Dimension 2 von Größe 3. Der Einfachheit halber werden die Werte 1 bis 6 in der Dimension 0 repliziert.

Ähnlich wie im 2D-Beispiel können wir nur eine Dimension reduzieren. Wenn wir beispielsweise die Dimension 0 reduzieren, erhalten wir ein Rang-2-Array, bei dem alle Werte der Dimension 0 zu einem Skalar zusammengefasst wurden:

|  4   8  12 |
| 16  20  24 |

Wenn wir Dimension 2 reduzieren, erhalten wir auch ein Rang-2-Array, bei dem alle Werte der Dimension 2 zu einem Skalar zusammengefasst wurden:

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

Die relative Reihenfolge zwischen den verbleibenden Dimensionen in der Eingabe wird in der Ausgabe beibehalten, einigen Dimensionen werden jedoch möglicherweise neue Zahlen zugewiesen, da sich der Rang ändert.

Wir können auch mehrere Dimensionen reduzieren. Aus den Dimensionen 0 und 1 ergibt sich das 1D-Array [20, 28, 36].

Wenn das 3D-Array über alle Dimensionen hinweg reduziert wird, wird das skalare 84-Objekt erstellt.

Variadische Reduzierung

Wenn N > 1, ist die Funktion „Reduce“ etwas komplexer, da sie gleichzeitig auf alle Eingaben angewendet wird. Die Operanden werden in der folgenden Reihenfolge an die Berechnung übergeben:

  • Reduzierter Wert für ersten Operanden wird ausgeführt
  • ...
  • Reduzierter Wert für den n-ten Operanden wird ausgeführt
  • Eingabewert für den ersten Operanden
  • ...
  • Eingabewert für den n-ten Operanden

Betrachten Sie beispielsweise die folgende Reduktionsfunktion, mit der der Maximalwert und argmax eines 1D-Arrays parallel berechnet werden können:

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

Bei den 1D-Eingabearrays V = Float[N], K = Int[N] und den Initialisierungswerten I_V = Float, I_K = Int entspricht das Ergebnis f_(N-1) der Reduzierung über die einzige Eingabedimension der folgenden rekursiven Anwendung:

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

Wenn Sie diese Reduzierung auf ein Array von Werten und ein Array von sequenziellen Indizes (iota) anwenden, wird dies über die Arrays hinweg ausgeführt und gibt ein Tupel zurück, das den höchsten Wert und den übereinstimmenden Index enthält.

ReducePrecision

Weitere Informationen finden Sie unter XlaBuilder::ReducePrecision.

Modelliert die Auswirkungen der Konvertierung von Gleitkommawerten in ein Format mit geringerer Genauigkeit (z. B. IEEE-FP16) und zurück in das ursprüngliche Format. Die Anzahl der Exponenten- und Mantissen-Bits im Format mit geringerer Genauigkeit kann beliebig angegeben werden, auch wenn möglicherweise nicht alle Bitgrößen in allen Hardwareimplementierungen unterstützt werden.

ReducePrecision(operand, mantissa_bits, exponent_bits)

Argumente Typ Semantik
operand XlaOp Array des Gleitkommatyps T.
exponent_bits int32 Anzahl der Exponentenbits im Format mit geringerer Genauigkeit
mantissa_bits int32 Anzahl der Mantissen-Bits im Format mit geringerer Genauigkeit

Das Ergebnis ist ein Array vom Typ T. Die Eingabewerte werden auf den nächsten Wert gerundet, der mit der gegebenen Anzahl von Mantissen-Bits (unter Verwendung der gleichmäßigen Semantik) dargestellt werden kann. Alle Werte, die den durch die Anzahl der Exponentenbits angegebenen Bereich überschreiten, werden auf positive oder negative unendliche Werte gesetzt. NaN-Werte bleiben erhalten, können aber in kanonische NaN-Werte konvertiert werden.

Das Format mit niedrigerer Genauigkeit muss mindestens ein Exponentenbit haben (um einen Nullwert von einer Unendlichkeit zu unterscheiden, da beide eine Mantissen Null haben) und eine nicht negative Anzahl von Mantissen-Bits haben. Die Anzahl der Exponenten- oder Mantissenbits kann den entsprechenden Wert für den Typ T überschreiten. Der entsprechende Teil der Umwandlung ist dann einfach eine Nulloperation.

ReduceScatter

Weitere Informationen finden Sie unter XlaBuilder::ReduceScatter.

ReduceScatter ist ein kollektiver Vorgang, der einen AllReduce-Vorgang ausführt und das Ergebnis dann aufteilt, indem es in shard_count-Blöcke entlang der scatter_dimension aufgeteilt wird. Das Replikat i in der Replikatgruppe erhält den ith-Shard.

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

Argumente Typ Semantik
operand XlaOp Array oder ein nicht leeres Tupel von Arrays, um die Anzahl der Replikate zu reduzieren.
computation XlaComputation Reduktionsberechnung
scatter_dimension int64 Zu streuende Dimension.
shard_count int64 Anzahl der zu teilenden Blöcke in scatter_dimension
replica_groups Vektor der Vektoren von int64 Gruppen, zwischen denen die Reduzierungen durchgeführt werden
channel_id int64 (optional) Optionale Kanal-ID für modulübergreifende Kommunikation
  • Wenn operand ein Tupel von Arrays ist, wird die Reduce-Streuung für jedes Element des Tupels ausgeführt.
  • replica_groups ist eine Liste von Replikatgruppen, zwischen denen die Reduzierung durchgeführt wird. Die Replikat-ID für das aktuelle Replikat kann mit ReplicaId abgerufen werden. Die Reihenfolge der Replikate in jeder Gruppe bestimmt die Reihenfolge, in der das Ergebnis der Gesamtreduzierung verteilt wird. replica_groups muss entweder leer sein (in diesem Fall gehören alle Replikate zu einer einzigen Gruppe) oder muss dieselbe Anzahl von Elementen wie die Anzahl der Replikate enthalten. Wenn mehrere Replikatgruppen vorhanden sind, müssen alle dieselbe Größe haben. Beispielsweise führt replica_groups = {0, 2}, {1, 3} eine Reduktion zwischen den Replikaten 0 und 2 sowie 1 und 3 durch und teilt dann das Ergebnis.
  • shard_count ist die Größe der einzelnen Replikatgruppen. Wir benötigen dies, wenn replica_groups leer ist. Wenn replica_groups nicht leer ist, muss shard_count der Größe der einzelnen Replikatgruppen entsprechen.
  • channel_id wird für die modulübergreifende Kommunikation verwendet: Nur reduce-scatter-Vorgänge mit derselben channel_id können miteinander kommunizieren.

Die Ausgabeform ist die Eingabeform, wobei der scatter_dimension shard_count-mal kleiner ist. Wenn beispielsweise zwei Replikate vorhanden sind und der Operand den Wert [1.0, 2.25] bzw. [3.0, 5.25] auf den beiden Replikaten hat, ist der Ausgabewert dieses Vorgangs, bei dem scatter_dim 0 ist, [4.0] für das erste Replikat und [7.5] für das zweite Replikat.

ReduceWindow

Weitere Informationen finden Sie unter XlaBuilder::ReduceWindow.

Wendet eine Reduktionsfunktion auf alle Elemente in jedem Fenster einer Sequenz von n mehrdimensionalen Arrays an, wodurch ein einzelnes oder ein Tupel von n mehrdimensionalen Arrays als Ausgabe erzeugt wird. Jedes Ausgabearray hat die gleiche Anzahl von Elementen wie die Anzahl der gültigen Positionen des Fensters. Eine Pooling-Ebene kann als ReduceWindow ausgedrückt werden. Ähnlich wie bei Reduce wird der angewendete computation immer den init_values-Wert auf der linken Seite übergeben.

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

Argumente Typ Semantik
operands N XlaOps Eine Folge von n mehrdimensionalen Arrays der Typen T_0,..., T_{N-1}, die jeweils die Grundfläche darstellen, auf der das Fenster platziert wird.
init_values N XlaOps N Startwerte für die Reduktion, einer für jeden der N Operanden. Weitere Informationen finden Sie unter Reduzieren.
computation XlaComputation Reduktionsfunktion vom Typ T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}), die auf Elemente in jedem Fenster aller Eingabeoperanden angewendet wird.
window_dimensions ArraySlice<int64> Array von Ganzzahlen für Fensterdimensionswerte
window_strides ArraySlice<int64> Array von Ganzzahlen für Fensterschrittwerte
base_dilations ArraySlice<int64> Array von Ganzzahlen für Basisdilatationswerte
window_dilations ArraySlice<int64> Array von Ganzzahlen für Fensterdilatationswerte
padding Padding Padding-Typ für Fenster („Padding::kSame“, das so auffüllt, dass es dieselbe Ausgabeform wie die Eingabe hat, wenn der Schritt 1 ist, oder Padding::kValid, das keine Auffüllung verwendet und das Fenster „stoppt“, sobald es nicht mehr passt)

Wobei:

  • N muss größer oder gleich 1 sein.
  • Alle Eingabearrays müssen dieselben Abmessungen haben.
  • Wenn N = 1, ist Collate(T) T.
  • Bei N > 1 ist Collate(T_0, ..., T_{N-1}) ein Tupel von N-Elementen vom Typ (T0,...T{N-1}).

Der folgende Code und die Abbildung zeigen ein Beispiel für die Verwendung von ReduceWindow. Die Eingabe ist eine Matrix der Größe [4x6] und sowohl window_dimensions als auch window_stride_dimensions sind [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);

Der Schritt von 1 in einer Dimension gibt an, dass die Position eines Fensters in der Dimension ein Element vom benachbarten Fenster entfernt ist. Um anzugeben, dass sich keine Fenster überschneiden, sollte „window_stride_dimensions“ gleich „window_dimensions“ sein. Die folgende Abbildung zeigt die Verwendung von zwei verschiedenen Schrittwerten. Das Padding wird auf jede Dimension der Eingabe angewendet und die Berechnungen sind dieselben wie bei der Eingabe mit den Dimensionen, die sie nach dem Auffüllen hat.

Für ein nicht triviales Padding-Beispiel könnten Sie den Mindestwert für das Reduzieren des Fensters (Ausgangswert ist MAX_FLOAT) mit der Dimension 3 und dem Schritt 2 über dem Eingabearray [10000, 1000, 100, 10, 1] berechnen. Das Padding von kValid berechnet die Mindestwerte über zwei gültige Fenster: [10000, 1000, 100] und [100, 10, 1], was zur Ausgabe [100, 1] führt. Das Padding von kSame füllt zuerst das Array auf, sodass die Form nach dem Reduce-Fenster dieselbe wäre wie die Eingabe für Schritt 1. Dazu werden Anfangselemente auf beiden Seiten hinzugefügt, wobei [MAX_VALUE, 10000, 1000, 100, 10, 1, MAX_VALUE] abgerufen wird. Wenn Sie das Fenster „Reduzierung des Fensters“ über das aufgefüllte Array ausführen, werden die drei Fenster [MAX_VALUE, 10000, 1000], [1000, 100, 10], [10, 1, MAX_VALUE] und „yields“ [1000, 10, 1] ausgeführt.

Die Auswertungsreihenfolge der Reduktionsfunktion ist beliebig und kann nicht deterministisch sein. Daher sollte die Reduktionsfunktion nicht zu empfindlich auf Neuverknüpfungen reagieren. Weitere Informationen finden Sie in der Diskussion über Assoziativität im Kontext von Reduce.

ReplicaId

Weitere Informationen finden Sie unter XlaBuilder::ReplicaId.

Gibt die eindeutige ID (U32-Skalar) des Replikats zurück.

ReplicaId()

Die eindeutige ID jedes Replikats ist eine vorzeichenlose Ganzzahl im Intervall [0, N), wobei N die Anzahl der Replikate ist. Da alle Replikate dasselbe Programm ausführen, gibt ein ReplicaId()-Aufruf im Programm einen anderen Wert für jedes Replikat zurück.

Form ändern

Weitere Informationen finden Sie unter XlaBuilder::Reshape und Collapse-Vorgang.

Wandelt die Abmessungen eines Arrays in eine neue Konfiguration um.

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

Argumente Typ Semantik
operand XlaOp Array vom Typ T
dimensions Vektor: int64 Reihenfolge, in der Dimensionen minimiert werden
new_sizes Vektor: int64 Vektor von Größen mit neuen Abmessungen

Konzeptionell wird ein Array durch die Umform zuerst zu einem eindimensionalen Vektor von Datenwerten vereinfacht und dann in einer neuen Form verfeinert. Die Eingabeargumente sind ein beliebiges Array vom Typ T, ein Vektor der Compile-Zeitkonstante der Dimensionsindizes und ein Compile-Zeitkonstanter Vektor der Dimensionsgrößen für das Ergebnis. Die Werte im Vektor dimension, sofern angegeben, müssen eine Permutation aller Dimensionen von T sein. Wenn nicht angegeben, ist der Standardwert {0, ..., rank - 1}. Die Reihenfolge der Dimensionen in dimensions reicht von der langsamsten Dimension (größte) bis zur am schnellsten variierenden Dimension (die geringste) in der Schleifenverschachtelung, wodurch das Eingabearray in eine einzige Dimension minimiert wird. Der Vektor new_sizes bestimmt die Größe des Ausgabearrays. Der Wert bei Index 0 in new_sizes ist die Größe der Dimension 0, der Wert bei Index 1 ist die Größe von Dimension 1 usw. Das Produkt der new_size-Dimensionen muss dem Produkt der Dimensionsgrößen des Operanden entsprechen. Wenn Sie das minimierte Array auf das durch new_sizes definierte mehrdimensionale Array verfeinern, werden die Dimensionen in new_sizes von der langsamsten (größten) bis zur schnellsten variierenden (größten Abweichung) sortiert.

Ist v beispielsweise ein Array mit 24 Elementen:

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

In einem Sonderfall kann ein Array mit einem einzelnen Element durch die Umformung in einen Skalar umgewandelt werden und umgekehrt. Beispiel:

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

Umsatz (umgekehrt)

Weitere Informationen finden Sie unter XlaBuilder::Rev.

Rev(operand, dimensions)

Argumente Typ Semantik
operand XlaOp Array vom Typ T
dimensions ArraySlice<int64> Umzukehrende Dimensionen

Kehrt die Reihenfolge der Elemente im Array operand entlang des angegebenen dimensions um und generiert ein Ausgabearray derselben Form. Jedes Element des Operandenarrays bei einem mehrdimensionalen Index wird bei einem transformierten Index im Ausgabearray gespeichert. Der mehrdimensionale Index wird durch Umkehren des Index in jeder umgekehrten Dimension transformiert. Wenn also eine Dimension der Größe N eine der umgekehrten Dimensionen ist, wird ihr Index i in N – 1 – i umgewandelt.

Eine Verwendung der Rev-Operation besteht darin, das Faltungsgewichtsarray entlang der beiden Fensterdimensionen während der Gradientenberechnung in neuronalen Netzwerken umzukehren.

RngNormal

Weitere Informationen finden Sie unter XlaBuilder::RngNormal.

Konstruiert eine Ausgabe einer bestimmten Form mit Zufallszahlen, die nach der \(N(\mu, \sigma)\) -Normalverteilung generiert werden. Die Parameter \(\mu\) und \(\sigma\)sowie die Ausgabeform müssen einen Elementtyp „Gleitkommazahl“ haben. Außerdem müssen die Parameter skalarwert sein.

RngNormal(mu, sigma, shape)

Argumente Typ Semantik
mu XlaOp Skalar des Typs T, der den Mittelwert von generierten Zahlen angibt
sigma XlaOp Skalar vom Typ T, der die Standardabweichung der generierten
shape Shape Ausgabeform vom Typ T

RngUniform

Weitere Informationen finden Sie unter XlaBuilder::RngUniform.

Konstruiert eine Ausgabe einer bestimmten Form mit Zufallszahlen, die nach der gleichmäßigen Verteilung über das Intervall \([a,b)\)generiert werden. Die Parameter und der Ausgabeelementtyp müssen boolesche, ganzzahlige oder Gleitkommatypen sein. Außerdem müssen die Typen einheitlich sein. Die CPU- und GPU-Back-Ends unterstützen derzeit nur F64, F32, F16, BF16, S64, U64, S32 und U32. Außerdem müssen die Parameter einen skalaren Wert haben. Wenn \(b <= a\) , ist das Ergebnis implementiert.

RngUniform(a, b, shape)

Argumente Typ Semantik
a XlaOp Skalar des Typs T, der die Untergrenze des Intervalls angibt
b XlaOp Skalar des Typs T, der die Obergrenze des Intervalls angibt
shape Shape Ausgabeform vom Typ T

RngBitGenerator

Erzeugt mithilfe des angegebenen Algorithmus (oder des Back-End-Standardwerts) eine Ausgabe mit einer bestimmten Form, die mit einheitlichen zufälligen Bits gefüllt ist, und gibt einen aktualisierten Status (mit der gleichen Form wie der Anfangszustand) sowie die generierten Zufallsdaten zurück.

Der Anfangszustand ist der Anfangszustand der aktuellen Generierung von Zufallszahlen. Er, die erforderliche Form und gültige Werte hängen vom verwendeten Algorithmus ab.

Es wird garantiert, dass die Ausgabe eine deterministische Funktion des Anfangszustands ist. Es kann jedoch nicht garantiert werden, dass sie zwischen Back-Ends und verschiedenen Compiler-Versionen deterministisch ist.

RngBitGenerator(algorithm, key, shape)

Argumente Typ Semantik
algorithm RandomAlgorithm zu verwendender PRNG-Algorithmus.
initial_state XlaOp Ausgangsstatus für den PRNG-Algorithmus.
shape Shape Ausgabeform für generierte Daten.

Verfügbare Werte für algorithm:

Streudiagramm

Der XLA-Streuvorgang generiert eine Folge von Ergebnissen, die die Werte des Eingabearrays operands mit mehreren durch scatter_indices festgelegten Indexen sind, die mithilfe von update_computation mit der Reihenfolge der Werte in updates aktualisiert werden.

Weitere Informationen finden Sie unter XlaBuilder::Scatter.

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

Argumente Typ Semantik
operands Sequenz von N XlaOp N Arrays vom Typ T_0, ..., T_N, in die gestreut werden soll.
scatter_indices XlaOp Array mit den Startindexen der Segmente, auf die gestreut werden muss.
updates Sequenz von N XlaOp N Arrays vom Typ T_0, ..., T_N updates[i] enthält die Werte, die für die Streuung von operands[i] verwendet werden müssen.
update_computation XlaComputation Berechnung, die verwendet werden soll, um die vorhandenen Werte im Eingabearray und die Aktualisierungen während der Streuung zu kombinieren. Diese Berechnung muss vom Typ T_0, ..., T_N, T_0, ..., T_N -> Collate(T_0, ..., T_N) sein.
index_vector_dim int64 Die Dimension in scatter_indices, die die Startindexe enthält.
update_window_dims ArraySlice<int64> Die Menge der Abmessungen in der updates-Form, bei denen es sich um die Fensterabmessungen handelt.
inserted_window_dims ArraySlice<int64> Die Fensterabmessungen, die in die Form updates eingefügt werden müssen.
scatter_dims_to_operand_dims ArraySlice<int64> Eine Dimensionszuordnung von den Streuindexen bis zum Operandenindexbereich Dieses Array wird als Zuordnung von i zu scatter_dims_to_operand_dims[i] interpretiert . Sie muss eine Person und ein Gesamtbild sein.
indices_are_sorted bool Gibt an, ob die Indizes garantiert vom Aufrufer sortiert werden.
unique_indices bool Gibt an, ob die Indizes garantiert eindeutig sind für den Aufrufer.

Wobei:

  • N muss größer oder gleich 1 sein.
  • operands[0], ..., operands[N-1] müssen alle dieselben Abmessungen haben.
  • updates[0], ..., updates[N-1] müssen alle dieselben Abmessungen haben.
  • Wenn N = 1, ist Collate(T) T.
  • Bei N > 1 ist Collate(T_0, ..., T_N) ein Tupel von N-Elementen vom Typ T.

Wenn index_vector_dim gleich scatter_indices.rank ist, wird scatter_indices implizit als nachgestellte 1-Dimension betrachtet.

Wir definieren update_scatter_dims vom Typ ArraySlice<int64> als Gruppe von Dimensionen in aufsteigender Reihenfolge in der Form updates, die nicht update_window_dims haben.

Die Argumente von Streudiagrammen sollten folgenden Einschränkungen folgen:

  • Jedes updates-Array muss den Rang update_window_dims.size + scatter_indices.rank - 1 haben.

  • Die Grenzen der Dimension i in jedem updates-Array müssen Folgendem entsprechen:

    • Wenn i in update_window_dims vorhanden ist (d. h. update_window_dims[k] für einen Teil k), darf die Grenze der Dimension i in updates die entsprechende Grenze von operand nach Berücksichtigung von inserted_window_dims nicht überschreiten. Beispiel: adjusted_window_bounds[k], wobei adjusted_window_bounds die Grenzen von operand enthält, wobei die Grenzen bei den Indexen inserted_window_dims entfernt wurden.
    • Wenn i in update_scatter_dims vorhanden ist (d.h. update_scatter_dims[k] für einen Teil k), muss die Grenze der Dimension i in updates der entsprechenden Grenze von scatter_indices entsprechen und index_vector_dim wird übersprungen (d.h. scatter_indices.shape.dims[k], wenn k < index_vector_dim und ansonsten scatter_indices.shape.dims[k+1]).
  • update_window_dims muss in aufsteigender Reihenfolge angegeben werden, darf keine sich wiederholenden Dimensionsnummern haben und im Bereich [0, updates.rank) liegen.

  • inserted_window_dims muss in aufsteigender Reihenfolge angegeben werden, darf keine sich wiederholenden Dimensionsnummern haben und im Bereich [0, operand.rank) liegen.

  • operand.rank muss der Summe von update_window_dims.size und inserted_window_dims.size entsprechen.

  • scatter_dims_to_operand_dims.size muss gleich scatter_indices.shape.dims[index_vector_dim] sein und die Werte müssen im Bereich [0, operand.rank) liegen.

Für einen bestimmten Index U in jedem updates-Array wird der entsprechende Index I im entsprechenden operands-Array, auf das diese Aktualisierung angewendet werden muss, so berechnet:

  1. Lassen Sie G = { U[k] für k in update_scatter_dims }. Verwenden Sie G, um einen Indexvektor S im Array scatter_indices zu suchen, sodass S[i] = scatter_indices[Kombination(G, i)], wobei Kombinieren(A, b) b an den Positionen index_vector_dim in A einfügt.
  2. Erstellen Sie mit S einen Index Sin in operand, indem Sie S mithilfe der scatter_dims_to_operand_dims-Karte verteilen. Formeller:
    1. Sin[scatter_dims_to_operand_dims[k]] = S[k], wenn k < scatter_dims_to_operand_dims.size.
    2. SAndernfalls in[_] = 0.
  3. Erstellen Sie einen Index Win in jedem operands-Array, indem Sie die Indexe bei update_window_dims in U gemäß inserted_window_dims verteilen. Formeller:
    1. Win[window_dims_to_operand_dims(k)] = U[k], wenn k in update_window_dims ist, wobei window_dims_to_operand_dims die monotone Funktion mit der Domain [0, update_window_dims.size) und dem Bereich [0, operand.rank) \ inserted_window_dims ist. Wenn update_window_dims.size beispielsweise 4 ist, operand.rank 6 und inserted_window_dims {0, 2} ist, dann ist window_dims_to_operand_dims {01, 13, 24, 35}.
    2. WAndernfalls in[_] = 0.
  4. I ist Win + Sin, wobei + für das elementweise Addieren steht.

Zusammenfassend lässt sich die Streuoperation wie folgt definieren.

  • Initialisieren Sie output mit operands, d.h. für alle Indexe J und für alle Indexe O im Array operands[J]:
    output[J][O] = operands[J][O]
  • Für jeden Index-U im Array updates[J] und den entsprechenden Index O im Array operand[J], wenn O ein gültiger Index für output ist:
    (output[0][O], ..., output[N-1][O]) =update_computation(output[0][O], ..., ,output[N-1][O],updates[0][U], ...,updates[N-1][U])

Die Reihenfolge, in der Aktualisierungen angewendet werden, ist nicht deterministisch. Wenn also mehrere Indizes in updates auf denselben Index in operands verweisen, ist der entsprechende Wert in output nicht deterministisch.

Der erste Parameter, der an das update_computation-Array übergeben wird, ist immer der aktuelle Wert aus dem output-Array. Der zweite Parameter ist immer der Wert aus dem updates-Array. Das ist insbesondere dann wichtig, wenn update_computation nicht kommutativ ist.

Wenn indices_are_sorted auf „true“ gesetzt ist, kann XLA davon ausgehen, dass start_indices (in aufsteigender start_index_map-Reihenfolge) nach dem Nutzer sortiert werden. Andernfalls ist die Semantik Implementierung definiert.

Wenn unique_indices auf „true“ gesetzt ist, kann XLA davon ausgehen, dass alle verstreuten Elemente eindeutig sind. XLA könnte also nicht-atomare Vorgänge verwenden. Wenn unique_indices auf „true“ gesetzt ist und die Indizes, auf die verteilt werden, nicht eindeutig sind, ist die Semantik Implementierung definiert.

Informell kann der Streuvorgang als Inverse des Erfassungsvorgangs betrachtet werden, d.h., der Streuvorgang aktualisiert die Elemente in der Eingabe, die von der entsprechenden Erfassungsoperation extrahiert werden.

Eine ausführliche Beschreibung und Beispiele finden Sie im Abschnitt „Informelle Beschreibung“ unter Gather.

Auswählen

Weitere Informationen finden Sie unter XlaBuilder::Select.

Konstruiert ein Ausgabearray aus Elementen von zwei Eingabearrays, basierend auf den Werten eines Prädikatarrays.

Select(pred, on_true, on_false)

Argumente Typ Semantik
pred XlaOp Array vom Typ PRED
on_true XlaOp Array vom Typ T
on_false XlaOp Array vom Typ T

Die Arrays on_true und on_false müssen dieselbe Form haben. Dies ist auch die Form des Ausgabearrays. Das Array pred muss dieselbe Dimensionalität wie on_true und on_false mit dem Elementtyp PRED haben.

Für jedes Element P von pred wird das entsprechende Element des Ausgabearrays aus on_true übernommen, wenn der Wert von P true ist, und aus on_false, wenn der Wert von P false ist. Als eingeschränkte Form der Übertragung kann pred ein Skalar vom Typ PRED sein. In diesem Fall wird das Ausgabearray vollständig aus on_true genommen, wenn pred gleich true ist, und aus on_false, wenn pred gleich false ist.

Beispiel mit nicht skalarem pred:

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

Beispiel mit skalarem 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};

Die Auswahl zwischen Tupeln wird unterstützt. Tupel werden zu diesem Zweck als skalare Typen betrachtet. Wenn on_true und on_false Tupel sind (die dieselbe Form haben müssen), muss pred ein Skalar vom Typ PRED sein.

SelectAndScatter

Weitere Informationen finden Sie unter XlaBuilder::SelectAndScatter.

Diese Operation kann als zusammengesetzte Operation betrachtet werden, bei der zuerst ReduceWindow für das Array operand berechnet wird, um ein Element aus jedem Fenster auszuwählen. Anschließend wird das Array source an die Indizes der ausgewählten Elemente verteilt, um ein Ausgabearray mit derselben Form wie das Operandenarray zu erstellen. Die binäre select-Funktion wird verwendet, um ein Element aus jedem Fenster auszuwählen, indem es auf jedes Fenster angewendet wird. Sie wird mit der Eigenschaft aufgerufen, dass der Indexvektor des ersten Parameters lexikografisch kleiner als der Indexvektor des zweiten Parameters ist. Die select-Funktion gibt true zurück, wenn der erste Parameter ausgewählt ist, und false, wenn der zweite Parameter ausgewählt ist. Die Funktion muss die Transitivität enthalten, d. h., wenn select(a, b) und select(b, c) true sind, dann ist select(a, c) ebenfalls true. So ist das ausgewählte Element nicht von der Reihenfolge der Elemente abhängig, die in einem bestimmten Fenster durchlaufen werden.

Die Funktion scatter wird auf jeden ausgewählten Index im Ausgabearray angewendet. Sie verwendet zwei Skalarparameter:

  1. Aktueller Wert am ausgewählten Index im Ausgabearray
  2. Der Streuwert aus source, der für den ausgewählten Index gilt

Die Funktion kombiniert die beiden Parameter und gibt einen Skalarwert zurück, mit dem der Wert am ausgewählten Index im Ausgabearray aktualisiert wird. Zu Beginn sind alle Indizes des Ausgabearrays auf init_value gesetzt.

Das Ausgabearray hat dieselbe Form wie das operand-Array und das source-Array muss dieselbe Form wie das Ergebnis der Anwendung eines ReduceWindow-Vorgangs auf das operand-Array haben. Mit SelectAndScatter können die Gradientenwerte für eine Pooling-Ebene in einem neuronalen Netzwerk rückpropagiert werden.

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

Argumente Typ Semantik
operand XlaOp Array vom Typ T, über das die Fenster
select XlaComputation Binäre Berechnung des Typs T, T -> PRED, die auf alle Elemente in jedem Fenster angewendet wird. Gibt true zurück, wenn der erste Parameter ausgewählt ist, und false, wenn der zweite Parameter ausgewählt ist.
window_dimensions ArraySlice<int64> Array von Ganzzahlen für Fensterdimensionswerte
window_strides ArraySlice<int64> Array von Ganzzahlen für Fensterschrittwerte
padding Padding Abstandstyp für Fenster (Padding::kSame oder Padding::kValid)
source XlaOp Array vom Typ T mit den zu streuenden Werten
init_value XlaOp Skalarwert vom Typ T für den Anfangswert des Ausgabearrays
scatter XlaComputation Binäre Berechnung des Typs T, T -> T, um jedes Streuquellenelement auf sein Zielelement anzuwenden

Die folgende Abbildung zeigt Beispiele für die Verwendung von SelectAndScatter, wobei die Funktion select den maximalen Wert aus ihren Parametern berechnet. Wenn sich die Fenster wie in der nachfolgenden Abbildung (2) überschneiden, kann ein Index des Arrays operand von verschiedenen Fenstern mehrmals ausgewählt werden. In der Abbildung wird das Element mit dem Wert 9 durch die beiden oberen Fenster (blau und rot) ausgewählt und die Binäradditionsfunktion scatter erzeugt das Ausgabeelement des Werts 8 (2 + 6).

Die Auswertungsreihenfolge der scatter-Funktion ist beliebig und kann nicht deterministisch sein. Daher sollte die scatter-Funktion nicht zu empfindlich auf Neuverknüpfungen reagieren. Weitere Informationen finden Sie in der Diskussion über Assoziativität im Kontext von Reduce.

Senden

Weitere Informationen finden Sie unter XlaBuilder::Send.

Send(operand, channel_handle)

Argumente Typ Semantik
operand XlaOp Zu sendende Daten (Array vom Typ T)
channel_handle ChannelHandle eindeutige Kennung für jedes Sende-/Empfangspaar

Sendet die angegebenen Operandendaten in einer anderen Berechnung mit demselben Kanal-Handle an einen Recv-Befehl. Gibt keine Daten zurück.

Ähnlich wie beim Recv-Vorgang stellt die Client API des Vorgangs Send die synchrone Kommunikation dar und wird intern in zwei HLO-Anweisungen (Send und SendDone) aufgeteilt, um asynchrone Datenübertragungen zu ermöglichen. Weitere Informationen finden Sie unter HloInstruction::CreateSend und HloInstruction::CreateSendDone.

Send(HloInstruction operand, int64 channel_id)

Initiiert eine asynchrone Übertragung des Operanden an die durch die Anweisung Recv zugewiesenen Ressourcen mit derselben Kanal-ID. Gibt einen Kontext zurück, der von der folgenden SendDone-Anweisung verwendet wird, um auf den Abschluss der Datenübertragung zu warten. Der Kontext ist ein Tupel von {operand (shape), request Identifier (U32)} und kann nur von einer SendDone-Anweisung verwendet werden.

SendDone(HloInstruction context)

Wenn ein Kontext vorhanden ist, der durch eine Send-Anweisung erstellt wurde, wird auf den Abschluss der Datenübertragung gewartet. Die Anweisung gibt keine Daten zurück.

Planung der Kanalanweisungen

Die Ausführungsreihenfolge der vier Anleitungen für jeden Kanal (Recv, RecvDone, Send, SendDone) sieht so aus:

  • Recv findet vor dem Send statt
  • Send findet vor dem RecvDone statt
  • Recv findet vor dem RecvDone statt
  • Send findet vor dem SendDone statt

Wenn die Back-End-Compiler einen linearen Zeitplan für jede Berechnung generieren, die über Kanalanweisungen kommuniziert, dürfen keine Zyklen zwischen den Berechnungen sein. Die folgenden Zeitpläne führen beispielsweise zu Deadlocks.

Slice

Weitere Informationen finden Sie unter XlaBuilder::Slice.

Durch die Segmentierung wird ein Teilarray aus dem Eingabearray extrahiert. Das Subarray hat den gleichen Rang wie die Eingabe und enthält die Werte innerhalb eines Begrenzungsrahmens innerhalb des Eingabearrays, wobei die Abmessungen und Indizes des Begrenzungsrahmens als Argumente für den Segmentierungsvorgang angegeben werden.

Slice(operand, start_indices, limit_indices, strides)

Argumente Typ Semantik
operand XlaOp N-dimensionales Array vom Typ T
start_indices ArraySlice<int64> Liste von N Ganzzahlen, die die Startindizes des Segments für jede Dimension enthalten. Werte müssen größer oder gleich null sein.
limit_indices ArraySlice<int64> Liste von N Ganzzahlen, die die Endindizes (ausschließlich) für das Segment für jede Dimension enthalten. Jeder Wert muss größer oder gleich dem jeweiligen start_indices-Wert für die Dimension und kleiner oder gleich der Größe der Dimension sein.
strides ArraySlice<int64> Liste mit N Ganzzahlen, die den Eingabeschritt des Slice festlegt. Mit dem Segment wird jedes strides[d]-Element der Dimension d ausgewählt.

Beispiel für eindimensionales Element:

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

Beispiel für zweidimensionale Darstellung:

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

Sortieren

Weitere Informationen finden Sie unter XlaBuilder::Sort.

Sort(operands, comparator, dimension, is_stable)

Argumente Typ Semantik
operands ArraySlice<XlaOp> Die Operanden, die sortiert werden sollen.
comparator XlaComputation Die zu verwendende Vergleichsberechnung.
dimension int64 Die Dimension, nach der sortiert werden soll.
is_stable bool Gibt an, ob eine stabile Sortierung verwendet werden soll.

Wenn nur ein Operand angegeben ist:

  • Wenn der Operand ein Rang-1-Tensor (ein Array) ist, ist das Ergebnis ein sortiertes Array. Wenn Sie das Array in aufsteigender Reihenfolge sortieren möchten, sollte der Vergleichsoperator einen Kleiner-als-Vergleich durchführen. Formal gilt das Array nach dem Sortieren für alle Indexpositionen i, j mit i < j, die entweder comparator(value[i], value[j]) = comparator(value[j], value[i]) = false oder comparator(value[i], value[j]) = true sind.

  • Wenn der Operand einen höheren Rang hat, wird der Operand entlang der angegebenen Dimension sortiert. Bei einem Rang-2-Tensor (Matrix) wird beispielsweise mit dem Dimensionswert 0 jede Spalte unabhängig sortiert und beim Dimensionswert 1 jede Zeile unabhängig voneinander. Ist keine Dimensionsnummer angegeben, wird standardmäßig die letzte Dimension ausgewählt. Für die sortierte Dimension gilt dieselbe Sortierreihenfolge wie für den Rang 1.

Wenn n > 1-Operanden angegeben sind:

  • Alle n-Operanden müssen Tensoren mit denselben Dimensionen sein. Die Elementtypen der Tensoren können unterschiedlich sein.

  • Alle Operanden werden zusammen, nicht einzeln sortiert. Konzeptionell werden die Operanden als Tupel behandelt. Bei der Überprüfung, ob die Elemente jedes Operanden an den Indexpositionen i und j vertauscht werden müssen, wird der Vergleichsoperator mit skalaren 2 * n-Parametern aufgerufen, wobei der Parameter 2 * k dem Wert an Position i aus dem Operanden k-th und der Parameter 2 * k + 1 dem Wert an Position j vom Operanden k-th entspricht. Normalerweise würde der Vergleichsoperator somit die Parameter 2 * k und 2 * k + 1 miteinander vergleichen und möglicherweise andere Parameterpaare als Stichpunkte verwenden.

  • Das Ergebnis ist ein Tupel, das aus den Operanden in sortierter Reihenfolge (entlang der angegebenen Dimension, wie oben) besteht. Der Operand i-th des Tupels entspricht dem Operanden i-th von „Sort“.

Wenn es beispielsweise die drei Operanden operand0 = [3, 1], operand1 = [42, 50] und operand2 = [-3.0, 1.1] gibt und der Vergleichsoperator nur die Werte von operand0 mit „kleiner als“ vergleicht, ist die Ausgabe der Sortierung das Tupel ([1, 3], [50, 42], [1.1, -3.0]).

Wenn is_stable auf „true“ gesetzt ist, ist die Sortierung garantiert stabil. Wenn es also Elemente gibt, die vom Vergleichsoperator als gleich angesehen werden, wird die relative Reihenfolge der Gleichheitswerte beibehalten. Die beiden Elemente e1 und e2 sind nur dann gleich, wenn comparator(e1, e2) = comparator(e2, e1) = false vorhanden ist. Die Standardeinstellung für is_stable ist „false“.

Transponieren

Weitere Informationen finden Sie im tf.reshape-Vorgang.

Transpose(operand)

Argumente Typ Semantik
operand XlaOp Der Operand, der transponiert werden soll.
permutation ArraySlice<int64> So ändern Sie die Dimensionen:

Permutiert die Operandendimensionen mit der angegebenen Permutation, sodass ∀ i . 0 ≤ i < rank ⇒ input_dimensions[permutation[i]] = output_dimensions[i].

Dies entspricht „Reshape(operand, permutation, Permute(permutation, operand.shape.dimensions))“.

TriangularSolve

Weitere Informationen finden Sie unter XlaBuilder::TriangularSolve.

Löst lineare Gleichungssysteme mit unteren oder oberen Dreieckskoeffizienten durch Vorwärts- oder RückwärtsSubstitution. Bei dieser Routine, die entlang führender Dimensionen gesendet wird, wird eines der Matrixsysteme op(a) * x = b oder x * op(a) = b für die Variable x gelöst, sofern a und b, wobei op(a) entweder op(a) = a, op(a) = Transpose(a) oder op(a) = Conj(Transpose(a)) ist.

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

Argumente Typ Semantik
a XlaOp Ein Array vom Typ „Rang > 2“ eines komplexen oder Gleitkommatyps mit der Form [..., M, M].
b XlaOp ein Array mit Rang > 2 desselben Typs mit der Form [..., M, K], wenn left_side wahr ist, andernfalls [..., K, M].
left_side bool gibt an, ob ein System der Form op(a) * x = b (true) oder x * op(a) = b (false) gelöst werden soll.
lower bool ob das obere oder untere Dreieck von a verwendet werden soll.
unit_diagonal bool Wenn true festgelegt ist, wird angenommen, dass die diagonalen Elemente von a 1 sind. Auf die diagonalen Elemente wird nicht zugegriffen.
transpose_a Transpose a kann unverändert verwendet oder transponiert oder konjugiert werden.

Eingabedaten werden je nach Wert von lower nur aus dem unteren/oberen Dreieck von a gelesen. Werte aus dem anderen Dreieck werden ignoriert. Die Ausgabedaten werden im gleichen Dreieck zurückgegeben. Die Werte im anderen Dreieck sind implementierungsdefiniert und können beliebig sein.

Wenn der Rang von a und b größer als 2 ist, werden sie als Batches von Matrizen behandelt, wobei alle außer den zwei untergeordneten Dimensionen Batchdimensionen sind. a und b müssen dieselben Batchdimensionen haben.

Tupel

Weitere Informationen finden Sie unter XlaBuilder::Tuple.

Ein Tupel mit einer variablen Anzahl von Daten-Handles, von denen jedes eine eigene Form hat.

Dies entspricht std::tuple in C++. Konzeptionell:

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

Tupel können über den Vorgang GetTupleElement dekonstruiert (aufgerufen) werden.

Während

Weitere Informationen finden Sie unter XlaBuilder::While.

While(condition, body, init)

Argumente Typ Semantik
condition XlaComputation XlaComputation vom Typ T -> PRED, die die Beendigungsbedingung der Schleife definiert.
body XlaComputation XlaComputation vom Typ T -> T, der den Text der Schleife definiert.
init T Anfangswert für den Parameter von condition und body.

Führt body sequenziell aus, bis condition fehlschlägt. Dies ähnelt einer typischen „Loop“-Methode in vielen anderen Sprachen, mit Ausnahme der unten aufgeführten Unterschiede und Einschränkungen.

  • Ein While-Knoten gibt einen Wert vom Typ T zurück, der das Ergebnis der letzten Ausführung von body ist.
  • Die Form des Typs T wird statisch bestimmt und muss in allen Iterationen gleich sein.

Die T-Parameter der Berechnungen werden im ersten Durchlauf mit dem Wert init initialisiert und bei jedem weiteren Durchlauf automatisch auf das neue Ergebnis von body aktualisiert.

Ein wesentlicher Anwendungsfall des Knotens While ist die Implementierung der wiederholten Ausführung des Trainings in neuronalen Netzwerken. Unten wird ein vereinfachter Pseudocode mit einem Diagramm dargestellt, das die Berechnung darstellt. Den Code finden Sie in while_test.cc. Der Typ T in diesem Beispiel ist eine Tuple, die aus einem int32 für die Iterationsanzahl und einem vector[10] für den Akkumulator besteht. Für 1.000 Iterationen fügt die Schleife dem Akkumulator weiter einen konstanten Vektor hinzu.

// 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};
}