Vorgangssemantik

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

Ein Hinweis zur Nomenklatur: Der generalisierte Datentyp XLA ist ein n-dimensionales Array, das Elemente eines bestimmten Typs (z. B. 32-Bit-Gleitkommazahl) enthält. 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“ nimmt eine variable Anzahl von Tokens an und erzeugt ein einzelnes Token. Tokens sind primitive Typen, die zwischen Nebenwirkungsvorgängen angeordnet werden können, um eine Sortierung zu erzwingen. AfterAll kann als Join von Tokens verwendet werden, um einen Vorgang nach festgelegten Vorgängen zu sortieren.

AfterAll(operands)

Argumente Typ Semantik
operands XlaOp unterschiedliche Anzahl von Tokens

AllGather

Weitere Informationen finden Sie unter XlaBuilder::AllGather.

Führt eine Verkettung über alle Replikate hinweg durch.

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

Argumente Typ Semantik
operand XlaOp Array zur Verkettung über Replikate hinweg
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 erfolgt (die Replikat-ID für das aktuelle Replikat kann mit ReplicaId abgerufen werden). Die Reihenfolge der Replikate in den einzelnen Gruppen bestimmt die Reihenfolge, in der sich ihre Eingaben im Ergebnis befinden. replica_groups muss entweder leer sein (in diesem Fall gehören alle Replikate zu einer einzigen Gruppe, sortiert von 0 bis N - 1) oder 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-Wert, der shard_count-mal größer ist. Wenn beispielsweise zwei Replikate vorhanden sind und der Operand auf den beiden Replikaten den Wert [1.0, 2.5] bzw. [3.0, 5.25] hat, ist der Ausgabewert aus diesem Vorgang, bei dem all_gather_dim 0 ist, 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 durchgeführt werden,
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 (Replikat-ID für das aktuelle Replikat kann mit ReplicaId abgerufen werden). replica_groups muss entweder leer sein (in diesem Fall alle Replikate zu einer einzigen Gruppe gehören) oder dieselbe Anzahl von Elementen enthalten wie die Anzahl der Replikate. Beispiel: replica_groups = {0, 2}, {1, 3} führt eine Reduzierung zwischen den Replikaten 0 und 2 sowie 1 und 3 durch.
  • 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 dieser 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 öfter als ein anderes ausführt, wartet das vorherige Replikat unbegrenzt. Da die Replikate alle dasselbe Programm ausführen, gibt es nicht viele Möglichkeiten, dies zu tun. Es ist jedoch möglich, dass die Bedingung einer Während-Schleife von Daten aus der Einspeisung abhängt und die eingespeisten Daten dazu führen, dass die Während-Schleife mehr Mal ein Replikat wiederholt als ein anderes.

AllToAll

Weitere Informationen finden Sie unter XlaBuilder::AllToAll.

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

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

Die teilnehmenden Kerne können folgendermaßen 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 mithilfe von 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 innerhalb der Replikate {1, 2, 3} und in der Erfassungsphase angewendet wird und die empfangenen Blöcke in der gleichen Reihenfolge von 1, 2, 3 verkettet werden. Dann wird ein weiteres "AllToAll" in den Replikaten 4, 5, 0 angewendet und die Verkettungsreihenfolge ist ebenfalls 4, 5, 0. Wenn replica_groups leer ist, gehören alle Replikate in der Verkettungsreihenfolge 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 angibt, nach 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 für ReplicaGroup Jede Gruppe enthält eine Liste mit 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 am Alltoall 4 Kerne teil. Auf jedem Kern wird der Operand entlang der Dimension 0 in vier Teile geteilt, sodass jedes Teil die Form f32[4,4] hat. Die vier Teile sind auf alle Kerne verteilt. Dann verkettet jeder Kern die empfangenen Teile entlang der Dimension 1 in der Reihenfolge von Kern 0 bis 4. Die Ausgabe auf jedem Kern hat also die Form f32[16,4].

BatchNormGrad

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

Berechnet Gradienten der Batchnorm.

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

Argumente Typ Semantik
operand XlaOp Zu normalisierendes n-dimensionales Array (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 Feature in der Feature-Dimension (feature_index ist der Index für die Feature-Dimension in operand) berechnet der Vorgang die Gradienten in Bezug auf operand, offset und scale über alle anderen Dimensionen hinweg. feature_index muss ein gültiger Index für die Feature-Dimension in operand sein.

Die drei Gradienten werden durch die folgenden Formeln definiert, wobei von einem vierdimensionalen Array als operand und einem Featuredimensionsindex l, der Batchgröße m und den räumlichen Größen w und h ausgegangen wird:

\[ \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 in 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 Gradient in Bezug auf die Eingabe offset($\nabla \beta$)

BatchNormInference

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

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 Merkmalsdimension 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 Aufruf von BatchNormTraining, ohne mean und variance für jeden Batch zu berechnen. Stattdessen werden die eingegebenen mean und variance als geschätzte Werte verwendet. Der Zweck dieses Vorgangs ist es, 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 Zu normalisierendes n-dimensionales Array (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 Merkmalsdimension 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 gilt 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 jedes 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 Batchvarianz \(\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 verschiebt: \(y_{ijkl}=\frac{\gamma_l(x_{ijkl}-\mu_l)}{\sqrt[2]{\sigma^2_l+\epsilon} }+\beta_l\)

Der Epsilon-Wert, normalerweise eine kleine Zahl, wird addiert, um Fehler durch 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 mit den obigen Formeln berechnet werden.

BitcastConvertType

Weitere Informationen finden Sie unter XlaBuilder::BitcastConvertType.

Führt ähnlich wie bei einem tf.bitcast in TensorFlow eine elementweise Bitcast-Operation 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 Dimmung D
new_element_type PrimitiveType Typ U

Die Abmessungen des Operanden und der Zielform müssen übereinstimmen, abgesehen von 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, in dem 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.

Erstens: Wenn B > B', erhält die Ausgabeform eine neue kleinere Dimension 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 gleich:

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

Alternativ erfordert die Anweisung für B' > B, dass die letzte logische Dimension der Eingabeform gleich B'/B ist. Diese Dimension wird während 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 erfolgen.

Nachricht an alle

Weitere Informationen finden Sie unter XlaBuilder::Broadcast.

Fügt einem Array Dimensionen durch Duplizieren der Daten im Array hinzu.

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, d.h., wenn broadcast_sizes die Werte {a0, ..., aN} und die Operandenform die Dimensionen {b0, ..., bM} hat, hat die Form der Ausgabe die Dimensionen {a0, ..., aN, b0, ..., bM}.

Der neue Dimensionsindex wird in Kopien des Operanden, d.h.

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

Beispiel: Wenn operand ein skalarer f32 mit dem Wert 2.0f ist und broadcast_sizes gleich {2, 3} ist, 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 durch Duplizieren der Daten im Array.

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 jede Dimension der Operandenform entspricht

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

Das operand wird an die durch 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 die Größe 1 oder dieselbe Größe wie die Dimension in der Ausgabeform haben, der sie zugeordnet sind. Die übrigen Dimensionen werden mit Abmessungen der Größe 1 gefüllt. Beim Broadcasting mit degenerierenden Dimensionen wird es dann entlang dieser degenerierten Dimensionen übertragen, 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 eines beliebigen Typs
args Sequenz von N XlaOp s N Argumente eines beliebigen Typs

Die ARity und die Typen von args müssen mit den Parametern von computation übereinstimmen. Es darf keine args enthalten.

Cholesky

Weitere Informationen finden Sie unter XlaBuilder::Cholesky.

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

Cholesky(a, lower)

Argumente Typ Semantik
a XlaOp Array > 2 (Rang > 2) eines komplexen oder Gleitkommatyps
lower bool Gibt an, ob das obere oder untere Dreieck von a verwendet werden soll.

Wenn lower true ist, werden Matrizen für kleine Dreiecke l so berechnet, dass $a = l ist . l^T$. Wenn lower false ist, werden die Matrizen des oberen Dreiecks u so berechnet:\(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. Ausgabedaten werden im selben 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 ein Batch von Matrizen behandelt, wobei alle außer den zwei weiteren Dimensionen Batchdimensionen sind.

Wenn a nicht symmetrisch (Hermitian) positiv eindeutig ist, ist das Ergebnis von der Implementierung definiert.

Klemme

Weitere Informationen finden Sie unter XlaBuilder::Clamp.

Bindet einen Operanden innerhalb des Bereichs 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 Mindest- und Höchstwert wird der Operand zurückgegeben, wenn er im Bereich zwischen Mindest- und Höchstwert liegt. Andernfalls wird der Mindestwert zurückgegeben, wenn der Operand unter diesem Bereich liegt, oder den Höchstwert, 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 Übertragung können min und/oder max ein Skalar vom Typ T sein.

Beispiel mit dem Skalar 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.

Minimiert die Dimensionen eines Arrays in eine Dimension.

Collapse(operand, dimensions)

Argumente Typ Semantik
operand XlaOp Array vom Typ T
dimensions Vektor für int64 Aufeinanderfolgende Teilmenge der T-Dimensionen.

Bei der Minimierung wird die angegebene Teilmenge der Operandendimensionen durch eine einzelne Dimension ersetzt. Die Eingabeargumente sind ein beliebiges Array vom Typ T und ein Vektor der kompilierten Zeitkonstante von Dimensionsindexen. Die Dimensionsindexe müssen eine fortlaufende Teilmenge der T-Dimensionen (niedrige bis hohe Dimensionszahlen) sein. Somit 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. Die neue Dimensionsgröße entspricht dem Produkt der ursprünglichen Abmessungen. Die niedrigste Dimensionszahl in dimensions ist die langsamste variierende Dimension (die größte) in der Schleifenverschachtelung, bei der diese Dimensionen minimiert werden. Die höchste Dimensionsnummer variiert am schnellsten (die geringste). Wenn eine allgemeinere Minimierungsreihenfolge erforderlich ist, verwenden Sie den Operator tf.reshape.

Nehmen wir z. B. „v“ als 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 über mehrere Replikate hinweg sendet und empfängt.

CollectivePermute(operand, source_target_pairs)

Argumente Typ Semantik
operand XlaOp n-dimensionales Eingabearray
source_target_pairs Vektor für <int64, int64> Eine Liste mit Paaren aus Quelle_Replikat_ID und Ziel-Replikat_ID. Für jedes Paar wird der Operand vom Quellreplikat zum Zielreplikat gesendet.

Für source_target_pair gelten die folgenden Einschränkungen:

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

Verketten

Weitere Informationen finden Sie unter XlaBuilder::ConcatInDim.

Verketten besteht aus mehreren Array-Operanden. Das Array hat denselben Rang wie die einzelnen Operanden des Eingabearrays (die denselben Rang haben müssen wie die anderen) 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. Werte mit dem Rang 0 können nicht verkettet werden, da es nicht möglich ist, die Dimension zu benennen, in der die Verkettung erfolgt.

Beispiel für eindimensionales Format:

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 den Wert true hat, und false_computation, wenn pred 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, die denselben Typ haben müssen. false_computation muss ein einzelnes Argument vom Typ \(T_1\) annehmen und wird mit false_operand aufgerufen, die denselben Typ haben müssen. 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 XlaComputations 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.

Jede 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 der zurückgegebenen Werte 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.

Als 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, dass es keine Auffüllung gibt.

ConvWithGeneralPadding (Faltung)

Weitere Informationen finden Sie unter XlaBuilder::ConvWithGeneralPadding.

Berechnet eine Faltung von der Art, die in neuronalen Netzwerken verwendet wird. Hier kann man sich eine Faltung als ein n-dimensionales Fenster vorstellen, das sich über eine n-dimensionale Grundfläche bewegt und für jede mögliche Position des Fensters eine Berechnung durchgeführt wird.

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 Kernel-Schritten
padding ArraySlice< pair<int64,int64>> n-d-Array des Abstands (niedrig, hoch)
lhs_dilation ArraySlice<int64> n-d-lhs-Erweiterungsfaktor-Array
rhs_dilation ArraySlice<int64> n-d rhs-Dilationsfaktor-Array
feature_group_count int64 Anzahl der Funktionsgruppen
batch_group_count int64 Anzahl der Batchgruppen

Es soll n die Anzahl der räumlichen Dimensionen sein. Das Argument lhs ist ein Rang-n+2-Array, das die Grundfläche beschreibt. Diese wird als Eingabe bezeichnet, auch wenn „rhs“ ebenfalls 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, durch die sich das Fenster bewegt.

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

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

Das Argument window_strides gibt den Schritt des Faltungsfensters 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 auf die Basisfläche angewendeten Auffüllwert null an. 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 müssen. padding[0] gibt den Abstand für die Dimension y an und padding[1] gibt den Abstand für die Dimension x an. Jedes Paar hat als erstes Element einen niedrigen und einen hohen Abstand als zweites. Der niedrige Abstand wird in Richtung niedrigerer Indizes angewendet, während der hohe Abstand in Richtung der höheren Indizes angewendet wird. Wenn padding[1] beispielsweise (2,3) ist, enthält die zweite räumliche Dimension einen Abstand von zwei Nullen auf der linken Seite und drei Nullen rechts auf der zweiten räumlichen Dimension. Das Padding 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 die Größe des Arrays vergrößert wird. Die Löcher werden mit einem No-Op-Wert gefüllt, der für Faltung Nullen bedeutet.

Die Dilatation der RHS wird auch als atroische Faltung bezeichnet. Weitere Informationen finden Sie unter tf.nn.atrous_conv2d. Die Dilatation des Lps 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 und 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 lhs-Eingabefeature-Dimension geteilt durch feature_group_count entsprechen (also hat sie 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 miteinander verkettet.

Für die tiefe Faltung wird das Argument feature_group_count auf die Eingabefeature-Dimension 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 finden Sie 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 Dimension des Ausgabe-Batches die Größe input batch / batch_group_count haben muss. 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 muss der Größe der Dimension batch in LHs entsprechen.
  • z: Gleiche Größe wie output-z im Kernel (rhs).
  • spatial_dims: Ein Wert für jede gültige Platzierung des Faltungsfensters.

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

Die gültigen Platzierungen des Faltungsfensters werden durch die Schritte und die Größe der Grundfläche nach dem Auffüllen bestimmt.

Betrachten Sie eine 2d-Faltung und wählen Sie einige feste batch-, z-, y- und x-Koordinaten in der Ausgabe aus, um zu beschreiben, was eine Faltung bewirkt. 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 der Grundfläche, in dem jeder 2D-Punkt mit einem 1D-Vektor verknüpft ist, sodass wir eine 3D-Box erhalten. Da wir im Convolutional Kernel die Ausgabekoordinate z korrigiert haben, haben wir auch eine 3D-Box. Die beiden Boxen haben dieselben Abmessungen, sodass wir die Summe der elementweisen Produkte zwischen den beiden Feldern ermitteln können (ähnlich einem Punktprodukt). Das ist der Ausgabewert.

Wenn output-z z.B. 5, dann erzeugt jede Position des Fensters 5 Werte in der Ausgabe in der 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 einem anderen Filter vorstellen.

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

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 bei einem elementweisen static_cast in C++ einen elementweisen Konvertierungsvorgang von einer Datenform in eine Zielform aus. Die Dimensionen müssen übereinstimmen und die Konvertierung erfolgt elementweise. So werden z. B. s32-Elemente über eine s32-zu-f32-Konvertierungsroutine zu f32-Elementen.

ConvertElementType(operand, new_element_type)

Argumente Typ Semantik
operand XlaOp Array vom Typ T mit Dimmung 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 durch, 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 durch.

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 abzielt.
args Sequenz von N XlaOp s N Argumente eines beliebigen Typs, die an die Funktion übergeben werden.
shape Shape Ausgabeform der Funktion

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

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

Beispiel: 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 muss idempotent ausgeführt werden.

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 Vektorpunktprodukt
Matrix [m x k] dot Vektor [k] Vektor [m] Matrix-Vektor-Multiplikation
Matrix [m x k] dot Matrix [k x n] Matrix [m x n] 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 ausgefü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 für lhs und rhs Vertrags- und Batch-Dimensionsnummern angegeben werden.

DotDimensionNumbers-Felder Typ Semantik
lhs_contracting_dimensions wiederholter int64-Wert lhs Zahlen der Vertragsmaße
rhs_contracting_dimensions wiederholter int64-Wert rhs Zahlen der Vertragsmaße
lhs_batch_dimensions wiederholter int64-Wert lhs Batch-Dimensionsnummern
rhs_batch_dimensions wiederholter int64-Wert rhs Batch-Dimensionsnummern

DotGeneral führt die Summe der Produkte über die in dimension_numbers angegebenen Reduzierungsdimensionen aus.

Die zugehörigen vertraglich vereinbarten Dimensionsnummern von lhs und rhs müssen nicht identisch sein, müssen aber dieselben Dimensionsgrößen haben.

Beispiel für das Minimieren von 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 (Batch-Größ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 Vertragsbindung/ohne Batch und schließlich mit der Dimension rhs ohne Vertragsbindung/kein Batch.

DynamicSlice

Weitere Informationen finden Sie unter XlaBuilder::DynamicSlice.

DynamicSlice extrahiert ein Teilarray aus dem Eingabearray bei dynamischem start_indices. Die Größe des Segments in jeder Dimension wird in size_indices übergeben, das den Endpunkt exklusiver Segmentintervalle in jeder Dimension angibt: [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 des Typs T
start_indices Sequenz von N XlaOp Liste von n skalaren Ganzzahlen, die die Startindexe 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 Modulo-Dimensionsgröße nicht umbrochen wird.

Die effektiven Slice-Indizes 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 immer innerhalb der Grenzen des Operandenarrays liegt. Wenn das Slice innerhalb der Grenzen liegt, bevor die Transformation angewendet wird, hat die Transformation keine Auswirkungen.

Beispiel für eindimensionales Format:

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 des Typs T
update XlaOp N-dimensionales Array des Typs T, das die Segmentaktualisierung enthält Jede Dimension der Aktualisierungsform muss größer als null sein und die Größe von Start und Aktualisierung muss kleiner oder gleich der Operandengröße für jede Dimension sein, um zu vermeiden, dass die Aktualisierungsindexe außerhalb des gültigen Bereichs liegen.
start_indices Sequenz von N XlaOp Liste von n skalaren Ganzzahlen, die die Startindexe des Segments für jede Dimension enthalten. Wert muss größer oder gleich null sein.

Die effektiven Slice-Indizes 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 Segment immer innerhalb der Grenzen des Operandenarrays liegt. Wenn das Slice innerhalb der Grenzen liegt, bevor die Transformation angewendet wird, hat die Transformation keine Auswirkungen.

Beispiel für eindimensionales Format:

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 (Maximum), Min (Minimum), LogicalAnd (logisches UND) oder LogicalOr (logisches ODER).

Argumente Typ Semantik
lhs XlaOp Linker Operand: 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. In der Broadcasting-Dokumentation erfahren Sie, was es bedeutet, wenn Formen kompatibel sind. 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 den Wert Rem hat, wird das Vorzeichen des Ergebnisses vom Dividenden entnommen. Der absolute Wert des Ergebnisses ist immer kleiner als der absolute Wert des Divisors.

Der Ganzzahl-Divisionsüberlauf (vorzeichenbehaftete/unsignierte Division/Rest durch null oder vorzeichenbehaftete Division/Restung von INT_SMIN mit -1) erzeugt einen Implementierungswert.

Für diese Vorgänge gibt es eine alternative Variante, die Broadcasting unterschiedlicher Ränge unterstützt:

Op(lhs, rhs, broadcast_dimensions)

Dabei gilt: Op ist der gleiche wie oben. Diese Variante der Operation sollte für arithmetische Operationen zwischen Arrays mit verschiedenen Rängen verwendet werden (z. B. zum Hinzufügen einer Matrix zu einem Vektor).

Der zusätzliche broadcast_dimensions-Operand ist ein Segment von Ganzzahlen, mit dem der Rang des Operanden mit niedrigerem Rang bis zum Rang des Operanden mit höherem Rang erweitert wird. broadcast_dimensions ordnet die Abmessungen der Form mit niedrigerem Rang den Abmessungen der Form mit höherem Rang zu. Die nicht zugeordneten Abmessungen der maximierten Form werden mit Maßen der Größe 1 ausgefüllt. Beim Broadcasting nach degenerierenden Dimensionen werden dann die Formen entlang dieser Degenerationsdimensionen 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 nach IEEE 754 die Standardsemantik für den Gleitkommavergleich gilt.

Op(lhs, rhs)

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

Argumente Typ Semantik
lhs XlaOp Linker Operand: 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. In der Broadcasting-Dokumentation erfahren Sie, was es bedeutet, wenn Formen kompatibel sind. Das Ergebnis einer Operation hat eine Form, die das Ergebnis der Übertragung der beiden Eingabearrays mit dem Elementtyp PRED ist. 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 Broadcasting unterschiedlicher Ränge unterstützt:

Op(lhs, rhs, broadcast_dimensions)

Dabei gilt: Op ist der gleiche wie oben. Diese Variante des Vorgangs sollte für Vergleichsvorgänge zwischen Arrays mit unterschiedlichen Rängen verwendet werden, z. B. zum Hinzufügen einer Matrix zu einem Vektor.

Der zusätzliche broadcast_dimensions-Operand ist ein Segment von Ganzzahlen, die die Dimensionen für die Übertragung der Operanden angeben. 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) Elementweise abs. x -> |x|.

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

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

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

Floor(operand) Elementweise Etage x -> ⌊x⌋.

Imag(operand) Elementweise 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 ist, d.h. 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 elementweise logistische Funktionen x -> logistic(x).

PopulationCount(operand) Berechnet die Anzahl der Bits, die in jedem Element 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 Kehrzahl der Quadratwurzeloperation x -> 1.0 / sqrt(x).

Sign(operand) Vorgang für elementweises Signieren 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}\]

mithilfe des Vergleichsoperators des Elementtyps operand.

Sqrt(operand) Elementweiser Quadratwurzelvorgang x -> sqrt(x).

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

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

Round(operand) Elementweise Rundung, gleicht Null.

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

Argumente Typ Semantik
operand XlaOp Der Operand der Funktion

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

FFT

Der XLA-FFT-Vorgang implementiert die Vorwärts- und Inverse Fourier-Transformationen für echte 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 korrekt 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 Leitet reelle bis komplexe FFT-Werte 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 jenseits der Nyquist-Frequenz weggelassen wird.
IRFFT Umgekehrte reelle zu komplexe FFT (nimmt komplex an, 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 von 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 der Anwendung einer Kaskade von FFT-Vorgängen auf jede der innersten Achsen. Beachten Sie, dass für die reellen und komplexesten Achsen die Transformation der innersten Achse (effektiv) an erster Stelle (RFFT; zuletzt für IRFFT) ausgeführt wird. Deshalb ändert sich die Größe der innersten Achse. Andere Achsentransformationen sind dann komplex->komplex.

Implementierungsdetails

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

Zusammentragen

Der XLA-Abrufvorgang fügt mehrere Segmente eines Eingabearrays zusammen (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 uns sammeln.
start_indices XlaOp Array mit den Startindexen der erfassten 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> Die Dimensionen in der Ausgabeform, die in ein vom Operanden unterteiltes Array verschoben werden.
slice_sizes ArraySlice<int64> slice_sizes[i] sind die Grenzen für das Segment bei der Dimension i.
collapsed_slice_dims ArraySlice<int64> Die Abmessungen in jedem Segment, die weggeklappt sind. Diese Abmessungen müssen Größe 1 haben.
start_index_map ArraySlice<int64> Eine Map, die beschreibt, wie Indexe in start_indices den rechtlichen Indizes in Operanden zugeordnet werden.
indices_are_sorted bool Gibt an, ob die Indizes garantiert vom Aufrufer sortiert werden.
unique_indices bool Gibt an, ob die Indizes für den Aufrufer garantiert eindeutig sind.

Der Einfachheit halber benennen wir Dimensionen im Ausgabearray, die nicht in offset_dims enthalten sind, 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 (d.h., wenn start_indices die Form [6,7] hat und index_vector_dim 2 ist, wird die Form von start_indices implizit als [6,7,1] betrachtet.

Die Grenzen des Ausgabearrays 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. offset_dims[k] für einige k), wird die entsprechende Grenze von slice_sizes nach Berücksichtigung von collapsed_slice_dims ausgewählt. Das heißt, wir wählen adjusted_slice_sizes[k] aus, wobei adjusted_slice_sizes gleich slice_sizes ist und die Grenzen bei den Indexen collapsed_slice_dims entfernt wurden.

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

  1. Lassen Sie G = { Out[k] für k in batch_dims }. Verwenden Sie G, um einen Vektor S so herauszuschneiden, dass S[i] = start_indices[Zusammenführen(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, gilt S = start_indices.

  2. Erstellen Sie einen Startindex Sin in operand mit S. 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 in den Versatzdimensionen in Out gemäß dem Satz von collapsed_slice_dims verteilen. Genauer gesagt:

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

    2. OAndernfalls in[_] = 0.

  4. In ist Oin + Sin, wobei + für die elementweise Addition 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}. remapped_offset_dims ist dann {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 Nutzer sortiert sind. Andernfalls ist die Semantik die 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 verstreut sind, nicht eindeutig sind, ist die Semantik die Implementierung definiert.

Informelle Beschreibung und Beispiele

Inoffiziell entspricht jeder 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 abzurufen.

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

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

  • Wir passen die Form des Segments an, 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 Segmente der Form [8,6] aus einem [16,11]-Array gesammelt werden, um einen Eindruck davon zu bekommen, wie all dies zusammenpasst. Die Position eines Segments im Array [16,11] kann als Indexvektor der Form S64[2] dargestellt werden, sodass der Satz 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, und diesen folgendermaßen 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 bei Index [X+O0,Y+O1].

slice_sizes ist [8,6], wodurch der Bereich von O0 und O1 bestimmt wird, was wiederum die Grenzen des Segments bestimmt.

Dieser Erfassungsvorgang fungiert als dynamisches Batchsegment mit G als Batchdimension.

Die Indizes zum Zusammenführen können mehrdimensional sein. In einer allgemeineren Version des obigen Beispiels mit einem Array zum Sammeln von Indexen mit der Form [4,5,2] werden Indexe beispielsweise so übersetzt:

Auch dies fungiert als ein dynamisches Batch-Segment G0 und G1 als Batchdimensionen. Die Segmentgröße beträgt immer noch [8,6].

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

  1. Wir können konfigurieren, welche Dimensionen in der Ausgabeform die Versatzabmessungen sind (Dimensionen, die im letzten Beispiel O0, O1 enthalten). Die Ausgabe-Batch-Dimensionen (Dimensionen, die im letzten Beispiel G0 und G1 enthalten) sind als Ausgabedimensionen definiert, die keine Offset-Dimensionen 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 Segmentgröße von 1 haben, ist der einzige gültige Index für sie 0. Das Ausschließen dieser Segmente führt nicht zu Mehrdeutigkeit.

  3. Das aus dem Array "Sammel-Indexe" (im letzten Beispiel X, Y) extrahierte Segment hat möglicherweise weniger Elemente als der Eingabearrayrang. Eine explizite Zuordnung legt fest, wie der Index erweitert werden soll, damit er den gleichen Rang wie die Eingabe hat.

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

G0 und G1 werden verwendet, um einen Startindex wie gewohnt aus dem Array für die Erfassung der Indexe herauszuschneiden, mit der Ausnahme, dass der Startindex nur ein Element, X, enthält. Analog dazu gibt es nur einen Ausgabe-Offset-Index mit dem Wert O0. Bevor sie jedoch als Indizes im Eingabearray verwendet werden, werden sie gemäß „Sammel-Indexzuordnung“ (start_index_map in der formalen Beschreibung) und „Offset Mapping“ (remapped_offset_dims in der formalen Beschreibung) zu [X,0] bzw. [0,O0] erweitert, was bis zu [X,O0] ergibt. Mit anderen Worten: Der Ausgabeindex [O1], der den Ausgabeindex [G1], der den Ausgabeindex [G1] zu [G1], der den Ausgabeindex [G1GatherIndices aus [G1], {2, der Eingabe-Index [G1GatherIndices ergibt, {2, der Eingabe-Index [G1], der den Ausgabeindex [G1GatherIndices aus [G1], der Ausgabeindex [G1GatherIndices ergibt, bevor sie als Indexe im Eingabearray verwendet werden, erweitert.00000Otf.gather_nd

slice_sizes ist für diesen Fall [1,11]. Dies bedeutet, dass jeder Index-X im Array der Indexe abrufen 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 angegebenen Dimension des Operanden zurück. Der Operand muss Array-Form 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 gegebenen XlaOp-Dimension fest. Der Operand muss Array-Form sein.

SetDimensionSize(operand, size, dimension)

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

Sie leiten den Operanden als Ergebnis weiter, 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 in 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 InFeed-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-Schnittstelle 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 jedoch eine Gesamtreihenfolge vorhanden sein. Beispiel: Zwei Infeeds im folgenden Code haben 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 managementfrei und wird fortgesetzt, ohne Daten aus dem Infeed des Geräts zu lesen.

Iota

Weitere Informationen finden Sie unter XlaBuilder::Iota.

Iota(shape, iota_dimension)

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

Argumente Typ Semantik
shape Shape Form des durch Iota() erstellten Arrays
iota_dimension int64 Die Dimension, die erhöht werden soll.

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 Sequenz 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 vom Typ T und M beliebigen Typs
dimensions int64-Array Array mit Kartenmaßen

Wendet eine Skalarfunktion auf die angegebenen operands-Arrays an und erzeugt ein Array derselben Dimensionen, wobei jedes Element das Ergebnis der zugeordneten Funktion ist, die auf die entsprechenden Elemente in den Eingabearrays angewendet wurde.

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 Abmessungen wie die Operanden, mit der Ausnahme, dass der Elementtyp T durch S ersetzt wird.

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

OptimizationBarrier

Blockiert alle Optimierungsdurchläufe, sodass Berechnungen über die Barriere hinweg erfolgen.

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 hinzugefügten Abstand auszufüllen
padding_config PaddingConfig Abstand an beiden Rändern (niedrig, hoch) und zwischen den Elementen jeder Abmessung

Erweitert 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 den 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 Randabstands 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. Das Innen-Padding erfolgt logisch vor dem Rand-Padding. Bei einem negativen Rand-Padding werden also Elemente aus dem Operanden mit Auffüllung entfernt.

Dieser Vorgang ist ein Leerbefehl, wenn die Paare zum Auffüllen der Ränder alle (0, 0) und die Werte des Innenrands 0 sind. Die folgende Abbildung zeigt Beispiele für verschiedene edge_padding- und interior_padding-Werte für ein zweidimensionales Array.

Recv

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 mit der angegebenen Form von einer Send-Anweisung in einer anderen Berechnung mit demselben Kanal-Handle. Gibt ein 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 einem RecvDone-Befehl verwendet werden.

RecvDone(HloInstruction context)

Wartet bei einem durch eine Recv-Anweisung erstellten Kontext, bis die Datenübertragung abgeschlossen ist, 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 Ungeordnetes Array von Dimensionen, 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 Dimensionen 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 diese Operation 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 Abmessungen unten beschrieben werden.

Die Berechnung der Reduktion kann von unterschiedlichen 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 Anwendungsfälle fast aus, um assoziativ zu sein.

Beispiele

Bei der Reduzierung einer Dimension in einem einzelnen 1D-Array mit den Werten [10, 11, 12, 13] und der Reduktionsfunktion f (dies ist computation) kann dies wie folgt berechnet werden:

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

Es gibt aber auch 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 die Reduktion implementiert werden kann, wobei die Addition 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, Dimension 0 der Größe 2 und Dimension 1 der Größe 3:

Ergebnisse der Reduzierung von Dimensionen 0 oder 1 mit einer Funktion zum Hinzufügen:

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

Für ein komplexeres Beispiel sehen Sie 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 über die Dimension 0 hinweg repliziert.

Ähnlich wie im 2D-Beispiel können wir nur eine Dimension reduzieren. Wenn wir beispielsweise die Dimension 0 reduzieren, erhalten wir ein Array vom Typ Rang 2, in 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, in dem alle Werte in 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 bleibt in der Ausgabe erhalten, 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, die addierend reduziert werden, ergibt sich das 1D-Array [20, 28, 36].

Wenn das 3D-Array über alle seine Abmessungen reduziert wird, ergibt sich der skalare 84.

Variadische Reduzierung

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

  • Laufender reduzierter Wert für ersten Operanden
  • ...
  • Aktiver reduzierter Wert für den N-ten Operanden
  • Eingabewert für den ersten Operanden
  • ...
  • Eingabewert für den n-ten Operanden

Betrachten Sie beispielsweise die folgende Reduktionsfunktion, mit der der Maximalwert und der Wert 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 1D-Eingabearrays V = Float[N], K = Int[N] und Initialisierungswerten I_V = Float, I_K = Int entspricht das Ergebnis f_(N-1) der Reduzierung für 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 diese Reduzierung auf ein Array von Werten und ein Array mit sequenziellen Indizes (d.h. iota) angewendet wird, erfolgt die Zusammenarbeit über die Arrays 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 Umwandlung von Gleitkommawerten in ein Format mit geringerer Genauigkeit (z. B. IEEE-FP16) und wieder zurück in das Originalformat. 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 in einem Format mit geringerer Genauigkeit

Das Ergebnis ist ein Array vom Typ T. Die Eingabewerte werden auf den nächsten Wert gerundet, der mit der angegebenen Anzahl von Mantissenbits dargestellt werden kann (mithilfe der Semantik "Gleichheitsgrad" verwendet wird). 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 enthalten und eine nicht negative Anzahl von Mantissenbits haben, um einen Nullwert von einem unendlichen Wert zu unterscheiden. Die Anzahl der Exponenten- oder Mantissen-Bits kann den entsprechenden Wert für den Typ T überschreiten. Der entsprechende Teil der Konvertierung ist dann eine Nulloperation.

ReduceScatter

Weitere Informationen finden Sie unter XlaBuilder::ReduceScatter.

ReduceScatter ist ein kollektiver Vorgang, bei dem AllReduce effektiv ausgeführt und das Ergebnis dann in shard_count-Blöcke entlang des 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, die in allen Replikaten reduziert werden sollen.
computation XlaComputation Reduktionsberechnung
scatter_dimension int64 Zu streuende Dimension.
shard_count int64 Anzahl der zu teilenden Blöcke: 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 die modulübergreifende Kommunikation
  • Wenn operand ein Tupel von Arrays ist, wird die Streuung zum Reduzieren für jedes Element des Tupels ausgeführt.
  • replica_groups ist eine Liste von Replikatgruppen, zwischen denen die Reduzierung durchgeführt wird (Replikat-ID für das aktuelle Replikat kann mit ReplicaId abgerufen werden). Die Reihenfolge der Replikate in den einzelnen Gruppen 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 die Anzahl der Elemente muss mit der Anzahl der Replikate übereinstimmen. Wenn mehrere Replikatgruppen vorhanden sind, müssen alle dieselbe Größe haben. Beispiel: replica_groups = {0, 2}, {1, 3} führt eine Reduzierung zwischen den Replikaten 0 und 2 sowie 1 und 3 durch und verteilt 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 mit scatter_dimension shard_count-mal kleiner. Wenn beispielsweise zwei Replikate vorhanden sind und der Operand auf den beiden Replikaten den Wert [1.0, 2.25] bzw. [3.0, 5.25] hat, ist der Ausgabewert dieses Vorgangs mit 0 für das erste Replikat [4.0] und für das zweite Replikat [7.5].scatter_dim

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 ausgegeben wird. Jedes Ausgabearray hat dieselbe 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 dem angewendeten computation immer die init_values auf der linken Seite übergeben.

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

Argumente Typ Semantik
operands N XlaOps Eine Sequenz aus n mehrdimensionalen Arrays vom Typ T_0,..., T_{N-1}, von denen jedes die Grundfläche darstellt, auf der sich das Fenster befindet.
init_values N XlaOps Die 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 Werte der Basiserweiterung
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 kein Auffüllen verwendet und das Fenster „stoppt“, sobald es nicht mehr passt)

Wobei:

  • N muss größer oder gleich 1 sein.
  • Alle Eingabearrays müssen dieselben Dimensionen 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 "1" in einer Dimension gibt an, dass die Position eines Fensters in der Dimension 1 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 die gleichen, als ob die Eingabe mit den Dimensionen, die sie nach dem Padding hat, enthalten wäre.

Für ein nicht triviales Padding-Beispiel könnten Sie den Mindestwert für das Reduzieren des Fensters berechnen (der Anfangswert ist MAX_FLOAT) mit der Dimension 3 und dem Schritt 2 über das Eingabearray [10000, 1000, 100, 10, 1]. Das Auffüllen von kValid berechnet Mindestwerte über zwei gültige Fenster: [10000, 1000, 100] und [100, 10, 1], was zur Ausgabe [100, 1] führt. Durch das Auffüllen von kSame wird das Array zuerst so aufgefüllt, dass die Form nach dem Reduce-Window dieselbe ist wie die Eingabe für Schritt 1, indem anfängliche Elemente auf beiden Seiten hinzugefügt werden und [MAX_VALUE, 10000, 1000, 100, 10, 1, MAX_VALUE] abgerufen wird. Wenn Sie das Reduce-Window über das aufgefüllte Array hinweg ausführen, werden die drei Fenster [MAX_VALUE, 10000, 1000], [1000, 100, 10], [10, 1, MAX_VALUE] und Renditen [1000, 10, 1] ausgeführt.

Die Auswertungsreihenfolge der Reduktionsfunktion ist beliebig und kann nicht deterministisch sein. Daher sollte die Reduktionsfunktion nicht allzu empfindlich auf Neuzuordnungen 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.

Formt 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 für int64 Reihenfolge, in der Dimensionen minimiert werden
new_sizes Vektor für int64 Vektor mit Größen neuer Dimensionen

Konzeptionell wird ein Array bei der Umformung 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 von Dimensionsindexen und ein Compile-Zeitkonstanten-Vektor der Dimensionsgrößen für das Ergebnis. Die Werte im Vektor dimension, falls angegeben, müssen eine Permutation aller T-Dimensionen sein. Wenn nicht angegeben, muss der Standardwert {0, ..., rank - 1} sein. Die Dimensionen in dimensions werden in der Schleifenverschachtelung von der langsamsten Dimension (größte große) bis zur am schnellsten variierenden Dimension (größtenteils) in der Schleifenverschachtelung sortiert, 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 Dimensionen new_size muss dem Produkt der Dimensionsgrößen des Operanden entsprechen. Beim Verfeinern des minimierten Arrays in das mehrdimensionale Array, das durch new_sizes definiert ist, werden die Dimensionen in new_sizes von der langsamsten (größten) bis zur schnellsten Variante (sehr geringfügig) sortiert.

Nehmen wir z. B. „v“ als 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 die Umformung ein Array mit einem einzelnen Element in einen Skalar umwandeln 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 im Ausgabearray in einem transformierten Index gespeichert. Der mehrdimensionale Index wird transformiert, indem der Index in jeder Dimension umgedreht wird. Wenn also eine Dimension der Größe N eine der Umkehrdimensionen ist, wird ihr Index i in N - 1 - i transformiert.

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.

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

RngNormal(mu, sigma, shape)

Argumente Typ Semantik
mu XlaOp Skalar des Typs T, der den Mittelwert der 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.

Erstellt eine Ausgabe einer bestimmten Form mit Zufallszahlen, die anhand der gleichmäßigen Verteilung über das Intervall \([a,b)\)generiert werden. Die Parameter und der Ausgabeelementtyp müssen boolesche, ganzzahlige oder Gleitkommatypen sein und die Typen müssen 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 skalarwert sein. Wenn \(b <= a\) , ist das Ergebnis implementierungsdefiniert.

RngUniform(a, b, shape)

Argumente Typ Semantik
a XlaOp Skalar vom Typ T, der untere Intervallgrenze angibt
b XlaOp Skalar des Typs T, der den oberen Grenzwert des Intervalls angibt
shape Shape Ausgabeform vom Typ T

RngBitGenerator

Erzeugt unter Verwendung des angegebenen Algorithmus (oder Back-End-Standards) eine Ausgabe mit einer bestimmten Form, die mit einheitlichen Zufallsbits gefüllt ist, und gibt einen aktualisierten Zustand (mit der gleichen Form wie der Anfangszustand) und 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.

Die Ausgabe ist garantiert eine deterministische Funktion des Anfangszustands, aber es ist nicht garantiert, 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 Anfangsstatus 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 den Werten des Eingabearrays operands entsprechen, wobei mehrere Segmente (bei durch scatter_indices angegebenen Indexen) mit der Reihenfolge der Werte in updates mithilfe von update_computation 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 sollen.
scatter_indices XlaOp Array mit den Anfangsindexen 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 zum Kombinieren der vorhandenen Werte im Eingabearray und der Aktualisierungen während der Streuung. Diese Berechnung sollte 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 Abmessungen in der updates-Form, die Fensterabmessungen sind.
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 zwischen den Streuindizes und dem Operandenindexbereich Dieses Array wird als Zuordnung von i zu scatter_dims_to_operand_dims[i] interpretiert . Sie muss eine Eins-zu-eins-eins-zu-eins und eine Gesamtsumme darstellen.
indices_are_sorted bool Gibt an, ob die Indizes garantiert vom Aufrufer sortiert werden.

Wobei:

  • N muss größer oder gleich 1 sein.
  • operands[0], ..., operands[N-1] müssen alle dieselbe Größe haben.
  • updates[0], ..., updates[N-1] müssen alle dieselbe Größe 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, betrachten wir scatter_indices implizit als eine nachgestellte 1-Dimension.

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

Für die Argumente der Streuung gelten folgende Einschränkungen:

  • 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 so aussehen:

    • Wenn i in update_window_dims vorhanden ist (d. h. für einen k-Wert update_window_dims[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] enthält die Grenzen von operand, wobei die Grenzen bei den Indexen inserted_window_dims entfernt wurden.adjusted_window_bounds
    • Wenn i in update_scatter_dims vorhanden ist (d.h. für einen k-Wert update_scatter_dims[k]), muss die Grenze der Dimension i in updates der entsprechenden Grenze von scatter_indices entsprechen, wobei index_vector_dim übersprungen wird (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 und darf keine sich wiederholenden Dimensionsnummern enthalten. Außerdem muss der Wert im Bereich [0, updates.rank) liegen.

  • inserted_window_dims muss in aufsteigender Reihenfolge angegeben werden und darf keine sich wiederholenden Dimensionsnummern enthalten. Außerdem muss der Wert 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 seine 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 nachzuschlagen, 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. Dazu teilen Sie S mit der scatter_dims_to_operand_dims-Karte auf. 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 enthalten 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 die elementweise Addition 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 Indexe in updates auf denselben Index in operands verweisen, ist der entsprechende Wert in output nicht deterministisch.

Der erste Parameter, der an update_computation übergeben wird, ist immer der aktuelle Wert aus dem Array output. Der zweite Parameter ist immer der Wert aus dem Array updates. Dies ist vor allem 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 Nutzer sortiert sind. Andernfalls ist die Semantik die Implementierung definiert.

Informellerweise kann der Streuvorgang als Inverse des Sammelvorgangs angesehen werden, d.h., er aktualisiert die Elemente in der Eingabe, die vom entsprechenden Collect-Vorgang extrahiert werden.

Eine detaillierte informelle Beschreibung und Beispiele finden Sie im Abschnitt „Informelle Description“ unter Gather.

Auswählen

Weitere Informationen finden Sie unter XlaBuilder::Select.

Erstellt ein Ausgabearray aus Elementen zweier 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 verwendet, wenn pred gleich true ist, und aus on_false, wenn pred gleich false ist.

Beispiel mit einem nicht skalaren 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 dem skalaren 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 gelten für diesen Zweck als skalare Typen. 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 operand-Array berechnet wird, um ein Element aus jedem Fenster auszuwählen, und dann das source-Array auf die Indizes der ausgewählten Elemente verteilt wird, 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 gibt false zurück, 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, sodass das ausgewählte Element nicht von der Reihenfolge der Elemente abhängt, die in einem bestimmten Fenster durchlaufen werden.

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

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

Die Funktion kombiniert die beiden Parameter und gibt einen skalaren Wert 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 Array operand und das Array source muss dieselbe Form haben, wie das Ergebnis nach Anwendung eines ReduceWindow-Vorgangs auf das Array operand. SelectAndScatter kann verwendet werden, um die Gradientenwerte für eine Pooling-Ebene in einem neuronalen Netzwerk rückzupropagieren.

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 gleiten
select XlaComputation Binäre Berechnung vom Typ 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 Padding-Typ 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 vom Typ T, T -> T, um jedes Streuquellenelement mit seinem 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 überschneiden, wie in Abbildung (2) unten gezeigt, 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äre Additionsfunktion scatter erzeugt das Ausgabeelement des Werts 8 (2 + 6).

Die Auswertungsreihenfolge der scatter-Funktion ist beliebig und kann nicht deterministisch sein. Daher sollte die Funktion scatter nicht allzu 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. Es werden keine Daten zurückgegeben.

Ähnlich wie beim Recv-Vorgang stellt die Client-API des Send-Vorgangs 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 Ressourcen, die durch den Befehl Recv mit derselben Kanal-ID zugewiesen wurden. 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 einem SendDone-Befehl verwendet werden.

SendDone(HloInstruction context)

Bei einem mit einer Send-Anweisung erstellten Kontext wird auf den Abschluss der Datenübertragung gewartet. Die Anweisung gibt keine Daten zurück.

Planung der Kanalanleitung

Die Ausführungsreihenfolge der 4 Anweisungen 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, darf es keine Zyklen zwischen den Berechnungen geben. Die folgenden Zeitpläne führen beispielsweise zu Deadlocks.

Slice

Weitere Informationen finden Sie unter XlaBuilder::Slice.

Beim Aufteilen 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 Dimensionen und Indexe des Begrenzungsrahmens als Argumente für den Aufteilungsvorgang angegeben werden.

Slice(operand, start_indices, limit_indices, strides)

Argumente Typ Semantik
operand XlaOp N-dimensionales Array des Typs T
start_indices ArraySlice<int64> Liste von N Ganzzahlen, die die Startindexe des Segments für jede Dimension enthalten. Werte müssen größer oder gleich null sein.
limit_indices ArraySlice<int64> Liste von N Ganzzahlen mit den Endindexen (ausschließlich) für das Segment für jede Dimension. 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 von N Ganzzahlen, die den Eingabeschritt des Segments bestimmt. Im Segment wird jedes strides[d]-Element der Dimension d ausgewählt.

Beispiel für eindimensionales Format:

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 zu sortierenden Operanden
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. Formell, nachdem das Array sortiert wurde, gilt es 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 (einer Matrix) wird beispielsweise mit dem Dimensionswert 0 jede Spalte unabhängig sortiert und bei einem Dimensionswert von 1 jede Zeile unabhängig voneinander. Wenn keine Dimensionsnummer angegeben ist, 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. Die Operanden werden konzeptionell als Tupel behandelt. Bei der Prü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 k-th-Operanden und der Parameter 2 * k + 1 dem Wert an Position j vom k-th-Operanden entspricht. Normalerweise würde der Vergleichsoperator also 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 i-th-Operand des Tupels entspricht dem i-th-Operanden 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 einem Kleiner-als-Wert 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. Das heißt, wenn Elemente vorhanden sind, die vom Vergleichsoperator als gleich angesehen werden, wird die relative Reihenfolge der gleichen Werte beibehalten. Die zwei Elemente e1 und e2 sind genau dann gleich, wenn comparator(e1, e2) = comparator(e2, e1) = false. Standardmäßig ist is_stable auf "false" gesetzt.

Transponieren

Weitere Informationen finden Sie unter tf.reshape.

Transpose(operand)

Argumente Typ Semantik
operand XlaOp Der Operand, der transponiert werden soll.
permutation ArraySlice<int64> So passen Sie die Abmessungen an.

Aktiviert 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 dreieckigen Koeffizientenmatrizen durch Vorwärts- oder RückwärtsSubstitution. Diese Routine wird entlang führender Dimensionen übertragen und löst eines der Matrixsysteme op(a) * x = b oder x * op(a) = b für die Variable x unter der Voraussetzung von 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 aus dem 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, sonst [..., 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 Gibt an, ob das obere oder untere Dreieck von a verwendet werden soll.
unit_diagonal bool bei true wird angenommen, dass die diagonalen Elemente von a 1 sind und es nicht darauf zugegriffen wird.
transpose_a Transpose Ob a in der vorliegenden Form verwendet werden soll oder ob das Transponieren oder Transponieren verwendet werden soll.

Eingabedaten werden je nach Wert von lower nur aus dem unteren/oberen Dreieck von a gelesen. Werte aus dem anderen Dreieck werden ignoriert. Ausgabedaten werden im selben 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 mit Ausnahme der beiden 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 GetTupleElement-Vorgang zerlegt (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, die den Hauptteil der Schleife definiert.
init T Anfangswert für den Parameter condition und body.

Führt body sequenziell aus, bis condition fehlschlägt. Dies ähnelt einer typischen „Während“-Schleife 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 nachfolgenden Durchlauf automatisch auf das neue Ergebnis von body aktualisiert.

Ein Hauptanwendungsfall des While-Knotens besteht darin, die wiederholte Ausführung des Trainings in neuronalen Netzwerken zu implementieren. Unten ist 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 weiterhin 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};
}