Vorgangssemantik

bookmark_border Mit Sammlungen den Überblick behalten Sie können Inhalte basierend auf Ihren Einstellungen speichern und kategorisieren.

Im Folgenden wird die Semantik der Vorgänge beschrieben, die in der XlaBuilder-Benutzeroberfläche definiert sind. Normalerweise werden diese Vorgänge Eins-zu-Eins den in der RPC-Schnittstelle in xla_data.proto definierten Vorgängen zugeordnet.

Hinweis zur Nomenklatur: Der generalisierte Datentyp, mit dem XLA arbeitet, ist ein n-dimensionales Array mit Elementen eines einheitlichen Typs (z. B. 32‑Bit-Float). In der gesamten Dokumentation wird Array verwendet, um ein Array mit beliebiger Dimension zu bezeichnen. Für den praktischen Gebrauch haben Sonderfälle 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 entgegen und gibt ein einzelnes Token aus. Tokens sind primitive Typen, die zwischen Operationen mit Nebenwirkungen eingeschoben werden können, um die Reihenfolge zu erzwingen. AfterAll kann als Token-Join verwendet werden, um eine Operation nach einer Set-Operation anzuordnen.

AfterAll(operands)

AllGather

Weitere Informationen finden Sie unter XlaBuilder::AllGather.

Führt eine Konkatenierung über Replikate hinweg aus.

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

• replica_groups ist eine Liste von Replikagruppen, zwischen denen die Koncatenate ausgeführt wird. Die Replika-ID für das aktuelle Replikat kann mit ReplicaId abgerufen werden. Die Reihenfolge der Replicas in jeder Gruppe bestimmt die Reihenfolge, in der ihre Eingaben im Ergebnis enthalten sind. 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. Beispiel: replica_groups = {0, 2}, {1, 3} führt eine Koncatenate zwischen den Replikas 0 und 2 sowie 1 und 3 aus.
• shard_count ist die Größe der einzelnen Replikationsgruppen. Wir benötigen diese Informationen, wenn replica_groups leer ist.
• channel_id wird für die moduleübergreifende Kommunikation verwendet: Nur all-gather-Vorgänge mit demselben channel_id können miteinander kommunizieren.

Die Ausgabeform ist die Eingabeform, bei der das all_gather_dim shard_count-mal größer ist. Wenn es beispielsweise zwei Replikate gibt und der Operand in den beiden Replikaten die Werte [1.0, 2.5] und [3.0, 5.25] hat, ist der Ausgabewert dieser Operation, bei der all_gather_dim = 0 ist, in beiden Replikaten [1.0, 2.5, 3.0, 5.25].

AllReduce

Weitere Informationen finden Sie unter XlaBuilder::AllReduce.

Führt eine benutzerdefinierte Berechnung über alle Replikate hinweg aus.

AllReduce(operand, computation, replica_group_ids, channel_id)

• Wenn operand ein Tupel von Arrays ist, wird die All-Reduce-Operation auf jedes Element des Tupels angewendet.
• replica_groups ist eine Liste von Replikationsgruppen, zwischen denen die Reduzierung durchgeführt wird. Die Replikations-ID für das aktuelle Replik kann mit ReplicaId abgerufen werden. replica_groups muss entweder leer sein (in diesem Fall gehören alle Repliken zu einer einzigen Gruppe) oder dieselbe Anzahl von Elementen wie die Anzahl der Repliken enthalten. Beispiel: replica_groups = {0, 2}, {1, 3} führt eine Reduzierung zwischen den Replicas 0 und 2 sowie 1 und 3 durch.
• channel_id wird für die moduleübergreifende Kommunikation verwendet: Nur all-reduce-Vorgänge mit demselben channel_id können miteinander kommunizieren.

Die Ausgabeform entspricht der Eingabeform. Wenn es beispielsweise zwei Repliken gibt und der Operand in den beiden Repliken jeweils den Wert [1.0, 2.5] und [3.0, 5.25] hat, ist der Ausgabewert dieses Operators und der Summenberechnung in beiden Repliken [4.0, 7.75]. Wenn die Eingabe ein Tupel ist, ist die Ausgabe auch ein Tupel.

Für die Berechnung des Ergebnisses von AllReduce ist eine Eingabe von jedem Replikat erforderlich. Wenn ein Replikat einen AllReduce-Knoten also häufiger ausführt als ein anderes, wartet das erste Replikat ewig. Da auf allen Replicas dasselbe Programm ausgeführt wird, gibt es nicht viele Möglichkeiten, wie das passieren kann. Es ist jedoch möglich, wenn die Bedingung einer While-Schleife von Daten aus dem Infeed abhängt und die Daten im Infeed dazu führen, dass die While-Schleife in einem Replikat mehrmals als in einem anderen iteriert.

AllToAll

Weitere Informationen finden Sie unter XlaBuilder::AllToAll.

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

1. Die Streuphase. Auf jedem Kern wird der Operand entlang der split_dimensions in split_count Blöcke aufgeteilt und die Blöcke werden auf alle Kerne verteilt, z.B. wird der i-te Block an den i-ten Kern gesendet.
2. Die Phase der Informationsbeschaffung. Jeder Kern verknüpft die empfangenen Blöcke entlang der concat_dimension.

Die teilnehmenden Kerne können so konfiguriert werden:

• replica_groups: Jede ReplicaGroup enthält eine Liste der Replikat-IDs, die an der Berechnung teilnehmen. Die Replikat-ID für das aktuelle Replikat kann mit ReplicaId abgerufen werden. AllToAll wird innerhalb von Untergruppen in der angegebenen Reihenfolge angewendet. Beispiel: replica_groups = { {1,2,3}, {4,5,0} } bedeutet, dass eine AllToAll-Operation innerhalb der Repliken {1, 2, 3} und in der Sammelphase angewendet wird und die empfangenen Blöcke in der Reihenfolge 1, 2, 3 zusammengefügt werden. Anschließend wird eine weitere AllToAll-Berechnung auf die Repliken 4, 5 und 0 angewendet. Die Reihenfolge der Zusammenführung ist ebenfalls 4, 5 und 0. Wenn replica_groups leer ist, gehören alle Replikatdateien derselben Gruppe an, in der Reihenfolge, in der sie erscheinen.

Voraussetzungen:

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

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

Unten sehen Sie ein Beispiel für Alltoall.

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

In diesem Beispiel nehmen 4 Kerne an der Alltoall-Operation teil. Auf jedem Kern wird der Operand entlang der Dimension 1 in vier Teile aufgeteilt. Jeder Teil hat also die Form f32[4,4]. Die vier Teile werden auf alle Kerne verteilt. Anschließend verknüpft jeder Kern die empfangenen Teile entlang der Dimension 0 in der Reihenfolge von Kern 0 bis 4. Die Ausgabe jedes Kerns hat also die Form f32[16,4].

BatchNormGrad

Eine ausführliche Beschreibung des Algorithmus finden Sie auch unter XlaBuilder::BatchNormGrad und im ursprünglichen Artikel zur Batchnormalisierung.

Berechnet die Gradienten der Batchnorm.

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

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

Die drei Gradienten werden durch die folgenden Formeln definiert (vorausgesetzt, es wird ein vierdimensionales Array als operand mit dem Index der Featuredimension l, der Batchgröße m und den räumlichen Größen w und h verwendet):

$$\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 für Batch- und räumliche Dimensionen dar.

Der Ausgabetyp ist ein Tupel mit drei Handles:

BatchNormInference

Eine ausführliche Beschreibung des Algorithmus finden Sie auch unter XlaBuilder::BatchNormInference und im ursprünglichen Artikel zur Batchnormalisierung.

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

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

Für jedes Merkmal in der Merkmalsdimension (feature_index ist der Index für die Merkmalsdimension in operand) werden der Mittelwert und die Varianz über alle anderen Dimensionen berechnet. Anhand dieser Werte wird jedes Element in operand normalisiert. feature_index muss ein gültiger Index für die Dimensionseigenschaft in operand sein.

BatchNormInference entspricht dem Aufruf von BatchNormTraining, ohne dass mean und variance für jeden Batch berechnet werden. Stattdessen werden die Eingaben mean und variance als geschätzte Werte verwendet. Der Zweck dieser Operation besteht darin, die Latenz bei der Inferenz zu reduzieren. Daher der Name BatchNormInference.

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

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)

Für jedes Merkmal in der Merkmalsdimension (feature_index ist der Index für die Merkmalsdimension in operand) werden der Mittelwert und die Varianz über alle anderen Dimensionen berechnet. Anhand dieser Werte wird jedes Element in operand normalisiert. feature_index muss ein gültiger Index für die Dimensionseigenschaft in operand sein.

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

• Hier wird die Batch-Mittelung $\mu_l$ für jedes Merkmal l in der Merkmaldimension berechnet: $\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}$

• Batchvarianz berechnen $\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 Epsilonwert, in der Regel eine kleine Zahl, wird hinzugefügt, um Divisionsfehler durch Null zu vermeiden.

Der Ausgabetyp ist ein Tupel aus drei XlaOps:

batch_mean und batch_var sind Momente, die anhand der oben genannten Formeln über die Batch- und räumlichen Dimensionen berechnet werden.

BitcastConvertType

Weitere Informationen finden Sie unter XlaBuilder::BitcastConvertType.

Ähnlich wie bei einer tf.bitcast in TensorFlow wird eine elementweise Bitcast-Operation von einer Datenform in eine Zielform ausgeführt. Eingabe- und Ausgabegröße müssen übereinstimmen: Beispiel: s32-Elemente werden über die Bitcast-Routine in f32-Elemente umgewandelt und ein s32-Element wird in vier s8-Elemente umgewandelt. Bitcast wird als Low-Level-Cast implementiert, sodass Maschinen mit unterschiedlichen Gleitkommadarstellungen unterschiedliche Ergebnisse liefern.

BitcastConvertType(operand, new_element_type)

Die Dimensionen des Operanden und der Zielform müssen mit Ausnahme der letzten Dimension übereinstimmen. Diese ändert sich entsprechend dem Verhältnis der ursprünglichen Größe vor und nach der Umwandlung.

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

Bitcast-Konvertierung in einen primitiven Typ mit anderer Breite

Die BitcastConvert-HLO-Anweisung unterstützt den Fall, dass die Größe des Ausgabeelementtyps T' nicht der Größe des Eingabeelements T entspricht. Da der gesamte Vorgang konzeptionell ein Bitcast ist und die zugrunde liegenden Bytes 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.

Bei B > B' erhält die Ausgabeform zuerst eine neue kleinste Dimension mit der Größe B/B'. Beispiel:

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

Für effektive Skalare gilt dieselbe Regel:

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

Alternativ muss für B' > B die letzte logische Dimension der Eingabeform B'/B entsprechen. Diese Dimension wird bei der Umwandlung entfernt:

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

Beachten Sie, dass die Umwandlungen zwischen verschiedenen Bitbreiten nicht elementweise erfolgen.

Nachricht an alle

Weitere Informationen finden Sie unter XlaBuilder::Broadcast.

Fügen Sie einem Array Dimensionen hinzu, indem Sie die Daten im Array duplizieren.

Broadcast(operand, broadcast_sizes)

Die neuen Dimensionen werden links eingefügt. Wenn broadcast_sizes die Werte {a0, ..., aN} hat und die Operandenform die Dimensionen {b0, ..., bM}, hat die Ausgabeform die Dimensionen {a0, ..., aN, b0, ..., bM}.

Die neuen Dimensionen werden in Kopien des Operanden indexiert, d.h.

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

Wenn operand beispielsweise ein Skalar f32 mit dem Wert 2.0f und broadcast_sizes = {2, 3} ist, ist das Ergebnis ein Array mit der Form f32[2, 3] und allen Werten 2.0f.

BroadcastInDim

Weitere Informationen finden Sie unter XlaBuilder::BroadcastInDim.

Erweitert die Größe und Anzahl der Dimensionen eines Arrays, indem die Daten im Array dupliziert werden.

BroadcastInDim(operand, out_dim_size, broadcast_dimensions)

Ähnlich wie „Broadcast“, aber es können überall Dimensionen hinzugefügt und vorhandene Dimensionen mit der Größe „1“ erweitert werden.

Die operand wird an die durch out_dim_size beschriebene Form gesendet. Bei broadcast_dimensions werden die Dimensionen von operand den Dimensionen des Zieltyps zugeordnet. Das bedeutet, dass die i. Dimension des Operanden der i. Dimension des Ausgabetyps mit broadcast_dimension[i] zugeordnet wird. Die Dimensionen von operand müssen die Größe 1 haben oder der Größe der Dimension in der Ausgabeform entsprechen, der sie zugeordnet sind. Die verbleibenden Dimensionen werden mit Dimensionen der Größe 1 ausgefüllt. Bei der Übertragung mit degenerierten Dimensionen werden dann diese degenerierten Dimensionen verwendet, um die Ausgabeform zu erreichen. Die Semantik wird auf der Seite „Broadcasting“ ausführlich beschrieben.

Anruf

Weitere Informationen finden Sie unter XlaBuilder::Call.

Ruft eine Berechnung mit den angegebenen Argumenten auf.

Call(computation, args...)

Die Arithmetik und die Typen der args müssen mit den Parametern der computation übereinstimmen. Es ist zulässig, kein args zu haben.

CompositeCall

Weitere Informationen finden Sie unter XlaBuilder::CompositeCall.

Kapselt einen Vorgang ein, der aus anderen StableHLO-Vorgängen besteht, Eingaben und zusammengesetzte Attribute annimmt und Ergebnisse liefert. Die Semantik der Operation wird durch das Dekompositionsattribut implementiert. Die zusammengesetzte Operation kann durch ihre Dekomposition ersetzt werden, ohne die Programmsemantik zu ändern. Wenn die Einbettung der Dekomposition nicht dieselbe Op-Semantik bietet, sollten Sie stattdessen „custom_call“ verwenden.

Im Feld „Version“ (Standardwert: 0) wird angegeben, wann sich die Semantik eines Composites ändert.

Diese Operation wird als kCall mit dem Attribut is_composite=true implementiert. Das Feld decomposition wird durch das Attribut computation angegeben. Die restlichen Attribute werden in den Frontend-Attributen gespeichert, die mit composite. beginnen.

Beispiel für eine CompositeCall-Operation:

f32[] call(f32[] %cst), to_apply=%computation, is_composite=true,
frontend_attributes = {
  composite.name="foo.bar",
  composite.attributes={n = 1 : i32, tensor = dense<1> : tensor<i32>},
  composite.version="1"
}

Call(computation, args..., name, composite_attributes, version)

Cholesky

Weitere Informationen finden Sie unter XlaBuilder::Cholesky.

Berechnet die Cholesky-Zerlegung einer Gruppe symmetrischer (hermitescher) positiv definiter Matrizen.

Cholesky(a, lower)

Wenn lower = true ist, werden untere Dreiecksmatrizen l berechnet, sodass $a = l ist . l^T$. Wenn lower = false ist, werden obere Dreiecksmatrizen u berechnet, sodass$a = u^T . u$.

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

Wenn a mehr als zwei Dimensionen hat, wird a als Batch von Matrizen behandelt, wobei alle Dimensionen mit Ausnahme der zwei untergeordneten Dimensionen Batchdimensionen sind.

Wenn a nicht symmetrisch (Hermitian) positiv definit ist, ist das Ergebnis implementierungsabhängig.

Einschränken

Weitere Informationen finden Sie unter XlaBuilder::Clamp.

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

Clamp(min, operand, max)

Gibt bei einem Operanden und Mindest- und Höchstwerten den Operanden zurück, wenn er sich im Bereich zwischen Mindest- und Höchstwert befindet. Andernfalls wird der Mindestwert zurückgegeben, wenn der Operand unter diesem Bereich liegt, oder der 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 dieselbe Form haben. Alternativ können min und/oder max als eingeschränkte Form der Broadcasting-Funktion als Skalar vom Typ T verwendet werden.

Beispiel mit den Skalaren min und max:

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

Minimieren

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

Die Dimensionen eines Arrays werden in eine Dimension zusammengefasst.

Collapse(operand, dimensions)

Durch Zusammenführen wird die angegebene Teilmenge der Dimensionen des Operanden durch eine einzelne Dimension ersetzt. Die Eingabeargumente sind ein beliebiges Array vom Typ T und ein Compilezeit-konstanter Vektor von Dimensionsindizes. Die Dimensionsindizes müssen in der richtigen Reihenfolge (niedrigste bis höchste Dimensionsnummer) eine fortlaufende Teilmenge der Dimensionen von T sein. Daher sind {0, 1, 2}, {0, 1} oder {1, 2} gültige Dimensionsgruppen, {1, 0} oder {0, 2} jedoch nicht. Sie werden durch eine einzelne neue Dimension ersetzt, die an derselben Position in der Dimensionssequenz wie die ersetzten Dimensionen steht. Die Größe der neuen Dimension entspricht dem Produkt der ursprünglichen Dimensionsgrößen. Die niedrigste Dimensionsnummer in dimensions ist die Dimension mit der langsamsten Änderung (höchste Hauptdimension) im Schleifennest, in dem diese Dimensionen minimiert werden. Die höchste Dimensionsnummer ist die Dimension mit der schnellsten Änderung (niedrigste untergeordnete Dimension). Weitere Informationen finden Sie unter tf.reshape, wenn eine allgemeinere Zusammenklappungsreihenfolge erforderlich ist.

Angenommen, v ist ein Array mit 24 Elementen:

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

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

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

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

CollectivePermute

Weitere Informationen finden Sie unter XlaBuilder::CollectivePermute.

CollectivePermute ist ein kollektiver Vorgang, bei dem Daten zwischen Replikas gesendet und empfangen werden.

CollectivePermute(operand, source_target_pairs)

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

• Keines der Paare darf dieselbe Zielreplikat-ID und dieselbe Quellreplikat-ID haben.
• Wenn eine Replikat-ID in keinem Paar als Ziel festgelegt ist, ist die Ausgabe für dieses Replikat ein Tensor mit Nullen in derselben Form wie die Eingabe.

Verketten

Weitere Informationen finden Sie unter XlaBuilder::ConcatInDim.

Mit „Verketten“ wird ein Array aus mehreren Arrayoperanden gebildet. Das Array hat dieselbe Anzahl von Dimensionen wie jeder der Eingabearrayoperanden (die dieselbe Anzahl von Dimensionen haben müssen) und enthält die Argumente in der Reihenfolge, in der sie angegeben wurden.

Concatenate(operands..., dimension)

Mit Ausnahme von dimension müssen alle Dimensionen gleich sein. Das liegt daran, dass XLA keine unvollständigen Arrays unterstützt. Außerdem können nulldimensionale Werte nicht zusammengeführt werden, da die Dimension, entlang derer die Zusammenführung erfolgt, nicht benannt werden kann.

Eindimensionales Beispiel:

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

2-dimensionales Beispiel:

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)

Führt true_computation aus, wenn pred gleich true ist, und false_computation, wenn pred gleich false ist. Das Ergebnis wird zurückgegeben.

true_computation muss ein einzelnes Argument vom Typ $T_0$ annehmen und wird mit true_operand aufgerufen, das vom selben Typ sein muss. false_computation muss ein einzelnes Argument vom Typ $T_1$ annehmen und wird mit false_operand aufgerufen, das vom selben Typ sein muss. Der Typ des zurückgegebenen Werts von true_computation und false_computation muss gleich sein.

Je nach Wert von pred wird nur eine von true_computation und false_computation ausgeführt.

Conditional(branch_index, branch_computations, branch_operands)

Führt branch_computations[branch_index] aus und gibt das Ergebnis zurück. Wenn branch_index eine S32 ist, die kleiner als 0 oder größer oder gleich 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, das vom selben Typ sein muss. Der Typ des zurückgegebenen Werts von branch_computations[b] muss bei allen Elementen gleich sein.

Je nach Wert von branch_index wird nur eine der branch_computations-Aktionen ausgeführt.

Conv (Faltung)

Weitere Informationen finden Sie unter XlaBuilder::Conv.

Wie ConvWithGeneralPadding, aber die Füllung wird kurz als „SAME“ oder „VALID“ angegeben. Bei der SAME-Ausrichtung wird die Eingabe (lhs) mit Nullen aufgefüllt, damit die Ausgabe dieselbe Form wie die Eingabe hat, wenn der Schritt nicht berücksichtigt wird. VALID padding bedeutet einfach, dass kein Padding verwendet wird.

ConvWithGeneralPadding (Faltung)

Weitere Informationen finden Sie unter XlaBuilder::ConvWithGeneralPadding.

Berechnet eine Convolution, wie sie in neuronalen Netzwerken verwendet wird. Hier kann eine Convolution als n-dimensionales Fenster betrachtet werden, das sich über einen n-dimensionalen Basisbereich bewegt. Für jede mögliche Position des Fensters wird eine Berechnung durchgeführt.

n ist die Anzahl der räumlichen Dimensionen. Das Argument lhs ist ein (n+2)-dimensionales Array, das die Basisfläche beschreibt. Dies wird als Eingabe bezeichnet, auch wenn natürlich auch der RHS eine Eingabe ist. In einem neuronalen Netzwerk sind dies 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 die Convolution durchgeführt wird.
• z/depth/features: Jede (y,x)-Position im Basisbereich ist mit einem Vektor verknüpft, der in diese Dimension eingeht.
• spatial_dims: Beschreibt die n räumlichen Dimensionen, die den Basisbereich definieren, über den sich das Fenster bewegt.

Das rhs-Argument ist ein (n+2)-dimensionales Array, das den Convolutional-Filter/-Kern/-Fenster beschreibt. Die Dimensionen sind in dieser Reihenfolge:

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

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

Das Argument padding gibt an, wie viel Null-Padding auf den Basisbereich angewendet werden soll. Die Größe des Paddings kann negativ sein. Der absolute Wert des negativen Paddings gibt an, wie viele Elemente aus der angegebenen Dimension vor der Durchführung der Convolution entfernt werden sollen. padding[0] gibt den Abstand für die Dimension y und padding[1] den Abstand für die Dimension x an. Jedes Paar hat das geringe Abstandselement als erstes Element und das große Abstandselement als zweites Element. Die geringe Lücke wird in Richtung niedrigerer Indizes angewendet, während die große Lücke in Richtung höherer Indizes angewendet wird. Wenn padding[1] beispielsweise (2,3) ist, wird in der zweiten räumlichen Dimension links mit zwei Nullen und rechts mit drei Nullen eingefügt. Das Hinzufügen von Padding entspricht dem Einfügen derselben Nullwerte in die Eingabe (lhs) vor der Durchführung der Convolution.

Mit den Argumenten lhs_dilation und rhs_dilation wird der Vergrößerungsfaktor angegeben, der in jeder räumlichen Dimension auf den linken bzw. rechten Ausdruck angewendet werden soll. Wenn der Erweiterungsfaktor in einer räumlichen Dimension d ist, werden zwischen jedem der Einträge in dieser Dimension implizit d-1 Lücken platziert, wodurch die Größe des Arrays erhöht wird. Die Löcher werden mit einem Nullwert gefüllt, was bei der Convolution Nullen bedeutet.

Die Dilatation der rechten Seite wird auch als atrosische Convolution bezeichnet. Weitere Informationen finden Sie unter tf.nn.atrous_conv2d. Die Dilatation des linken Glieds wird auch als transponierte Konvolution bezeichnet. Weitere Informationen finden Sie unter tf.nn.conv2d_transpose.

Das Argument feature_group_count (Standardwert 1) kann für gruppierte Convolutionen verwendet werden. feature_group_count muss sowohl für die Eingabe- als auch für die Ausgabe-Featuredimension ein Teiler sein. Wenn feature_group_count größer als 1 ist, bedeutet das, 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 fortlaufenden Folge von Features besteht. Die Eingabemerkmalsdimension rhs muss der Eingabemerkmalsdimension lhs geteilt durch feature_group_count entsprechen. Sie hat also bereits die Größe einer Gruppe von Eingabemerkmalen. Die i-ten Gruppen werden zusammen verwendet, um feature_group_count für viele separate Convolutionen zu berechnen. Die Ergebnisse dieser Faltungen werden in der Ausgabefeature-Dimension zusammengefügt.

Bei der richtungsabhängigen Convolution 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] umgewandelt. Weitere Informationen finden Sie unter tf.nn.depthwise_conv2d.

Das Argument batch_group_count (Standardwert 1) kann für gruppierte Filter während der Backpropagation verwendet werden. batch_group_count muss ein Teiler der Größe der Batchdimension lhs (Eingabe) sein. Wenn batch_group_count größer als 1 ist, sollte die Ausgabe-Batchdimension die Größe input batch / batch_group_count haben. Die batch_group_count muss ein Teiler der Größe der Ausgabeelemente sein.

Die Ausgabeform hat in dieser Reihenfolge die folgenden Dimensionen:

• batch: Die Größe dieser Dimension multipliziert mit batch_group_count muss der Größe der Dimension batch in lhs entsprechen.
• z: Entspricht der Größe von output-z im Kernel (rhs).
• spatial_dims: Ein Wert für jede gültige Platzierung des Convolutional-Fensters.

Die Abbildung oben zeigt, wie das Feld batch_group_count funktioniert. Wir teilen jeden lhs-Batch in batch_group_count Gruppen auf und machen dasselbe mit den Ausgabefeatures. Anschließend führen wir für jede dieser Gruppen paarweise Convolutionen durch und fügen die Ausgabe entlang der Ausgabemerkmalsdimension zusammen. Die betriebliche Semantik aller anderen Dimensionen (Feature- und standortbezogene) bleibt unverändert.

Die gültigen Platzierungen des konvolutionalen Fensters werden durch die Schritte und die Größe der Grundfläche nach dem Padding bestimmt.

Um die Funktionsweise einer 2D-Konvolution zu beschreiben, betrachten Sie eine 2D-Konvolution und wählen Sie einige feste Koordinaten batch, z, y, x in der Ausgabe aus. Dann ist (y,x) die Position einer Ecke des Fensters im Basisbereich (z.B. die linke obere Ecke, je nachdem, wie Sie die räumlichen Dimensionen interpretieren). Wir haben jetzt ein 2D-Fenster aus dem Basisbereich, bei dem jeder 2D-Punkt mit einem 1D-Vektor verknüpft ist. So erhalten wir einen 3D-Quader. Da wir die Ausgabekoordinate z für den Convolutional-Kernel festgelegt haben, haben wir auch einen 3D-Box. Die beiden Boxen haben dieselben Abmessungen, sodass wir die Summe der elementweisen Produkte zwischen den beiden Boxen ermitteln können (ähnlich wie bei einem Skalarprodukt). Das ist der Ausgabewert.

Wenn output-z beispielsweise 5, dann werden für jede Position des Fensters fünf Werte in der Ausgabe in der Dimension z generiert. Diese Werte unterscheiden sich darin, welcher Teil des Convolutional-Kernels verwendet wird. Für jede output-z-Koordinate gibt es einen separaten 3D-Box mit Werten. Sie können sich das also als fünf separate Convolutionen mit jeweils einem anderen Filter vorstellen.

Hier ist der Pseudocode für eine 2D-Konvolution mit Padding und Schrittweite:

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.

Ähnlich wie eine elementweise static_cast in C++ führt dieser Befehl eine elementweise Konvertierung von einer Datenform in eine Zielform aus. Die Dimensionen müssen übereinstimmen und die Conversion muss elementweise erfolgen. Beispiel: s32-Elemente werden über eine s32-zu-f32-Conversion-Routine in f32-Elemente umgewandelt.

ConvertElementType(operand, new_element_type)

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

Bei einer Konvertierung wie T=s32 in U=f32 wird eine normalisierende Konvertierungsroutine vom Typ „Int in Float“ ausgeführt, z. B. „auf die nächste gerade Zahl aufrunden“.

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

CrossReplicaSum

Führt AllReduce mit einer Summenberechnung aus.

CustomCall

Weitere Informationen finden Sie unter XlaBuilder::CustomCall.

Eine vom Nutzer bereitgestellte Funktion innerhalb einer Berechnung aufrufen.

CustomCall(target_name, args..., shape)

Die Funktionssignatur ist unabhängig von der Arity oder dem Typ von args identisch:

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

Beispiel:

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

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

Hier 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 Nebenwirkungen haben und ihre Ausführung muss idempotent sein.

Punkt

Weitere Informationen finden Sie unter XlaBuilder::Dot.

Dot(lhs, rhs)

Die genaue Bedeutung dieses Vorgangs hängt von den Rängen der Operanden ab:

Bei diesem Vorgang wird die Summe der Produkte über die zweite Dimension von lhs (oder die erste, wenn es nur eine Dimension hat) und die erste Dimension von rhs berechnet. Das sind die „zusammengeführten“ Dimensionen. Die minimierten Dimensionen von lhs und rhs müssen dieselbe Größe haben. In der Praxis kann es zum Ausführen von Skalarprodukten zwischen Vektoren, Vektor-/Matrixmultiplikationen oder Matrix-/Matrixmultiplikationen verwendet werden.

DotGeneral

Weitere Informationen finden Sie unter XlaBuilder::DotGeneral.

DotGeneral(lhs, rhs, dimension_numbers)

Ähnlich wie „Punkt“, aber es können sowohl für lhs als auch für rhs Zahlen für die Dimensionskontraktion und die Batchdimension angegeben werden.

Mit „DotGeneral“ wird die Summe der Produkte über die in dimension_numbers angegebenen Dimensionsgruppen berechnet.

Die zugehörigen Nummern der Vertragsdimensionen aus lhs und rhs müssen nicht identisch sein, aber dieselben Dimensionsgrößen haben.

Beispiel für schrumpfende 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} }

Die zugehörigen Nummern der Batchdimensionen aus lhs und rhs müssen dieselbe Dimensionsgröße haben.

Beispiel mit Zahlen für die Batchdimension (Batchgröße 2, Matrizen 2 × 2):

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

Die resultierende Dimensionsnummer beginnt also mit der Batch-Dimension, dann mit der lhs-Dimension ohne Schrumpfung/Batch und schließlich mit der rhs-Dimension ohne Schrumpfung/Batch.

DynamicSlice

Weitere Informationen finden Sie unter XlaBuilder::DynamicSlice.

Mit DynamicSlice wird ein Teilarray aus dem Eingabearray an der dynamischen Position start_indices extrahiert. Die Größe des Ausschnitts in jeder Dimension wird in size_indices übergeben, wobei der Endpunkt der exklusiven Ausschnittsintervalle in jeder Dimension angegeben wird: [start, start + size]. Die Form von start_indices muss eindimensional sein und die Dimensionsgröße muss der Anzahl der Dimensionen von operand entsprechen.

DynamicSlice(operand, start_indices, size_indices)

Die effektiven Intervallindices werden berechnet, indem vor dem Ausführen des Intervalls die folgende Transformation auf jeden Index i in [1, N) angewendet wird:

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

So wird sichergestellt, dass der extrahierte Ausschnitt immer innerhalb des Operandenarrays liegt. Wenn sich der Ausschnitt vor der Anwendung der Transformation innerhalb des Bereichs befindet, hat die Transformation keine Auswirkungen.

Eindimensionales Beispiel:

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

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

2-dimensionales Beispiel:

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 ein Ausschnitt update an start_indices überschrieben wird. Die Form von update bestimmt die Form des Teilarrays des Ergebnisses, das aktualisiert wird. Die Form von start_indices muss eindimensional sein und die Dimensionsgröße muss der Anzahl der Dimensionen von operand entsprechen.

DynamicUpdateSlice(operand, update, start_indices)

Die effektiven Intervallindices werden berechnet, indem vor dem Ausführen des Intervalls die folgende Transformation auf jeden Index i in [1, N) angewendet wird:

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

So wird sichergestellt, dass der aktualisierte Ausschnitt immer innerhalb des Operandenarrays liegt. Wenn sich der Ausschnitt vor der Anwendung der Transformation innerhalb des Bereichs befindet, hat die Transformation keine Auswirkungen.

Eindimensionales Beispiel:

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}

2-dimensionales Beispiel:

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 werden eine Reihe von elementweisen binären arithmetischen Operationen unterstützt.

Op(lhs, rhs)

Dabei steht Op für eine der folgenden Funktionen: Add (Addition), Sub(Subtraktion), Mul (Multiplikation), Div (Division), Pow (Potenz), Rem (Rest), Max (Maximum), Min (Minimum), And (logisches AND), Or (logisches OR), Xor (logisches XOR), ShiftLeft (Linksverschiebung), ShiftRightArithmetic (Arithmetische Rechtsverschiebung), ShiftRightLogical (logische Rechtsverschiebung), Atan2 (Arkustangens mit zwei Argumenten) oder Complex (Kombiniert den reellen und den imaginären Teil zu einer komplexen Zahl)

Die Formen der Argumente müssen entweder ähnlich oder kompatibel sein. In der Dokumentation zu Übertragungen erfahren Sie, was es bedeutet, dass Formen kompatibel sind. Das Ergebnis einer Operation hat die Form, die durch das Broadcasten der beiden Eingabearrays entsteht. Bei dieser Variante werden Operationen zwischen Arrays verschiedener Ränge nicht unterstützt, es sei denn, einer der Operanden ist ein Skalar.

Wenn Op = Rem ist, wird das Vorzeichen des Ergebnisses aus dem Dividenden übernommen und der absolute Wert des Ergebnisses ist immer kleiner als der absolute Wert des Nenners.

Bei einem Überlauf bei der Ganzzahldivision (Division/Rest mit Vorzeichen/ohne Vorzeichen durch Null oder Division/Rest mit Vorzeichen von INT_SMIN mit -1) wird ein implementierungsdefinierter Wert ausgegeben.

Für diese Vorgänge gibt es eine alternative Variante mit Unterstützung für die Übertragung verschiedener Dimensionen:

Op(lhs, rhs, broadcast_dimensions)

Dabei ist Op wie oben angegeben. Diese Variante des Vorgangs sollte für arithmetische Operationen zwischen Arrays unterschiedlicher Ränge verwendet werden, z. B. wenn einem Vektor eine Matrix hinzugefügt wird.

Der zusätzliche broadcast_dimensions-Operand ist ein Array von Ganzzahlen, mit dem die Anzahl der Dimensionen des Operanden mit niedrigerer Dimension auf die Anzahl der Dimensionen des Operanden mit höherer Dimension erweitert wird. broadcast_dimensionsordnet die Dimensionen der niedrigerdimensionalen Form den Dimensionen der höherdimensionalen Form zu. Die nicht zugeordneten Dimensionen der erweiterten Form werden mit Dimensionen der Größe 1 ausgefüllt. Bei der Übertragung mit degenerierten Dimensionen werden die Formen dann entlang dieser degenerierten Dimensionen übertragen, um die Formen beider Operanden anzugleichen. Die Semantik wird auf der Seite „Übertragungen“ ausführlich beschrieben.

Elementweise Vergleichsvorgänge

Weitere Informationen finden Sie unter XlaBuilder::Eq.

Es werden eine Reihe standardmäßiger elementweiser binärer Vergleichsvorgänge unterstützt. Beim Vergleich von Gleitkommatypen gilt die Standard-IEEE-754-Gleitkommavergleichssemantik.

Op(lhs, rhs)

Dabei steht Op für Eq (gleich), Ne (ungleich), Ge (größer oder gleich), Gt (größer), Le (kleiner oder gleich) oder Lt (kleiner). Die Operatoren „EqTotalOrder“, „NeTotalOrder“, „GeTotalOrder“, „GtTotalOrder“, „LeTotalOrder“ und „LtTotalOrder“ bieten dieselben Funktionen, unterstützen aber zusätzlich eine totale Ordnung über die Gleitkommazahlen, indem sie -NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN erzwingen.

Die Formen der Argumente müssen entweder ähnlich oder kompatibel sein. In der Dokumentation zu Übertragungen erfahren Sie, was es bedeutet, dass Formen kompatibel sind. Das Ergebnis einer Operation hat die Form, die durch das Broadcasten der beiden Eingabearrays mit dem Elementtyp PRED entsteht. Bei dieser Variante werden keine Operationen zwischen Arrays unterschiedlicher Ränge unterstützt, es sei denn, einer der Operanden ist ein Skalar.

Für diese Vorgänge gibt es eine alternative Variante mit Unterstützung für die Übertragung verschiedener Dimensionen:

Op(lhs, rhs, broadcast_dimensions)

Dabei ist Op wie oben angegeben. Diese Variante des Vorgangs sollte für Vergleichsvorgänge zwischen Arrays unterschiedlicher Ränge verwendet werden, z. B. wenn einem Vektor eine Matrix hinzugefügt wird.

Der zusätzliche Operand broadcast_dimensions ist ein Array von Ganzzahlen, das die Dimensionen angibt, die für die Übertragung der Operanden verwendet werden sollen. Die Semantik wird auf der Seite „Broadcasting“ ausführlich beschrieben.

Elementweise unary functions

XlaBuilder unterstützt die folgenden elementweisen unaryn Funktionen:

Abs(operand) Elementweiser Absolutwert x -> |x|.

Cbrt(operand) Elementweise Kubikwurzel x -> cbrt(x)

Ceil(operand) Elementweiser Wert x -> ⌈x⌉.

Clz(operand) Zählt führende Nullen elementweise.

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

Erf(operand) Elementweise Fehlerfunktion x -> erf(x) mit

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

Exp(operand) Elementweise natürliche Exponentialfunktion x -> e^x.

Expm1(operand) Elementweiser natürlicher Exponent minus eins x -> e^x - 1.

Floor(operand) Elementspezifischer Mindestpreis x -> ⌊x⌋.

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

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

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

Log1p(operand) Elementweise verschobener natürlicher Logarithmus x -> ln(1+x).

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

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

Not(operand) Elementweises logisches Nicht x -> !(x).

PopulationCount(operand) Berechnet die Anzahl der Bits, die in jedem Element von operand festgelegt sind.

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

Round(operand) Elementweise Rundung, bei Unentschieden von null weg.

RoundNearestEven(operand) Elementweises Runden auf die nächste gerade Zahl.

Rsqrt(operand) Elementweises Kehrwert der Quadratwurzelx -> 1.0 / sqrt(x).

Sign(operand) Elementweiser Signaturvorgang x -> sgn(x), wobei

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

mit dem Vergleichsoperator des Elementtyps operand.

Sin(operand) Elementweise Sinusfunktion x -> sin(x).

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

Tan(operand) Elementweise Tangente x -> tan(x).

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

Die Funktion wird auf jedes Element im operand-Array angewendet, was zu einem Array mit derselben Form führt. operand darf ein Skalar (0-dimensional) sein.

Fft

Der XLA-FFT-Vorgang implementiert die Vorwärts- und Rückwärts-Fourier-Transformation für reelle und komplexe Eingaben/Ausgaben. Mehrdimensionale FFTs auf bis zu 3 Achsen werden unterstützt.

Weitere Informationen finden Sie unter XlaBuilder::Fft.

Multidimensionale FFT

Wenn mehr als eine fft_length angegeben wird, entspricht dies der Anwendung einer Kaskade von FFT-Vorgängen auf jede der innersten Achsen. Beachten Sie, dass bei den Fällen „reell → komplex“ und „komplex → reell“ die Transformation der innersten Achse (RFFT; bei IRFFT als letztes) (effektiv) zuerst ausgeführt wird. Deshalb ändert sich die Größe der innersten Achse. Andere Achsentransformationen sind dann „komplex“ -> „komplex“.

Implementierungsdetails

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

Zusammentragen

Der XLA-Gattervorgang fügt mehrere Segmente (jedes Segment mit einem potenziell unterschiedlichen Laufzeitoffset) eines Eingabearrays zusammen.

Allgemeine Semantik

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

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

Der Einfachheit halber bezeichnen wir Dimensionen im Ausgabearray, die nicht in offset_dims enthalten sind, als batch_dims.

Die Ausgabe ist ein Array mit batch_dims.size + offset_dims.size Dimensionen.

operand.rank muss der Summe von offset_dims.size und collapsed_slice_dims.size entsprechen. Außerdem muss slice_sizes.size mit operand.rank übereinstimmen.

Wenn index_vector_dim = start_indices.rank ist, wird davon ausgegangen, dass start_indices eine nachgestellte 1-Dimension hat. Wenn start_indices also die Form [6,7] hat und index_vector_dim = 2 ist, wird davon ausgegangen, dass start_indices die Form [6,7,1] hat.

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

1. Wenn i in batch_dims enthalten ist (d. h. für eine bestimmte k mit batch_dims[k] übereinstimmt), werden die entsprechenden Dimensionsgrenzen aus start_indices.shape ausgewählt, wobei index_vector_dim übersprungen wird. Das bedeutet, dass start_indices.shape.dims[k] ausgewählt wird, wenn k < index_vector_dim, andernfalls start_indices.shape.dims[k+1].

2. Wenn i in offset_dims vorhanden ist (d. h. offset_dims[k] für eine bestimmte k entspricht), wählen wir die entsprechende Grenze aus slice_sizes aus, nachdem collapsed_slice_dims berücksichtigt wurde. Das heißt, wir wählen adjusted_slice_sizes[k] aus, wobei adjusted_slice_sizes slice_sizes ist, ohne die Grenzen an den Indizes collapsed_slice_dims.

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

1. Angenommen, G = {Out[k] für k in batch_dims}. Verwenden Sie G, um einen Vektor S zu schneiden, sodass S[i] = start_indices[Combine(G, i)] ist, wobei Combine(A, b) b an Position index_vector_dim in A einfügt. Dies ist auch dann eindeutig definiert, wenn G leer ist: Wenn G leer ist, ist S = start_indices.

2. Erstellen Sie mit S einen Startindex, Sin, in operand. Dazu müssen Sie S mit start_index_map streuen. Genauer gesagt:

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

   2. Sin[_] = 0 andernfalls.

3. Erstellen Sie einen Index Oin in operand, indem Sie die Indizes gemäß dem collapsed_slice_dims-Set auf die Offset-Dimensionen in Out 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. Oin[_] = 0 andernfalls.

4. In ist Oin + Sin, wobei „+“ die elementweise Addition ist.

remapped_offset_dims ist eine monoton steigende Funktion mit Definitionsbereich [0, offset_dims.size) und Wertebereich [0, operand.rank) \ collapsed_slice_dims. Beispiel: Wenn offset_dims.size = 4, operand.rank = 6 und collapsed_slice_dims = {0, 2}, dann ist remapped_offset_dims = {0→1, 1→3, 2→4, 3→5}.

Wenn indices_are_sorted auf „wahr“ gesetzt ist, kann XLA davon ausgehen, dass start_indices vom Nutzer sortiert wurde (in aufsteigender Reihenfolge, nachdem die Werte gemäß start_index_map verstreut wurden). Andernfalls werden die Semantiken von der Implementierung definiert.

Beschreibung und Beispiele

Ungezwungen ausgedrückt entspricht jeder Index Out im Ausgabearray einem Element E im Operandenarray, das so berechnet wird:

• Anhand der Batch-Dimensionen in Out wird ein Startindex aus start_indices abgerufen.

• Mit start_index_map wird der Startindex (dessen Größe kleiner als operand.rank sein kann) einem „vollständigen“ Startindex in operand zugeordnet.

• Wir schneiden dynamisch einen Ausschnitt mit der Größe slice_sizes mit dem vollständigen Startindex aus.

• Wir ändern die Form des Ausschnitts, indem wir die collapsed_slice_dims-Dimensionen minimieren. Da alle minimierten Dimensionsebenen eine Obergrenze von 1 haben müssen, ist diese Neuformatierung immer zulässig.

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

In allen folgenden Beispielen ist index_vector_dim auf start_indices.rank bis 1 festgelegt. Interessantere Werte für index_vector_dim ändern die Funktionsweise nicht grundlegend, erschweren aber die visuelle Darstellung.

Um zu verstehen, wie das alles zusammenhängt, sehen wir uns ein Beispiel an, bei dem fünf Scheiben der Form [8,6] aus einem [16,11]-Array erfasst werden. Die Position eines Slabs im [16,11]-Array kann als Indexvektor der Form S64[2] dargestellt werden. Die fünf Positionen können also als S64[5,2]-Array dargestellt werden.

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

Zuerst wählen wir mit G einen (X,Y)-Vektor aus dem Index-Array für die Gruppierung aus. Das Element im Ausgabearray mit dem Index [G,O0,O1] ist dann das Element im Eingabearray mit dem Index [X+O0,Y+O1].

slice_sizes ist [8,6], was den Bereich von O₀ und O₁ bestimmt, was wiederum die Grenzen des Slices bestimmt.

Dieser Vorgang dient als dynamischer Batch-Ausschnitt mit G als Batch-Dimension.

Die Aggregationsindizes können mehrdimensional sein. In einer allgemeineren Version des Beispiels oben mit einem „Indexe erfassen“-Array der Form [4,5,2] würden Indizes beispielsweise so übersetzt:

Auch hier handelt es sich um einen dynamischen Batch-Ausschnitt G0 mit den Batch-Dimensionen G1. Die Größe der Scheibe bleibt [8,6].

Der Vorgang „GATHER“ in XLA verallgemeinert die oben beschriebenen informellen Semantiken auf folgende Weise:

1. Wir können konfigurieren, welche Dimensionen im Ausgabeformat die Offset-Dimensionen sind (Dimensionen mit O0, O1 im letzten Beispiel). Die Ausgabe-Batchdimensionen (Dimensionen mit G0, G1 im letzten Beispiel) sind die Ausgabedimensionen, die keine Offset-Dimensionen sind.

2. Die Anzahl der Ausgabeoffset-Dimensionen, die explizit in der Ausgabeform vorhanden sind, kann kleiner als die Anzahl der Eingabedimensionen sein. Diese „fehlenden“ Dimensionen, die ausdrücklich als collapsed_slice_dims aufgeführt sind, müssen eine Segmentgröße von 1 haben. Da sie eine Segmentgröße von 1 haben, ist 0 der einzige gültige Index für sie. Wenn sie entfernt werden, führt das nicht zu Unklarheiten.

3. Der Ausschnitt, der aus dem Array „Indices zusammenführen“ extrahiert wird (X, Y im letzten Beispiel), kann weniger Elemente als die Anzahl der Dimensionen des Eingabearrays haben. Eine explizite Zuordnung legt fest, wie der Index erweitert werden soll, damit er dieselbe Anzahl von Dimensionen wie die Eingabe hat.

Als letztes Beispiel verwenden wir (2) und (3), um tf.gather_nd zu implementieren:

Mit G0 und G1 wird wie gewohnt ein Startindex aus dem Array mit den Indexen für die Datenerhebung extrahiert. Der Unterschied besteht darin, dass der Startindex nur ein Element hat: X. Ebenso gibt es nur einen Ausgabeoffsetindex mit dem Wert O0. Bevor sie jedoch als Indizes in das Eingabearray übernommen werden, werden sie gemäß „Indexzuordnung für die Datenerhebung“ (start_index_map in der formellen Beschreibung) und „Offsetzuordnung“ (remapped_offset_dims in der formellen Beschreibung) in [X,0] und [0,O0] erweitert. Die Summe dieser beiden Indizes ergibt [X,O0]. Mit anderen Worten: Der Ausgabeindex [G0,G1,O0] wird dem Eingabeindex [GatherIndices[G0,G1,0],O0] zugeordnet, was die Semantik für tf.gather_nd ergibt.

slice_sizes für diesen Fall ist [1,11]. Intuitiv bedeutet das, dass jeder Index X im Array „Indexe für Zusammenführung“ eine ganze Zeile auswählt und das Ergebnis die Koncatenate aller 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 ein Array sein.

GetDimensionSize(operand, dimension)

SetDimensionSize

Weitere Informationen finden Sie unter XlaBuilder::SetDimensionSize.

Legt die dynamische Größe der angegebenen Dimension von XlaOp fest. Der Operand muss ein Array sein.

SetDimensionSize(operand, size, dimension)

Der Operand wird als Ergebnis übergeben, wobei die dynamische Dimension vom Compiler erfasst wird.

Aufgefüllte Werte werden von nachfolgenden Reduzierungsaktionen 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 einem Tupel mit einem konstanten Wert zur Kompilierungszeit.

Der Wert muss eine Konstante zur Kompilierungszeit sein, damit die Forminferenz den Typ des resultierenden Werts bestimmen kann.

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

In-Feed

Weitere Informationen finden Sie unter XlaBuilder::Infeed.

Infeed(shape)

Liest ein einzelnes Datenelement aus der impliziten Infeed-Streaming-Schnittstelle des Geräts, interpretiert die Daten als die angegebene Form und das Layout und gibt eine XlaOp der Daten zurück. In einer Berechnung sind mehrere In-Feed-Vorgänge zulässig. Es muss jedoch eine Gesamtreihenfolge für die In-Feed-Vorgänge geben. Im folgenden Code haben beispielsweise zwei Infeeds eine Gesamtreihenfolge, da es eine Abhängigkeit zwischen den While-Schleifen gibt.

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 im Grunde ein Null-Vorgang und wird ausgeführt, ohne dass Daten aus dem Infeed des Geräts gelesen werden.

Iota

Weitere Informationen finden Sie unter XlaBuilder::Iota.

Iota(shape, iota_dimension)

Es wird ein konstantes Literal auf dem Gerät erstellt, anstatt eine potenziell große Hostübertragung. Erstellt ein Array mit einer bestimmten Form, das Werte enthält, die mit null beginnen und sich entlang der angegebenen Dimension um eins erhöhen. Bei Gleitkommatypen entspricht das generierte Array ConvertElementType(Iota(...)), wobei Iota einen Ganzzahltyp hat und die Umwandlung in den Gleitkommatyp erfolgt.

Beispiel: Iota(s32[4, 8], 0) gibt

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

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

Die zugeordnete Funktion ist eine beliebige Berechnung mit der Einschränkung, dass sie N Eingänge vom Skalartyp T und eine einzelne Ausgabe vom Typ S hat. Die Ausgabe hat dieselben Dimensionen wie die Operanden, mit der Ausnahme, dass der Elementtyp T durch S ersetzt wird.

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

OptimizationBarrier

Verhindert, dass Optimierungsdurchläufe Berechnungen über die Barriere hinweg verschieben.

Damit wird sichergestellt, dass alle Eingaben vor allen Operatoren ausgewertet werden, die von den Ausgaben der Barriere abhängen.

Pad

Weitere Informationen finden Sie unter XlaBuilder::Pad.

Pad(operand, padding_value, padding_config)

Erweitert das angegebene operand-Array, indem um das Array herum und zwischen den Elementen des Arrays mit dem angegebenen padding_value ein Abstand hinzugefügt wird. padding_config gibt die Größe des Randabstands und des Abstands im Inneren für jede Dimension an.

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

Mit edge_padding_low und edge_padding_high wird die Größe des Abstands angegeben, der am unteren Ende (neben Index 0) bzw. am oberen Ende (neben dem höchsten Index) jeder Dimension hinzugefügt wird. Die Größe des Randabstands kann negativ sein. Der absolute Wert des negativen Abstands gibt die Anzahl der Elemente an, die aus der angegebenen Dimension entfernt werden sollen.

interior_padding gibt den Abstand an, der zwischen zwei Elementen in jeder Dimension hinzugefügt wird. Er darf nicht negativ sein. Das Innenabstand wird logisch vor dem Randabstand angewendet. Bei einem negativen Randabstand werden Elemente aus dem Operanden mit Innenabstand entfernt.

Dieser Vorgang ist wirkungslos, wenn alle Ränder (0, 0) und alle Werte für den Innenabstand 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)

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

Die Client-API des Recv-Vorgangs steht für die synchrone Kommunikation. Die Anweisung wird jedoch intern in zwei HLO-Anweisungen (Recv und RecvDone) zerlegt, um asynchrone Datenübertragungen zu ermöglichen. Siehe auch HloInstruction::CreateRecv und HloInstruction::CreateRecvDone.

Recv(const Shape& shape, int64 channel_id)

Hier werden Ressourcen zugewiesen, 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 einer nachfolgenden RecvDone-Anweisung verwendet wird, um auf den Abschluss der Datenübertragung zu warten. Der Kontext ist ein Tupel von {Empfangsbuffer (Form), Anfrage-ID (U32)} und kann nur von einer RecvDone-Anweisung verwendet werden.

RecvDone(HloInstruction context)

Wartet bei einem durch eine Recv-Anweisung erstellten Kontext auf den Abschluss der Datenübertragung und gibt die empfangenen Daten zurück.

Einschränken

Weitere Informationen finden Sie unter XlaBuilder::Reduce.

Wendet eine Reduzierungsfunktion parallel auf ein oder mehrere Arrays an.

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

Wobei:

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

Bei diesem Vorgang werden eine oder mehrere Dimensionen jedes Eingabearrays in Skalare umgewandelt. Die Anzahl der Dimensionen jedes zurückgegebenen Arrays beträgt number_of_dimensions(operand) - len(dimensions). Die Ausgabe der Operation ist Collate(Q_0, ..., Q_N), wobei Q_i ein Array vom Typ T_i ist, dessen Dimensionen unten beschrieben werden.

Die Reduktionsberechnung kann mit verschiedenen Back-Ends neu verknüpft werden. Dies kann zu numerischen Abweichungen führen, da einige Reduktionsfunktionen wie die Addition für Gleitkommazahlen nicht assoziativ sind. Wenn der Bereich der Daten jedoch begrenzt ist, ist die Addition von Gleitkommazahlen für die meisten praktischen Anwendungen nahezu assoziativ.

Beispiele

Wenn Sie in einem einzelnen 1D-Array mit Werten [10, 11, 12, 13] über eine Dimension mit der Reduzierungsfunktion f (das ist computation) reduzieren, kann das so berechnet werden:

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

Es gibt aber auch viele andere Möglichkeiten, z.B.

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

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

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 number of dimensions 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 (einer Matrix): Die Form hat zwei Dimensionen: Dimension 0 mit der Größe 2 und Dimension 1 mit der Größe 3:

Ergebnisse der Reduzierung der Dimensionen 0 oder 1 mit einer Additionsfunktion:

Beachten Sie, dass beide Ergebnisse der Dimensionsreduzierung eindimensionale Arrays sind. Im Diagramm werden sie nur zu Veranschaulichung als Spalte und Zeile dargestellt.

Hier ein komplexeres Beispiel für ein 3D-Array: Die Anzahl der Dimensionen ist 3, Dimension 0 hat die Größe 4, Dimension 1 die Größe 2 und Dimension 2 die Größe 3. Der Übersichtlichkeit halber werden die Werte 1 bis 6 in Dimension 0 repliziert.

Ähnlich wie beim 2D-Beispiel können wir nur eine Dimension reduzieren. Wenn wir beispielsweise Dimension 0 reduzieren, erhalten wir ein zweidimensionales Array, in dem alle Werte in Dimension 0 in einen Skalar zusammengefasst wurden:

|  4   8  12 |
| 16  20  24 |

Wenn wir Dimension 2 reduzieren, erhalten wir auch ein zweidimensionales Array, in dem alle Werte in Dimension 2 in einen Skalar zusammengefasst wurden:

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

Die relative Reihenfolge der verbleibenden Dimensionen in der Eingabe wird in der Ausgabe beibehalten. Einigen Dimensionen werden jedoch möglicherweise neue Nummern zugewiesen, da sich die Anzahl der Dimensionen ändert.

Wir können auch mehrere Dimensionen reduzieren. Wenn Sie die Dimensionen 0 und 1 hinzufügen und reduzieren, erhalten Sie das 1‑dimensionale Array [20, 28, 36].

Wenn das 3D-Array über alle Dimensionen hinweg reduziert wird, ergibt sich der Skalar 84.

Variadic Reduce

Bei N > 1 ist die Anwendung der Funktion „Reduzieren“ etwas komplexer, da sie gleichzeitig auf alle Eingaben angewendet wird. Die Operanden werden der Berechnung in der folgenden Reihenfolge zugeführt:

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

Betrachten Sie beispielsweise die folgende Reduktionsfunktion, mit der der Maximalwert und der Argmax eines eindimensionalen 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 eindimensionalen Eingabearrays V = Float[N], K = Int[N] und Anfangswerten I_V = Float, I_K = Int entspricht das Ergebnis f_(N-1) der Reduzierung über die einzige Eingabedimension der folgenden rekursiven Anwendung:

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

Wenn Sie diese Reduktion auf ein Array von Werten und ein Array mit sequenziellen Indizes (z.B. „iota“) anwenden, werden die Arrays gemeinsam durchlaufen und ein Tupel mit dem maximalen Wert und dem entsprechenden Index zurückgegeben.

ReducePrecision

Weitere Informationen finden Sie unter XlaBuilder::ReducePrecision.

Hiermit wird die Umwandlung von Gleitkommawerten in ein Format mit niedrigerer Genauigkeit (z. B. IEEE-FP16) und zurück in das Originalformat simuliert. Die Anzahl der Exponenten- und Mantissenbits im Format mit niedrigerer Genauigkeit kann beliebig festgelegt werden. Allerdings werden möglicherweise nicht alle Bitgrößen von allen Hardwareimplementierungen unterstützt.

ReducePrecision(operand, mantissa_bits, exponent_bits)

Das Ergebnis ist ein Array vom Typ T. Die Eingabewerte werden auf den nächsten Wert gerundet, der mit der angegebenen Anzahl von Mantissenbits darstellbar ist (mit der Semantik „Gleichgewicht auf gerade Zahlen“). Alle Werte, die den durch die Anzahl der Exponentenbits angegebenen Bereich überschreiten, werden auf positive oder negative Unendlichkeit begrenzt. NaN-Werte werden beibehalten, können aber in kanonische NaN-Werte umgewandelt werden.

Das Format mit niedrigerer Genauigkeit muss mindestens ein Exponentenbit haben, um einen Nullwert von einer Unendlichkeit zu unterscheiden, da beide eine Mantisse von null haben. Außerdem muss es eine nicht negative Anzahl von Mantissenbits haben. Die Anzahl der Exponenten- oder Mantissenbits kann den entsprechenden Wert für den Typ T überschreiten. Der entsprechende Teil der Konvertierung ist dann einfach eine Null-Operation.

ReduceScatter

Weitere Informationen finden Sie unter XlaBuilder::ReduceScatter.

ReduceScatter ist ein kollektiver Vorgang, der effektiv ein AllReduce ausführt und das Ergebnis dann verteilt, indem es entlang der scatter_dimension in shard_count Blöcke aufgeteilt wird. Das ith-Speichersegment wird an das Replikat i in der Replikatgruppe gesendet.

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

• Wenn operand ein Tupel von Arrays ist, wird die Reduzierungs-/Scattering-Operation auf jedes Element des Tupels angewendet.
• replica_groups ist eine Liste von Replikationsgruppen, zwischen denen die Reduzierung durchgeführt wird. Die Replikations-ID für das aktuelle Replik kann mit ReplicaId abgerufen werden. Die Reihenfolge der Replikate in jeder Gruppe bestimmt die Reihenfolge, in der das Ergebnis der Allreduce-Funktion verteilt wird. replica_groups muss entweder leer sein (in diesem Fall gehören alle Replikate zu einer einzigen Gruppe) oder dieselbe Anzahl von Elementen wie die Anzahl der Replikate enthalten. Wenn es mehrere Replikationsgruppen gibt, müssen sie alle dieselbe Größe haben. Beispiel: replica_groups = {0, 2}, {1, 3} führt eine Reduzierung zwischen den Replicas 0 und 2 sowie 1 und 3 durch und verteilt dann das Ergebnis.
• shard_count ist die Größe der einzelnen Replikationsgruppen. Wir benötigen diese Informationen, wenn replica_groups leer ist. Wenn replica_groups nicht leer ist, muss shard_count der Größe jeder Replikatgruppe entsprechen.
• channel_id wird für die moduleübergreifende Kommunikation verwendet: Nur reduce-scatter-Vorgänge mit demselben channel_id können miteinander kommunizieren.

Die Ausgabeform ist die Eingabeform, bei der die scatter_dimension um shard_count mal kleiner ist. Wenn es beispielsweise zwei Replikats gibt und der Operand in den beiden Replikats jeweils den Wert [1.0, 2.25] und [3.0, 5.25] hat, ist der Ausgabewert dieser Operation, bei der scatter_dim = 0 ist, [4.0] für das erste Replikat und [7.5] für das zweite Replikat.

ReduceWindow

Weitere Informationen finden Sie unter XlaBuilder::ReduceWindow.

Wendet eine Reduktionsfunktion auf alle Elemente in jedem Fenster einer Sequenz von N mehrdimensionalen Arrays an und gibt ein einzelnes oder ein Tupel von N mehrdimensionalen Arrays als Ausgabe zurück. 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 der angewendete computation immer an den init_values auf der linken Seite übergeben.

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

Wobei:

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

Der folgende Code und die folgende Abbildung zeigen ein Beispiel für die Verwendung von ReduceWindow. Die Eingabe ist eine Matrix der Größe [4 x 6] und sowohl „window_dimensions“ als auch „window_stride_dimensions“ haben die Größe [2 x 3].

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

Ein Schritt von 1 in einer Dimension gibt an, dass sich die Position eines Fensters in der Dimension um ein Element vom benachbarten Fenster entfernt befindet. Wenn Sie angeben möchten, dass sich keine Fenster überschneiden, müssen „window_stride_dimensions“ und „window_dimensions“ gleich sein. Die Abbildung unten zeigt die Verwendung von zwei verschiedenen Schrittwerten. Das Padding wird auf jede Dimension der Eingabe angewendet. Die Berechnungen erfolgen so, als hätte die Eingabe die Dimensionen, die sie nach dem Padding hat.

Ein nicht triviales Beispiel für Padding ist die Berechnung des Minimums des reduzierten Fensters (Anfangswert MAX_FLOAT) mit der Dimension 3 und dem Schritt 2 über das Eingabearray [10000, 1000, 100, 10, 1]. Mit kValid werden Mindestwerte über zwei gültige Fenster berechnet: [10000, 1000, 100] und [100, 10, 1]. Das Ergebnis ist [100, 1]. Mit Padding kSame wird das Array zuerst so aufgefüllt, dass die Form nach dem Reduzierungsfenster derselben wie die Eingabe für Schritt 1 entspricht. Dazu werden an beiden Seiten Anfangselemente hinzugefügt, sodass [MAX_VALUE, 10000, 1000, 100, 10, 1, MAX_VALUE] entsteht. Wenn Sie „reduce-window“ auf das umrahmte Array anwenden, werden drei Fenster [MAX_VALUE, 10000, 1000], [1000, 100, 10] und [10, 1, MAX_VALUE] verwendet und es wird [1000, 10, 1] zurückgegeben.

Die Auswertungsreihenfolge der Reduzierungsfunktion ist beliebig und kann nicht deterministisch sein. Daher sollte die Reduzierungsfunktion nicht zu empfindlich auf eine erneute Verknüpfung reagieren. Weitere Informationen finden Sie in der Diskussion zur 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 Ganzzahl ohne Vorzeichen im Intervall [0, N), wobei N die Anzahl der Replikate ist. Da auf allen Replicas dasselbe Programm ausgeführt wird, gibt ein ReplicaId()-Aufruf im Programm auf jedem Replikat einen anderen Wert zurück.

Neu gestalten

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

Die Dimensionen eines Arrays werden in eine neue Konfiguration umgewandelt.

Reshape(operand, dimensions)

Bei der Umwandlung wird ein Array zuerst in einen eindimensionalen Datenvektor umgewandelt und dann in eine neue Form gebracht. Die Eingabeargumente sind ein beliebiges Array vom Typ T, ein Compilezeitkonstanter Vektor von Dimensionsindizes und ein Compilezeitkonstanter Vektor von Dimensionsgrößen für das Ergebnis. Der dimensions-Vektor bestimmt die Größe des Ausgabearrays. Der Wert an Index 0 in dimensions ist die Größe der Dimension 0, der Wert an Index 1 ist die Größe der Dimension 1 usw. Das Produkt der dimensions-Dimensionen muss dem Produkt der Dimensionsgrößen des Operanden entsprechen. Wenn Sie das minimierte Array in das von dimensions definierte mehrdimensionale Array verfeinern, werden die Dimensionen in dimensions von der langsamsten (höchste Hauptdimension) zur schnellsten (niedrigste untergeordnete Dimension) sortiert.

Angenommen, v ist ein Array mit 24 Elementen:

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

let v012_24 = Reshape(v, {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, {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} };

Als Sonderfall kann reshape ein Array mit einem Element in ein Skalar und umgekehrt umwandeln. Beispiel:

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

Rev (Rückwärts)

Weitere Informationen finden Sie unter XlaBuilder::Rev.

Rev(operand, dimensions)

Die Reihenfolge der Elemente im operand-Array wird entlang des angegebenen dimensions umgekehrt. Dadurch wird ein Ausgabearray derselben Form generiert. Jedes Element des Operandenarrays an einem mehrdimensionalen Index wird im Ausgabearray an einem transformierten Index gespeichert. Der mehrdimensionale Index wird transformiert, indem der Index in jeder Dimension umgekehrt wird, die umgekehrt werden soll. Wenn eine Dimension der Größe N eine der Dimensionen ist, die umgekehrt werden sollen, wird ihr Index i in N − 1 − i umgewandelt.

Der Rev-Vorgang kann beispielsweise verwendet werden, um das Array mit den Convolutionsgewichten 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 Gleitkommaelementartyp haben. Außerdem müssen die Parameter skalare Werte haben.

RngNormal(mu, sigma, shape)

RngUniform

Weitere Informationen finden Sie unter XlaBuilder::RngUniform.

Erstellt eine Ausgabe einer bestimmten Form mit Zufallszahlen, die gemäß der gleichmäßigen Verteilung über das Intervall $[a,b)$generiert werden. Die Parameter und der Ausgabeelementtyp müssen ein boolescher Typ, ein Ganzzahltyp oder ein Gleitkommatyp sein. Die Typen müssen übereinstimmen. 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 skalar sein. Wenn $b <= a$ das Ergebnis implementierungsabhängig ist.

RngUniform(a, b, shape)

RngBitGenerator

Er generiert eine Ausgabe mit einer bestimmten Form, die mithilfe des angegebenen Algorithmus (oder des Backend-Standards) mit gleichmäßig verteilten Zufallsbits gefüllt ist, und gibt einen aktualisierten Zustand (mit derselben Form wie der Anfangszustand) und die generierten Zufallsdaten zurück.

Der Anfangsstatus ist der Anfangsstatus der aktuellen Zufallszahlengenerierung. Die erforderliche Form und die gültigen Werte hängen vom verwendeten Algorithmus ab.

Die Ausgabe ist garantiert eine deterministische Funktion des Anfangszustands, aber nicht deterministisch zwischen Back-Ends und verschiedenen Compilerversionen.

RngBitGenerator(algorithm, key, shape)

Verfügbare Werte für algorithm:

• rng_default: Back-End-spezifischer Algorithmus mit back-end-spezifischen Formanforderungen.

• rng_three_fry: Der PRNG-Algorithmus „ThreeFry“ basiert auf einem Zähler. Die Form initial_state ist u64[2] mit beliebigen Werten. Salmon et al. SC 2011. Parallele Zufallszahlen: So einfach wie das Einmaleins.

• rng_philox: Philox-Algorithmus zum Generieren von Zufallszahlen parallel. Die Form initial_state ist u64[3] mit beliebigen Werten. Salmon et al. SC 2011. Parallele Zufallszahlen: So einfach wie das Einmaleins.

Streudiagramm

Der XLA-Speichervorgang generiert eine Sequenz von Ergebnissen, die die Werte des Eingabearrays operands sind. Dabei werden mehrere Scheiben (an den durch scatter_indices angegebenen Indizes) mit der Sequenz von Werten in updates mithilfe von update_computation aktualisiert.

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)

Wobei:

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

Wenn index_vector_dim mit scatter_indices.rank übereinstimmt, wird davon ausgegangen, dass scatter_indices eine nachgestellte 1-Dimension hat.

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

Die Argumente von „scatter“ müssen den folgenden Einschränkungen entsprechen:

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

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

  • Wenn i in update_window_dims vorhanden ist (d.h. update_window_dims[k] für eine bestimmte k entspricht), darf die Grenze der Dimension i in updates die entsprechende Grenze von operand nicht überschreiten, nachdem die inserted_window_dims berücksichtigt wurden (d.h. adjusted_window_bounds[k], wobei adjusted_window_bounds die Grenzen von operand ohne die Grenzen an den Indizes inserted_window_dims enthält).
  • Wenn i in update_scatter_dims vorhanden ist (d.h. für eine bestimmte k gleich 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, andernfalls scatter_indices.shape.dims[k+1]).

• update_window_dims muss in aufsteigender Reihenfolge vorliegen, darf keine sich wiederholenden Dimensionsnummern enthalten und muss im Bereich [0, updates.rank) liegen.

• inserted_window_dims muss in aufsteigender Reihenfolge vorliegen, darf keine sich wiederholenden Dimensionsnummern enthalten und muss 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 scatter_indices.shape.dims[index_vector_dim] entsprechen und die Werte müssen im Bereich [0, operand.rank) liegen.

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

1. Angenommen, G = {U[k] für k in update_scatter_dims}. Verwenden Sie G, um einen Indexvektor S im Array scatter_indices so abzurufen, dass S[i] = scatter_indices[Combine(G, i)] ist, wobei Combine(A, b) B an den Positionen index_vector_dim in A einfügt.
2. Erstellen Sie einen Index Sin in operand mit S, indem Sie S mithilfe der Karte scatter_dims_to_operand_dims streuen. Formeller:
   1. Sin[scatter_dims_to_operand_dims[k]] = S[k], wenn k < scatter_dims_to_operand_dims.size.
   2. Sin[_] = 0 andernfalls.
3. Erstellen Sie einen Index Win in jedem operands-Array, indem Sie die Indizes in 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 liegt. Dabei ist window_dims_to_operand_dims die monotone Funktion mit dem Definitionsbereich [0, update_window_dims.size) und dem Wertebereich [0, operand.rank) \ inserted_window_dims. Wenn update_window_dims.size beispielsweise 4, operand.rank 6 und inserted_window_dims {0, 2} ist, ist window_dims_to_operand_dims {0→1, 1→3, 2→4, 3→5}.
   2. Win[_] = 0 andernfalls.
4. I ist Win + Sin, wobei „+“ die elementweise Addition ist.

Zusammenfassend lässt sich der Streuungsvorgang so 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 Updates angewendet werden, ist nicht deterministisch. Wenn also mehrere Indizes in updates auf denselben Index in operands verweisen, ist der entsprechende Wert in output nicht deterministisch.

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

Wenn indices_are_sorted auf „wahr“ gesetzt ist, kann XLA davon ausgehen, dass scatter_indices vom Nutzer in aufsteigender Reihenfolge sortiert wurde (nachdem die Werte gemäß scatter_dims_to_operand_dims verstreut wurden). Andernfalls werden die Semantiken von der Implementierung definiert.

Wenn unique_indices auf „wahr“ gesetzt ist, kann XLA davon ausgehen, dass alle Elemente, auf die gesendet wird, eindeutig sind. So können für XLA nicht atomare Vorgänge verwendet werden. Wenn unique_indices auf „wahr“ gesetzt ist und die Indizes, auf die verteilt wird, nicht eindeutig sind, wird die Semantik durch die Implementierung definiert.

Die Streuvorgänge können als Umkehr der Sammelvorgänge betrachtet werden. Das bedeutet, dass die Streuvorgänge die Elemente in der Eingabe aktualisieren, die durch die entsprechenden Sammelvorgänge extrahiert werden.

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

Auswählen

Weitere Informationen finden Sie unter XlaBuilder::Select.

Erstellt ein Ausgabearray aus Elementen von zwei Eingabearrays basierend auf den Werten eines Prädikatsarrays.

Select(pred, on_true, on_false)

Die Arrays on_true und on_false müssen dieselbe Form haben. Dies ist auch die Form des Ausgabearrays. Das Array pred muss dieselbe Dimension wie on_true und on_false haben und den Elementtyp PRED.

Für jedes Element P von pred wird das entsprechende Element des Ausgabe-Arrays 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 Broadcasting-Technologie kann pred ein Skalar vom Typ PRED sein. In diesem Fall wird das Ausgabearray vollständig aus on_true übernommen, wenn pred den Wert true hat, und aus on_false, wenn pred den Wert false hat.

Beispiel mit nicht skalarem pred:

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

Beispiel mit skalarem pred:

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

Die Auswahl zwischen Tupeln wird unterstützt. Tupel gelten zu diesem Zweck als Skalartypen. 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.

Dieser Vorgang kann als zusammengesetzter Vorgang betrachtet werden, bei dem zuerst ReduceWindow auf dem 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. Mit der binären Funktion select wird ein Element aus jedem Fenster ausgewählt, indem sie 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 Funktion select gibt true zurück, wenn der erste Parameter ausgewählt ist, und false, wenn der zweite Parameter ausgewählt ist. Außerdem muss die Funktion transitiv sein, d. h., wenn select(a, b) und select(b, c) true sind, dann ist auch select(a, c) true. So hängt das ausgewählte Element nicht von der Reihenfolge der Elemente ab, die für ein bestimmtes Fenster durchlaufen werden.

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

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

Er kombiniert die beiden Parameter und gibt einen Skalarwert zurück, mit dem der Wert am ausgewählten Index im Ausgabearray aktualisiert wird. Anfänglich sind alle Indizes des Ausgabearrays auf init_value festgelegt.

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

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

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

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

Senden

Weitere Informationen finden Sie unter XlaBuilder::Send.

Send(operand, channel_handle)

Sendet die angegebenen Operandendaten an eine Recv-Anweisung in einer anderen Berechnung, die denselben Channel-Handle verwendet. Es werden keine Daten zurückgegeben.

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

Send(HloInstruction operand, int64 channel_id)

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

SendDone(HloInstruction context)

Wartet bei einem durch eine Send-Anweisung erstellten Kontext auf den Abschluss der Datenübertragung. Die Anweisung gibt keine Daten zurück.

Anleitung zum Planen von Kanälen

Die Ausführungsreihenfolge der vier Anweisungen für jeden Kanal (Recv, RecvDone, Send, SendDone) ist unten dargestellt.

• Recv geschieht vor dem Send
• Send geschieht vor dem RecvDone
• Recv geschieht vor dem RecvDone
• Send geschieht vor dem SendDone

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

Die Einschränkung für die Anweisungen gilt nur für TPUs zur Laufzeit. Auf der GPU werden send und recv blockiert und es werden erst nach einem Handshake zwischen den Quell- und Zielgeräten Daten gesendet.

Slice

Weitere Informationen finden Sie unter XlaBuilder::Slice.

Mit dem Slicing wird ein Teilarray aus dem Eingabearray extrahiert. Das Unterarray hat dieselbe Anzahl von Dimensionen wie die Eingabe und enthält die Werte innerhalb eines Begrenzungsrahmens im Eingabearray. Die Dimensionen und Indizes des Begrenzungsrahmens werden als Argumente für den Slice-Vorgang angegeben.

Slice(operand, start_indices, limit_indices, strides)

Eindimensionales Beispiel:

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

2-dimensionales Beispiel:

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)

Wenn nur ein Operand angegeben wird:

• Wenn der Operand ein eindimensionaler Tensor (ein Array) ist, ist das Ergebnis ein sortiertes Array. Wenn Sie das Array in aufsteigender Reihenfolge sortieren möchten, sollte der Vergleicher einen Vergleich durchführen, der kleiner als ist. Nach der Sortierung des Arrays gilt für alle Indexpositionen i, j mit i < j entweder comparator(value[i], value[j]) = comparator(value[j], value[i]) = false oder comparator(value[i], value[j]) = true.

• Wenn der Operand mehr Dimensionen hat, wird er nach der angegebenen Dimension sortiert. Bei einem zweidimensionalen Tensor (einer Matrix) werden beispielsweise mit dem Dimensionswert 0 alle Spalten unabhängig voneinander sortiert und mit dem Dimensionswert 1 alle Zeilen unabhängig voneinander. Wenn keine Dimensionsnummer angegeben wird, wird standardmäßig die letzte Dimension ausgewählt. Für die sortierte Dimension gilt dieselbe Sortierreihenfolge wie im eindimensionalen Fall.

Wenn n > 1-Operanden angegeben werden:

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

• Alle Operanden werden gemeinsam und nicht einzeln sortiert. Die Operanden werden konzeptionell als Tupel behandelt. Wenn geprüft werden soll, ob die Elemente der einzelnen Operanden an den Indexpositionen i und j vertauscht werden müssen, wird der Vergleicher mit 2 * n Skalarparametern aufgerufen. Dabei entspricht der Parameter 2 * k dem Wert an Position i des Operanden k-th und der Parameter 2 * k + 1 dem Wert an Position j des Operanden k-th. Normalerweise würde der Vergleicher also die Parameter 2 * k und 2 * k + 1 miteinander vergleichen und gegebenenfalls andere Parameterpaare als Entscheidungskriterium verwenden.

• Das Ergebnis ist ein Tupel, das die Operanden in sortierter Reihenfolge (wie oben angegeben) enthält. Der i-th-Operand des Tupels entspricht dem i-th-Operanden von Sort.

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

Wenn is_stable auf „wahr“ gesetzt ist, ist die Sortierung garantiert stabil. Das bedeutet, dass die relative Reihenfolge der gleichwertigen Werte beibehalten wird, wenn es Elemente gibt, die vom Vergleicher als gleich betrachtet werden. Zwei Elemente e1 und e2 sind genau dann gleich, wenn comparator(e1, e2) = comparator(e2, e1) = false. Standardmäßig ist is_stable auf „false“ festgelegt.

TopK

Weitere Informationen finden Sie unter XlaBuilder::TopK.

TopK ermittelt die Werte und Indizes der k größten oder kleinsten Elemente für die letzte Dimension des angegebenen Tensors.

TopK(operand, k, largest)

Ermittelt für einen eindimensionalen Eingabetensor (ein Array) die k größten oder kleinsten Einträge im Array und gibt ein Tupel aus zwei Arrays (values, indices) aus. values[j] ist also der j-größte/kleinste Eintrag in operand und sein Index ist indices[j].

Bei einem Eingabetensor mit mehr als einer Dimension werden die k höchsten Einträge entlang der letzten Dimension berechnet. Alle anderen Dimensionen (Zeilen) bleiben in der Ausgabe erhalten. Bei einem Operanden der Form [A, B, ..., P, Q] mit Q >= k ist die Ausgabe also ein Tupel (values, indices) mit den folgenden Eigenschaften:

values.shape = indices.shape = [A, B, ..., P, k]

Wenn zwei Elemente in einer Zeile übereinstimmen, wird das Element mit dem niedrigeren Index zuerst angezeigt.

Transponieren

Weitere Informationen finden Sie unter dem Vorgang tf.reshape.

Transpose(operand)

Die Operandendimensionen werden mit der angegebenen Permutation permutiert, also ∀ i . 0 ≤ i < number of dimensions ⇒ input_dimensions[permutation[i]] = output_dimensions[i].

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

TriangularSolve

Weitere Informationen finden Sie unter XlaBuilder::TriangularSolve.

Löst lineare Gleichungssysteme mit unter- oder oberdiagonalen Koeffizientenmatrizen durch Vorwärts- oder Rückwärtssubstitution. Durch das Broadcasten entlang der vorangestellten Dimensionen löst diese Routine eines der Matrixsysteme op(a) * x = b oder x * op(a) = b für die Variable x unter Berücksichtigung 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)