Im Folgenden wird die Semantik von Vorgängen beschrieben, die in der XlaBuilder
-Schnittstelle definiert sind. In der Regel sind diese Vorgänge 1:1-Vorgängen zugeordnet, die in der RPC-Schnittstelle in xla_data.proto
definiert sind.
Ein Hinweis zur Nomenklatur: Der generalisierte Datentyp XLA ist ein n-dimensionales Array, das Elemente eines bestimmten Typs (z. B. 32-Bit-Gleitkommazahl) enthält. In der gesamten Dokumentation wird array verwendet, um ein Array mit beliebiger Dimension anzugeben. Sonderfälle haben spezifischere und bekanntere Namen. Ein Vektor ist beispielsweise ein eindimensionales Array und eine Matrix ein zweidimensionales Array.
AfterAll
Weitere Informationen finden Sie unter XlaBuilder::AfterAll
.
„AfterAll“ nimmt eine variable Anzahl von Tokens an und erzeugt ein einzelnes Token. Tokens sind primitive Typen, die zwischen Nebenwirkungsvorgängen angeordnet werden können, um eine Sortierung zu erzwingen. AfterAll
kann als Join von Tokens verwendet werden, um einen Vorgang nach festgelegten Vorgängen zu sortieren.
AfterAll(operands)
Argumente | Typ | Semantik |
---|---|---|
operands |
XlaOp |
unterschiedliche Anzahl von Tokens |
AllGather
Weitere Informationen finden Sie unter XlaBuilder::AllGather
.
Führt eine Verkettung über alle Replikate hinweg durch.
AllGather(operand, all_gather_dim, shard_count, replica_group_ids,
channel_id)
Argumente | Typ | Semantik |
---|---|---|
operand
|
XlaOp
|
Array zur Verkettung über Replikate hinweg |
all_gather_dim |
int64 |
Verkettungsdimension |
replica_groups
|
Vektor der Vektoren von int64 |
Gruppen, zwischen denen die Verkettung durchgeführt wird, |
channel_id
|
int64 optional
|
Optionale Kanal-ID für die modulübergreifende Kommunikation |
replica_groups
ist eine Liste von Replikatgruppen, zwischen denen die Verkettung erfolgt (die Replikat-ID für das aktuelle Replikat kann mitReplicaId
abgerufen werden). Die Reihenfolge der Replikate in den einzelnen Gruppen bestimmt die Reihenfolge, in der sich ihre Eingaben im Ergebnis befinden.replica_groups
muss entweder leer sein (in diesem Fall gehören alle Replikate zu einer einzigen Gruppe, sortiert von0
bisN - 1
) oder dieselbe Anzahl von Elementen wie die Anzahl der Replikate enthalten.replica_groups = {0, 2}, {1, 3}
führt beispielsweise eine Verkettung zwischen den Replikaten0
und2
sowie1
und3
durch.shard_count
ist die Größe der einzelnen Replikatgruppen. Wir benötigen dies, wennreplica_groups
leer ist.channel_id
wird für die modulübergreifende Kommunikation verwendet: Nurall-gather
-Vorgänge mit derselbenchannel_id
können miteinander kommunizieren.
Die Ausgabeform ist die Eingabeform mit dem all_gather_dim
-Wert, der shard_count
-mal größer ist. Wenn beispielsweise zwei Replikate vorhanden sind und der Operand auf den beiden Replikaten den Wert [1.0, 2.5]
bzw. [3.0, 5.25]
hat, ist der Ausgabewert aus diesem Vorgang, bei dem all_gather_dim
0
ist, für beide Replikate [1.0, 2.5, 3.0,
5.25]
.
AllReduce
Weitere Informationen finden Sie unter XlaBuilder::AllReduce
.
Führt eine benutzerdefinierte Berechnung für alle Replikate durch.
AllReduce(operand, computation, replica_group_ids, channel_id)
Argumente | Typ | Semantik |
---|---|---|
operand
|
XlaOp
|
Array oder ein nicht leeres Tupel von Arrays, um die Anzahl der Replikate zu reduzieren |
computation |
XlaComputation |
Reduktionsberechnung |
replica_groups
|
Vektor der Vektoren von int64 |
Gruppen, zwischen denen die Reduzierungen durchgeführt werden, |
channel_id
|
int64 optional
|
Optionale Kanal-ID für die modulübergreifende Kommunikation |
- Wenn
operand
ein Tupel von Arrays ist, wird die All-Reduce-Funktion für jedes Element des Tupels ausgeführt. replica_groups
ist eine Liste von Replikatgruppen, zwischen denen die Reduzierung durchgeführt wird (Replikat-ID für das aktuelle Replikat kann mitReplicaId
abgerufen werden).replica_groups
muss entweder leer sein (in diesem Fall alle Replikate zu einer einzigen Gruppe gehören) oder dieselbe Anzahl von Elementen enthalten wie die Anzahl der Replikate. Beispiel:replica_groups = {0, 2}, {1, 3}
führt eine Reduzierung zwischen den Replikaten0
und2
sowie1
und3
durch.channel_id
wird für die modulübergreifende Kommunikation verwendet: Nurall-reduce
-Vorgänge mit derselbenchannel_id
können miteinander kommunizieren.
Die Ausgabeform ist mit der Eingabeform identisch. Wenn beispielsweise zwei Replikate vorhanden sind und der Operand in den beiden Replikaten den Wert [1.0, 2.5]
bzw. [3.0, 5.25]
hat, ist der Ausgabewert dieses Vorgangs und dieser Summenberechnung für beide Replikate [4.0, 7.75]
. Wenn die Eingabe ein Tupel ist, ist auch die Ausgabe ein Tupel.
Für die Berechnung des Ergebnisses von AllReduce
ist eine Eingabe von jedem Replikat erforderlich. Wenn also ein Replikat einen AllReduce
-Knoten öfter als ein anderes ausführt, wartet das vorherige Replikat unbegrenzt. Da die Replikate alle dasselbe Programm ausführen, gibt es nicht viele Möglichkeiten, dies zu tun. Es ist jedoch möglich, dass die Bedingung einer Während-Schleife von Daten aus der Einspeisung abhängt und die eingespeisten Daten dazu führen, dass die Während-Schleife mehr Mal ein Replikat wiederholt als ein anderes.
AllToAll
Weitere Informationen finden Sie unter XlaBuilder::AllToAll
.
AllToAll ist ein kollektiver Vorgang, bei dem Daten von allen Kernen an alle Kerne gesendet werden. Er besteht aus zwei Phasen:
- Die Streuphase. Auf jedem Kern wird der Operand in eine Anzahl von
split_count
Blöcken entlang dessplit_dimensions
aufgeteilt und die Blöcke werden auf alle Kerne verteilt, sodass der i-te-Block zum i-ten Kern gesendet wird. - Die Erfassungsphase. Jeder Kern verkettet die empfangenen Blöcke entlang des
concat_dimension
.
Die teilnehmenden Kerne können folgendermaßen konfiguriert werden:
replica_groups
: Jede Replikatgruppe enthält eine Liste der an der Berechnung beteiligten Replikat-IDs (die Replikat-ID für das aktuelle Replikat kann mithilfe vonReplicaId
abgerufen werden). „AllToAll“ wird in Untergruppen in der angegebenen Reihenfolge angewendet.replica_groups = { {1,2,3}, {4,5,0} }
bedeutet beispielsweise, dass ein AllToAll innerhalb der Replikate{1, 2, 3}
und in der Erfassungsphase angewendet wird und die empfangenen Blöcke in der gleichen Reihenfolge von 1, 2, 3 verkettet werden. Dann wird ein weiteres "AllToAll" in den Replikaten 4, 5, 0 angewendet und die Verkettungsreihenfolge ist ebenfalls 4, 5, 0. Wennreplica_groups
leer ist, gehören alle Replikate in der Verkettungsreihenfolge zu einer Gruppe.
Voraussetzungen:
- Die Dimensionsgröße des Operanden auf
split_dimension
ist durchsplit_count
teilbar. - Die Form des Operanden ist kein Tupel.
AllToAll(operand, split_dimension, concat_dimension, split_count,
replica_groups)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
n-dimensionales Eingabearray |
split_dimension
|
int64
|
Ein Wert im Intervall [0,
n) , der die Dimension angibt, nach der der Operand aufgeteilt wird |
concat_dimension
|
int64
|
Ein Wert im Intervall [0,
n) , der die Dimension angibt, mit der die Aufteilungsblöcke verkettet sind |
split_count
|
int64
|
Die Anzahl der Kerne, die an diesem Vorgang beteiligt sind. Wenn replica_groups leer ist, sollte dies die Anzahl der Replikate sein. Andernfalls sollte dies der Anzahl der Replikate in jeder Gruppe entsprechen. |
replica_groups
|
Vektor für ReplicaGroup
|
Jede Gruppe enthält eine Liste mit Replikat-IDs. |
Unten sehen Sie ein Beispiel für Alltoall.
XlaBuilder b("alltoall");
auto x = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {4, 16}), "x");
AllToAll(x, /*split_dimension=*/1, /*concat_dimension=*/0, /*split_count=*/4);
In diesem Beispiel nehmen am Alltoall 4 Kerne teil. Auf jedem Kern wird der Operand entlang der Dimension 0 in vier Teile geteilt, sodass jedes Teil die Form f32[4,4] hat. Die vier Teile sind auf alle Kerne verteilt. Dann verkettet jeder Kern die empfangenen Teile entlang der Dimension 1 in der Reihenfolge von Kern 0 bis 4. Die Ausgabe auf jedem Kern hat also die Form f32[16,4].
BatchNormGrad
Eine ausführliche Beschreibung des Algorithmus finden Sie auch unter XlaBuilder::BatchNormGrad
und im ursprünglichen Batch-Normalisierungspapier.
Berechnet Gradienten der Batchnorm.
BatchNormGrad(operand, scale, mean, variance, grad_output, epsilon, feature_index)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Zu normalisierendes n-dimensionales Array (x) |
scale |
XlaOp |
1 dimensionales Array (\(\gamma\)) |
mean |
XlaOp |
1 dimensionales Array (\(\mu\)) |
variance |
XlaOp |
1 dimensionales Array (\(\sigma^2\)) |
grad_output |
XlaOp |
Farbverläufe an BatchNormTraining (\(\nabla y\)) weitergegeben |
epsilon |
float |
Epsilon-Wert (\(\epsilon\)) |
feature_index |
int64 |
Index für Featuredimension in operand |
Für jedes Feature in der Feature-Dimension (feature_index
ist der Index für die Feature-Dimension in operand
) berechnet der Vorgang die Gradienten in Bezug auf operand
, offset
und scale
über alle anderen Dimensionen hinweg. feature_index
muss ein gültiger Index für die Feature-Dimension in operand
sein.
Die drei Gradienten werden durch die folgenden Formeln definiert, wobei von einem vierdimensionalen Array als operand
und einem Featuredimensionsindex l
, der Batchgröße m
und den räumlichen Größen w
und h
ausgegangen wird:
\[ \begin{split} c_l&= \frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \left( \nabla y_{ijkl} \frac{x_{ijkl} - \mu_l}{\sigma^2_l+\epsilon} \right) \\\\ d_l&= \frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \nabla y_{ijkl} \\\\ \nabla x_{ijkl} &= \frac{\gamma_{l} }{\sqrt{\sigma^2_{l}+\epsilon} } \left( \nabla y_{ijkl} - d_l - c_l (x_{ijkl} - \mu_{l}) \right) \\\\ \nabla \gamma_l &= \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \left( \nabla y_{ijkl} \frac{x_{ijkl} - \mu_l}{\sqrt{\sigma^2_{l}+\epsilon} } \right) \\\\\ \nabla \beta_l &= \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \nabla y_{ijkl} \end{split} \]
Die Eingaben mean
und variance
stellen Momentwerte in Batch- und räumlichen Dimensionen dar.
Der Ausgabetyp ist ein Tupel aus drei Handles:
Ausgaben | Typ | Semantik |
---|---|---|
grad_operand
|
XlaOp
|
Gradienten in Bezug auf die Eingabe operand ($\nabla
x$) |
grad_scale
|
XlaOp
|
Gradienten in Bezug auf die Eingabe scale ($\nabla
\gamma$) |
grad_offset
|
XlaOp
|
Gradient in Bezug auf die Eingabe offset ($\nabla
\beta$) |
BatchNormInference
Eine ausführliche Beschreibung des Algorithmus finden Sie auch unter XlaBuilder::BatchNormInference
und im ursprünglichen Batch-Normalisierungspapier.
Normalisiert ein Array über Batch- und räumliche Dimensionen hinweg.
BatchNormInference(operand, scale, offset, mean, variance, epsilon, feature_index)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
n-dimensionales Array, das normalisiert werden soll |
scale |
XlaOp |
1-dimensionales Array |
offset |
XlaOp |
1-dimensionales Array |
mean |
XlaOp |
1-dimensionales Array |
variance |
XlaOp |
1-dimensionales Array |
epsilon |
float |
Epsilon-Wert |
feature_index |
int64 |
Index für Featuredimension in operand |
Für jedes Element in der Featuredimension (feature_index
ist der Index für die Merkmalsdimension in operand
) berechnet der Vorgang den Mittelwert und die Varianz über alle anderen Dimensionen hinweg und verwendet den Mittelwert und die Varianz, um jedes Element in operand
zu normalisieren. feature_index
muss ein gültiger Index für die Featuredimension in operand
sein.
BatchNormInference
entspricht dem Aufruf von BatchNormTraining
, ohne mean
und variance
für jeden Batch zu berechnen. Stattdessen werden die eingegebenen mean
und variance
als geschätzte Werte verwendet. Der Zweck dieses Vorgangs ist es, die Latenz in der Inferenz zu verringern, daher der Name BatchNormInference
.
Die Ausgabe ist ein n-dimensionales, normalisiertes Array mit derselben Form wie die Eingabe-operand
.
BatchNormTraining
Eine ausführliche Beschreibung des Algorithmus finden Sie unter XlaBuilder::BatchNormTraining
und the original batch normalization paper
.
Normalisiert ein Array über Batch- und räumliche Dimensionen hinweg.
BatchNormTraining(operand, scale, offset, epsilon, feature_index)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Zu normalisierendes n-dimensionales Array (x) |
scale |
XlaOp |
1 dimensionales Array (\(\gamma\)) |
offset |
XlaOp |
1 dimensionales Array (\(\beta\)) |
epsilon |
float |
Epsilon-Wert (\(\epsilon\)) |
feature_index |
int64 |
Index für Featuredimension in operand |
Für jedes Element in der Featuredimension (feature_index
ist der Index für die Merkmalsdimension in operand
) berechnet der Vorgang den Mittelwert und die Varianz über alle anderen Dimensionen hinweg und verwendet den Mittelwert und die Varianz, um jedes Element in operand
zu normalisieren. feature_index
muss ein gültiger Index für die Featuredimension in operand
sein.
Der Algorithmus gilt für jeden Batch in operand
\(x\) , der m
-Elemente mit w
und h
als Größe der räumlichen Dimensionen enthält (vorausgesetzt, operand
ist ein vierdimensionales Array):
Berechnet den Batchmittelwert \(\mu_l\) für jedes Feature
l
in der Feature-Dimension: \(\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}\)Berechnet die Batchvarianz \(\sigma^2_l\): $\sigma^2l=\frac{1}{mwh}\sum{i=1}^m\sum{j=1}^w\sum{k=1}^h (x_{ijkl} - \mu_l)^2$
Normalisiert, skaliert und verschiebt: \(y_{ijkl}=\frac{\gamma_l(x_{ijkl}-\mu_l)}{\sqrt[2]{\sigma^2_l+\epsilon} }+\beta_l\)
Der Epsilon-Wert, normalerweise eine kleine Zahl, wird addiert, um Fehler durch Division durch Null zu vermeiden.
Der Ausgabetyp ist ein Tupel aus drei XlaOp
s:
Ausgaben | Typ | Semantik |
---|---|---|
output
|
XlaOp
|
n-dimensionales Array mit derselben Form wie die Eingabeoperand (y) |
batch_mean |
XlaOp |
1 dimensionales Array (\(\mu\)) |
batch_var |
XlaOp |
1 dimensionales Array (\(\sigma^2\)) |
batch_mean
und batch_var
sind Momente, die über die Batch- und räumlichen Dimensionen mit den obigen Formeln berechnet werden.
BitcastConvertType
Weitere Informationen finden Sie unter XlaBuilder::BitcastConvertType
.
Führt ähnlich wie bei einem tf.bitcast
in TensorFlow eine elementweise Bitcast-Operation von einer Datenform zu einer Zielform aus. Die Eingabe- und Ausgabegröße muss übereinstimmen: s32
-Elemente werden z.B. über die Bitcast-Routine zu f32
-Elementen und ein s32
-Element wird zu vier s8
-Elementen. Bitcast wird als Low-Level-Umwandlung implementiert, sodass Maschinen mit unterschiedlichen Gleitkommadarstellungen unterschiedliche Ergebnisse liefern.
BitcastConvertType(operand, new_element_type)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Array vom Typ T mit Dimmung D |
new_element_type |
PrimitiveType |
Typ U |
Die Abmessungen des Operanden und der Zielform müssen übereinstimmen, abgesehen von der letzten Dimension, die sich durch das Verhältnis der primitiven Größe vor und nach der Konvertierung ändert.
Die Quell- und Zielelementtypen dürfen keine Tupel sein.
Bitcast-Konvertierung in einen primitiven Typ mit unterschiedlicher Breite
Die HLO-Anweisung BitcastConvert
unterstützt den Fall, in dem die Größe des Ausgabeelementtyps T'
nicht mit der Größe des Eingabeelements T
übereinstimmt. Da der gesamte Vorgang konzeptionell ein Bitcast ist und die zugrunde liegenden Byte nicht geändert werden, muss sich die Form des Ausgabeelements ändern. Für B = sizeof(T), B' =
sizeof(T')
gibt es zwei mögliche Fälle.
Erstens: Wenn B > B'
, erhält die Ausgabeform eine neue kleinere Dimension der Größe B/B'
. Beispiel:
f16[10,2]{1,0} %output = f16[10,2]{1,0} bitcast-convert(f32[10]{0} %input)
Die Regel für effektive Skalare bleibt gleich:
f16[2]{0} %output = f16[2]{0} bitcast-convert(f32[] %input)
Alternativ erfordert die Anweisung für B' > B
, dass die letzte logische Dimension der Eingabeform gleich B'/B
ist. Diese Dimension wird während der Konvertierung entfernt:
f32[10]{0} %output = f32[10]{0} bitcast-convert(f16[10,2]{1,0} %input)
Beachten Sie, dass Konvertierungen zwischen verschiedenen Bitbreiten nicht elementweise erfolgen.
Nachricht an alle
Weitere Informationen finden Sie unter XlaBuilder::Broadcast
.
Fügt einem Array Dimensionen durch Duplizieren der Daten im Array hinzu.
Broadcast(operand, broadcast_sizes)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Das zu duplizierende Array |
broadcast_sizes |
ArraySlice<int64> |
Die Größen der neuen Dimensionen |
Die neuen Dimensionen werden links eingefügt, d.h., wenn broadcast_sizes
die Werte {a0, ..., aN}
und die Operandenform die Dimensionen {b0, ..., bM}
hat, hat die Form der Ausgabe die Dimensionen {a0, ..., aN, b0, ..., bM}
.
Der neue Dimensionsindex wird in Kopien des Operanden, d.h.
output[i0, ..., iN, j0, ..., jM] = operand[j0, ..., jM]
Beispiel: Wenn operand
ein skalarer f32
mit dem Wert 2.0f
ist und broadcast_sizes
gleich {2, 3}
ist, ist das Ergebnis ein Array mit der Form f32[2, 3]
und alle Werte im Ergebnis sind 2.0f
.
BroadcastInDim
Weitere Informationen finden Sie unter XlaBuilder::BroadcastInDim
.
Erweitert die Größe und den Rang eines Arrays durch Duplizieren der Daten im Array.
BroadcastInDim(operand, out_dim_size, broadcast_dimensions)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Das zu duplizierende Array |
out_dim_size |
ArraySlice<int64> |
Die Größen der Abmessungen der Zielform |
broadcast_dimensions |
ArraySlice<int64> |
welcher Dimension in der Zielform jede Dimension der Operandenform entspricht |
Ähnlich wie „Broadcast“, ermöglicht aber das Hinzufügen von Dimensionen an beliebiger Stelle und das Erweitern vorhandener Dimensionen mit Größe 1.
Das operand
wird an die durch out_dim_size
beschriebene Form übertragen.
broadcast_dimensions
ordnet die Abmessungen von operand
den Abmessungen der Zielform zu, d. h., die i-te Dimension des Operanden wird der broadcast_dimension[i]. Dimension der Ausgabeform zugeordnet. Die Abmessungen von operand
müssen die Größe 1 oder dieselbe Größe wie die Dimension in der Ausgabeform haben, der sie zugeordnet sind. Die übrigen Dimensionen werden mit Abmessungen der Größe 1 gefüllt. Beim Broadcasting mit degenerierenden Dimensionen wird es dann entlang dieser degenerierten Dimensionen übertragen, um die Ausgabeform zu erreichen. Die Semantik wird auf der Broadcasting-Seite ausführlich beschrieben.
Call
Weitere Informationen finden Sie unter XlaBuilder::Call
.
Löst eine Berechnung mit den angegebenen Argumenten aus.
Call(computation, args...)
Argumente | Typ | Semantik |
---|---|---|
computation |
XlaComputation |
Berechnung des Typs T_0, T_1, ..., T_{N-1} -> S mit n Parametern eines beliebigen Typs |
args |
Sequenz von N XlaOp s |
N Argumente eines beliebigen Typs |
Die ARity und die Typen von args
müssen mit den Parametern von computation
übereinstimmen. Es darf keine args
enthalten.
Cholesky
Weitere Informationen finden Sie unter XlaBuilder::Cholesky
.
Berechnet die Cholesky-Zerlegung eines Batches symmetrischer (Hermitianischer) positiver eindeutiger Matrizen.
Cholesky(a, lower)
Argumente | Typ | Semantik |
---|---|---|
a |
XlaOp |
Array > 2 (Rang > 2) eines komplexen oder Gleitkommatyps |
lower |
bool |
Gibt an, ob das obere oder untere Dreieck von a verwendet werden soll. |
Wenn lower
true
ist, werden Matrizen für kleine Dreiecke l
so berechnet, dass $a = l ist .
l^T$. Wenn lower
false
ist, werden die Matrizen des oberen Dreiecks u
so berechnet:\(a = u^T . u\).
Eingabedaten werden je nach Wert von lower
nur aus dem unteren/oberen Dreieck von a
gelesen. Werte aus dem anderen Dreieck werden ignoriert. Ausgabedaten werden im selben Dreieck zurückgegeben. Die Werte im anderen Dreieck sind implementierungsdefiniert und können beliebig sein.
Wenn der Rang von a
größer als 2 ist, wird a
als ein Batch von Matrizen behandelt, wobei alle außer den zwei weiteren Dimensionen Batchdimensionen sind.
Wenn a
nicht symmetrisch (Hermitian) positiv eindeutig ist, ist das Ergebnis von der Implementierung definiert.
Klemme
Weitere Informationen finden Sie unter XlaBuilder::Clamp
.
Bindet einen Operanden innerhalb des Bereichs zwischen einem Mindest- und Höchstwert.
Clamp(min, operand, max)
Argumente | Typ | Semantik |
---|---|---|
min |
XlaOp |
Array vom Typ T |
operand |
XlaOp |
Array vom Typ T |
max |
XlaOp |
Array vom Typ T |
Bei einem Operanden sowie Mindest- und Höchstwert wird der Operand zurückgegeben, wenn er im Bereich zwischen Mindest- und Höchstwert liegt. Andernfalls wird der Mindestwert zurückgegeben, wenn der Operand unter diesem Bereich liegt, oder den Höchstwert, wenn der Operand über diesem Bereich liegt. Der Wert ist clamp(a, x, b) = min(max(a, x), b)
.
Alle drei Arrays müssen die gleiche Form haben. Als eingeschränkte Form der Übertragung können min
und/oder max
ein Skalar vom Typ T
sein.
Beispiel mit dem Skalar min
und max
:
let operand: s32[3] = {-1, 5, 9};
let min: s32 = 0;
let max: s32 = 6;
==>
Clamp(min, operand, max) = s32[3]{0, 5, 6};
Minimieren
Weitere Informationen finden Sie unter XlaBuilder::Collapse
und tf.reshape
.
Minimiert die Dimensionen eines Arrays in eine Dimension.
Collapse(operand, dimensions)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Array vom Typ T |
dimensions |
Vektor für int64 |
Aufeinanderfolgende Teilmenge der T-Dimensionen. |
Bei der Minimierung wird die angegebene Teilmenge der Operandendimensionen durch eine einzelne Dimension ersetzt. Die Eingabeargumente sind ein beliebiges Array vom Typ T und ein Vektor der kompilierten Zeitkonstante von Dimensionsindexen. Die Dimensionsindexe müssen eine fortlaufende Teilmenge der T-Dimensionen (niedrige bis hohe Dimensionszahlen) sein. Somit sind {0, 1, 2}, {0, 1} oder {1, 2} gültige Dimensionssätze, {1, 0} oder {0, 2} jedoch nicht. Sie werden durch eine einzelne neue Dimension ersetzt, die sich an derselben Position in der Abfolge der Dimensionen befindet wie die, die sie ersetzen. Die neue Dimensionsgröße entspricht dem Produkt der ursprünglichen Abmessungen. Die niedrigste Dimensionszahl in dimensions
ist die langsamste variierende Dimension (die größte) in der Schleifenverschachtelung, bei der diese Dimensionen minimiert werden. Die höchste Dimensionsnummer variiert am schnellsten (die geringste). Wenn eine allgemeinere Minimierungsreihenfolge erforderlich ist, verwenden Sie den Operator tf.reshape
.
Nehmen wir z. B. „v“ als Array mit 24 Elementen:
let v = f32[4x2x3] { { {10, 11, 12}, {15, 16, 17} },
{ {20, 21, 22}, {25, 26, 27} },
{ {30, 31, 32}, {35, 36, 37} },
{ {40, 41, 42}, {45, 46, 47} } };
// Collapse to a single dimension, leaving one dimension.
let v012 = Collapse(v, {0,1,2});
then v012 == f32[24] {10, 11, 12, 15, 16, 17,
20, 21, 22, 25, 26, 27,
30, 31, 32, 35, 36, 37,
40, 41, 42, 45, 46, 47};
// Collapse the two lower dimensions, leaving two dimensions.
let v01 = Collapse(v, {0,1});
then v01 == f32[4x6] { {10, 11, 12, 15, 16, 17},
{20, 21, 22, 25, 26, 27},
{30, 31, 32, 35, 36, 37},
{40, 41, 42, 45, 46, 47} };
// Collapse the two higher dimensions, leaving two dimensions.
let v12 = Collapse(v, {1,2});
then v12 == f32[8x3] { {10, 11, 12},
{15, 16, 17},
{20, 21, 22},
{25, 26, 27},
{30, 31, 32},
{35, 36, 37},
{40, 41, 42},
{45, 46, 47} };
CollectivePermute
Weitere Informationen finden Sie unter XlaBuilder::CollectivePermute
.
CollectivePermute ist ein kollektiver Vorgang, der Daten über mehrere Replikate hinweg sendet und empfängt.
CollectivePermute(operand, source_target_pairs)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
n-dimensionales Eingabearray |
source_target_pairs |
Vektor für <int64, int64> |
Eine Liste mit Paaren aus Quelle_Replikat_ID und Ziel-Replikat_ID. Für jedes Paar wird der Operand vom Quellreplikat zum Zielreplikat gesendet. |
Für source_target_pair
gelten die folgenden Einschränkungen:
- Zwei Paare sollten nicht dieselbe Zielreplikat-ID und dieselbe ID des Quellreplikats haben.
- Wenn eine Replikat-ID kein Ziel in einem Paar ist, ist die Ausgabe für dieses Replikat ein Tensor, der aus Null(en) mit derselben Form wie die Eingabe besteht.
Verketten
Weitere Informationen finden Sie unter XlaBuilder::ConcatInDim
.
Verketten besteht aus mehreren Array-Operanden. Das Array hat denselben Rang wie die einzelnen Operanden des Eingabearrays (die denselben Rang haben müssen wie die anderen) und enthält die Argumente in der Reihenfolge, in der sie angegeben wurden.
Concatenate(operands..., dimension)
Argumente | Typ | Semantik |
---|---|---|
operands |
Sequenz von N XlaOp |
N Arrays vom Typ T mit den Dimensionen [L0, L1, ...]. Erfordert N >= 1. |
dimension |
int64 |
Ein Wert im Intervall [0, N) , der die Dimension angibt, die zwischen operands verkettet werden soll. |
Mit Ausnahme von dimension
müssen alle Dimensionen gleich sein. Das liegt daran, dass XLA keine unregelmäßigen Arrays unterstützt. Werte mit dem Rang 0 können nicht verkettet werden, da es nicht möglich ist, die Dimension zu benennen, in der die Verkettung erfolgt.
Beispiel für eindimensionales Format:
Concat({ {2, 3}, {4, 5}, {6, 7} }, 0)
>>> {2, 3, 4, 5, 6, 7}
Beispiel für zweidimensionale Darstellung:
let a = {
{1, 2},
{3, 4},
{5, 6},
};
let b = {
{7, 8},
};
Concat({a, b}, 0)
>>> {
{1, 2},
{3, 4},
{5, 6},
{7, 8},
}
Diagramm:
Bedingt
Weitere Informationen finden Sie unter XlaBuilder::Conditional
.
Conditional(pred, true_operand, true_computation, false_operand,
false_computation)
Argumente | Typ | Semantik |
---|---|---|
pred |
XlaOp |
Skalar des Typs PRED |
true_operand |
XlaOp |
Argument vom Typ \(T_0\) |
true_computation |
XlaComputation |
XlaComputation vom Typ \(T_0 \to S\) |
false_operand |
XlaOp |
Argument vom Typ \(T_1\) |
false_computation |
XlaComputation |
XlaComputation vom Typ \(T_1 \to S\) |
Führt true_computation
aus, wenn pred
den Wert true
hat, und false_computation
, wenn pred
false
ist, und gibt das Ergebnis zurück.
true_computation
muss ein einzelnes Argument vom Typ \(T_0\) annehmen und wird mit true_operand
aufgerufen, die denselben Typ haben müssen. false_computation
muss ein einzelnes Argument vom Typ \(T_1\) annehmen und wird mit false_operand
aufgerufen, die denselben Typ haben müssen. Der Typ des zurückgegebenen Werts von true_computation
und false_computation
muss identisch sein.
Je nach Wert von pred
wird nur entweder true_computation
oder false_computation
ausgeführt.
Conditional(branch_index, branch_computations, branch_operands)
Argumente | Typ | Semantik |
---|---|---|
branch_index |
XlaOp |
Skalar des Typs S32 |
branch_computations |
Sequenz von N XlaComputation |
XlaComputations vom Typ \(T_0 \to S , T_1 \to S , ..., T_{N-1} \to S\) |
branch_operands |
Sequenz von N XlaOp |
Argumente des Typs \(T_0 , T_1 , ..., T_{N-1}\) |
Führt branch_computations[branch_index]
aus und gibt das Ergebnis zurück. Wenn branch_index
ein S32
ist, der < 0 oder >= N ist, wird branch_computations[N-1]
als Standardzweig ausgeführt.
Jede branch_computations[b]
muss ein einzelnes Argument vom Typ \(T_b\) annehmen und wird mit branch_operands[b]
aufgerufen, die denselben Typ haben müssen. Der Typ der zurückgegebenen Werte für jede branch_computations[b]
muss gleich sein.
Abhängig vom Wert von branch_index
wird nur eine der branch_computations
ausgeführt.
Conv. (Faltung)
Weitere Informationen finden Sie unter XlaBuilder::Conv
.
Als ConvWithGeneralPadding, aber das Padding wird kurz als SAME oder VALID angegeben. Mit dem gleichen Padding wird die Eingabe (lhs
) mit Nullen aufgefüllt, sodass die Ausgabe dieselbe Form wie die Eingabe hat, wenn keine Streifen berücksichtigt werden. GÜLTIGES Padding bedeutet einfach, dass es keine Auffüllung gibt.
ConvWithGeneralPadding (Faltung)
Weitere Informationen finden Sie unter XlaBuilder::ConvWithGeneralPadding
.
Berechnet eine Faltung von der Art, die in neuronalen Netzwerken verwendet wird. Hier kann man sich eine Faltung als ein n-dimensionales Fenster vorstellen, das sich über eine n-dimensionale Grundfläche bewegt und für jede mögliche Position des Fensters eine Berechnung durchgeführt wird.
Argumente | Typ | Semantik |
---|---|---|
lhs |
XlaOp |
n+2-Array von Eingaben einstufen |
rhs |
XlaOp |
n+2-Array der Kernel-Gewichtungen einstufen |
window_strides |
ArraySlice<int64> |
n-d-Array von Kernel-Schritten |
padding |
ArraySlice< pair<int64,int64>> |
n-d-Array des Abstands (niedrig, hoch) |
lhs_dilation |
ArraySlice<int64> |
n-d-lhs-Erweiterungsfaktor-Array |
rhs_dilation |
ArraySlice<int64> |
n-d rhs-Dilationsfaktor-Array |
feature_group_count |
int64 | Anzahl der Funktionsgruppen |
batch_group_count |
int64 | Anzahl der Batchgruppen |
Es soll n die Anzahl der räumlichen Dimensionen sein. Das Argument lhs
ist ein Rang-n+2-Array, das die Grundfläche beschreibt. Diese wird als Eingabe bezeichnet, auch wenn „rhs“ ebenfalls eine Eingabe ist. In einem neuronalen Netzwerk sind das die Eingabeaktivierungen.
Die n+2-Dimensionen sind in dieser Reihenfolge:
batch
: Jede Koordinate in dieser Dimension stellt eine unabhängige Eingabe dar, für die Faltungen ausgeführt werden.z/depth/features
: Jeder (y,x)-Position im Grundbereich ist ein Vektor zugeordnet, der in diese Dimension einfließt.spatial_dims
: Beschreibt die räumlichenn
-Dimensionen, die die Grundfläche definieren, durch die sich das Fenster bewegt.
Das Argument rhs
ist ein n+2-Array, das den Faltungsfilter/Kernel/Fenster beschreibt. Die Dimensionen lauten in dieser Reihenfolge:
output-z
: Die Dimensionz
der Ausgabe.input-z
: Die Größe dieser Dimension multipliziert mitfeature_group_count
muss der Größe der Dimensionz
in LHs entsprechen.spatial_dims
: Beschreibt die räumlichen Dimensionen vonn
, die das n-d-Fenster definieren, das sich über die Grundfläche bewegt.
Das Argument window_strides
gibt den Schritt des Faltungsfensters in den räumlichen Dimensionen an. Wenn der Schritt in der ersten räumlichen Dimension beispielsweise 3 ist, kann das Fenster nur auf Koordinaten platziert werden, bei denen der erste räumliche Index durch 3 teilbar ist.
Das Argument padding
gibt den auf die Basisfläche angewendeten Auffüllwert null an. Der Umfang des Paddings kann negativ sein. Der absolute Wert des negativen Paddings gibt die Anzahl der Elemente an, die vor der Faltung aus der angegebenen Dimension entfernt werden müssen. padding[0]
gibt den Abstand für die Dimension y
an und padding[1]
gibt den Abstand für die Dimension x
an. Jedes Paar hat als erstes Element einen niedrigen und einen hohen Abstand als zweites. Der niedrige Abstand wird in Richtung niedrigerer Indizes angewendet, während der hohe Abstand in Richtung der höheren Indizes angewendet wird. Wenn padding[1]
beispielsweise (2,3)
ist, enthält die zweite räumliche Dimension einen Abstand von zwei Nullen auf der linken Seite und drei Nullen rechts auf der zweiten räumlichen Dimension. Das Padding entspricht dem Einfügen derselben Nullwerte in die Eingabe (lhs
) vor der Faltung.
Die Argumente lhs_dilation
und rhs_dilation
geben den Dilatationsfaktor an, der in jeder räumlichen Dimension auf die LHs bzw. RHS angewendet werden soll. Wenn der Dilatationsfaktor in einer räumlichen Dimension d ist, werden d-1-Löcher implizit zwischen den einzelnen Einträgen in dieser Dimension platziert, wodurch die Größe des Arrays vergrößert wird. Die Löcher werden mit einem No-Op-Wert gefüllt, der für Faltung Nullen bedeutet.
Die Dilatation der RHS wird auch als atroische Faltung bezeichnet. Weitere Informationen finden Sie unter tf.nn.atrous_conv2d
. Die Dilatation des Lps wird auch als transponierte Faltung bezeichnet. Weitere Informationen finden Sie unter tf.nn.conv2d_transpose
.
Das Argument feature_group_count
(Standardwert 1) kann für gruppierte Faltungen verwendet werden. feature_group_count
muss ein Divisor der Eingabe- und der Ausgabefeature-Dimension sein. Wenn feature_group_count
größer als 1 ist, bedeutet dies, dass die Eingabe- und Ausgabefeature-Dimension und die rhs
-Ausgabefeature-Dimension konzeptionell gleichmäßig in viele feature_group_count
-Gruppen aufgeteilt werden, wobei jede Gruppe aus einer aufeinanderfolgenden Untersequenz von Features besteht. Die Eingabefeature-Dimension von rhs
muss der lhs
-Eingabefeature-Dimension geteilt durch feature_group_count
entsprechen (also hat sie bereits die Größe einer Gruppe von Eingabefeatures). Die i-ten Gruppen werden zusammen verwendet, um feature_group_count
für viele separate Faltungen zu berechnen. Die Ergebnisse dieser Faltungen werden in der Ausgabefeature-Dimension miteinander verkettet.
Für die tiefe Faltung wird das Argument feature_group_count
auf die Eingabefeature-Dimension festgelegt und der Filter wird von [filter_height, filter_width, in_channels, channel_multiplier]
in [filter_height, filter_width, 1, in_channels * channel_multiplier]
umgeformt. Weitere Informationen finden Sie unter tf.nn.depthwise_conv2d
.
Das Argument batch_group_count
(Standardwert 1) kann während der Rückpropagierung für gruppierte Filter verwendet werden. batch_group_count
muss ein Divisor der Größe der Batchdimension lhs
(Eingabe) sein. Wenn batch_group_count
größer als 1 ist, bedeutet dies, dass die Dimension des Ausgabe-Batches die Größe input batch
/ batch_group_count
haben muss. batch_group_count
muss ein Divisor der Größe des Ausgabefeatures sein.
Die Ausgabeform hat folgende Abmessungen in dieser Reihenfolge:
batch
: Die Größe dieser Dimension multipliziert mitbatch_group_count
muss der Größe der Dimensionbatch
in LHs entsprechen.z
: Gleiche Größe wieoutput-z
im Kernel (rhs
).spatial_dims
: Ein Wert für jede gültige Platzierung des Faltungsfensters.
Die Abbildung oben zeigt, wie das Feld batch_group_count
funktioniert. Wir teilen jeden LH-Batch in batch_group_count
-Gruppen auf und tun dasselbe für die Ausgabefeatures. Dann führen wir für jede dieser Gruppen paarweise Faltungen durch und verketten die Ausgabe entlang der Ausgabefeature-Dimension. Die operative Semantik aller anderen Dimensionen (Feature und Räumlichkeit) bleibt gleich.
Die gültigen Platzierungen des Faltungsfensters werden durch die Schritte und die Größe der Grundfläche nach dem Auffüllen bestimmt.
Betrachten Sie eine 2d-Faltung und wählen Sie einige feste batch
-, z
-, y
- und x
-Koordinaten in der Ausgabe aus, um zu beschreiben, was eine Faltung bewirkt. Dann ist (y,x)
die Position einer Ecke des Fensters innerhalb des Grundbereichs (z.B. die obere linke Ecke, je nachdem, wie Sie die räumlichen Dimensionen interpretieren). Wir haben jetzt ein 2D-Fenster aus der Grundfläche, in dem jeder 2D-Punkt mit einem 1D-Vektor verknüpft ist, sodass wir eine 3D-Box erhalten. Da wir im Convolutional Kernel die Ausgabekoordinate z
korrigiert haben, haben wir auch eine 3D-Box. Die beiden Boxen haben dieselben Abmessungen, sodass wir die Summe der elementweisen Produkte zwischen den beiden Feldern ermitteln können (ähnlich einem Punktprodukt). Das ist der Ausgabewert.
Wenn output-z
z.B. 5, dann erzeugt jede Position des Fensters 5 Werte in der Ausgabe in der Dimension z
der Ausgabe. Diese Werte unterscheiden sich darin, welcher Teil des Convolutional Kernel verwendet wird. Für jede output-z
-Koordinate wird ein separates 3D-Feld mit Werten verwendet. Sie können sich das als fünf separate Faltungen
mit einem anderen Filter vorstellen.
Hier ist Pseudocode für eine 2D-Faltung mit Auffüllung und Streifen:
for (b, oz, oy, ox) { // output coordinates
value = 0;
for (iz, ky, kx) { // kernel coordinates and input z
iy = oy*stride_y + ky - pad_low_y;
ix = ox*stride_x + kx - pad_low_x;
if ((iy, ix) inside the base area considered without padding) {
value += input(b, iz, iy, ix) * kernel(oz, iz, ky, kx);
}
}
output(b, oz, oy, ox) = value;
}
ConvertElementType
Weitere Informationen finden Sie unter XlaBuilder::ConvertElementType
.
Führt ähnlich wie bei einem elementweisen static_cast
in C++ einen elementweisen Konvertierungsvorgang von einer Datenform in eine Zielform aus. Die Dimensionen müssen übereinstimmen und die Konvertierung erfolgt elementweise. So werden z. B. s32
-Elemente über eine s32
-zu-f32
-Konvertierungsroutine zu f32
-Elementen.
ConvertElementType(operand, new_element_type)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Array vom Typ T mit Dimmung D |
new_element_type |
PrimitiveType |
Typ U |
Die Abmessungen des Operanden und der Zielform müssen übereinstimmen. Die Quell- und Zielelementtypen dürfen keine Tupel sein.
Eine Konvertierung wie T=s32
in U=f32
führt eine normalisierende Int-to-Float-Konvertierungsroutine durch, z. B. „Round-to-Nearest-Even“.
let a: s32[3] = {0, 1, 2};
let b: f32[3] = convert(a, f32);
then b == f32[3]{0.0, 1.0, 2.0}
CrossReplicaSum
Führt AllReduce
mit einer Summenberechnung durch.
CustomCall
Weitere Informationen finden Sie unter XlaBuilder::CustomCall
.
Rufen Sie eine vom Nutzer bereitgestellte Funktion innerhalb einer Berechnung auf.
CustomCall(target_name, args..., shape)
Argumente | Typ | Semantik |
---|---|---|
target_name |
string |
Name der Funktion. Es wird eine Aufrufanweisung ausgegeben, die auf diesen Symbolnamen abzielt. |
args |
Sequenz von N XlaOp s |
N Argumente eines beliebigen Typs, die an die Funktion übergeben werden. |
shape |
Shape |
Ausgabeform der Funktion |
Die Funktionssignatur ist unabhängig von der Arity oder dem Typ der Argumente dieselbe:
extern "C" void target_name(void* out, void** in);
Beispiel: CustomCall wird wie folgt verwendet:
let x = f32[2] {1,2};
let y = f32[2x3] { {10, 20, 30}, {40, 50, 60} };
CustomCall("myfunc", {x, y}, f32[3x3])
Hier ist ein Beispiel für eine Implementierung von myfunc
:
extern "C" void myfunc(void* out, void** in) {
float (&x)[2] = *static_cast<float(*)[2]>(in[0]);
float (&y)[2][3] = *static_cast<float(*)[2][3]>(in[1]);
EXPECT_EQ(1, x[0]);
EXPECT_EQ(2, x[1]);
EXPECT_EQ(10, y[0][0]);
EXPECT_EQ(20, y[0][1]);
EXPECT_EQ(30, y[0][2]);
EXPECT_EQ(40, y[1][0]);
EXPECT_EQ(50, y[1][1]);
EXPECT_EQ(60, y[1][2]);
float (&z)[3][3] = *static_cast<float(*)[3][3]>(out);
z[0][0] = x[1] + y[1][0];
// ...
}
Die vom Nutzer bereitgestellte Funktion darf keine Nebeneffekte haben und muss idempotent ausgeführt werden.
Dot
Weitere Informationen finden Sie unter XlaBuilder::Dot
.
Dot(lhs, rhs)
Argumente | Typ | Semantik |
---|---|---|
lhs |
XlaOp |
Array vom Typ T |
rhs |
XlaOp |
Array vom Typ T |
Die genaue Semantik dieser Operation hängt von den Rängen der Operanden ab:
Eingabe | Ausgabe | Semantik |
---|---|---|
Vektor [n] dot Vektor [n] |
Skalar | Vektorpunktprodukt |
Matrix [m x k] dot Vektor [k] |
Vektor [m] | Matrix-Vektor-Multiplikation |
Matrix [m x k] dot Matrix [k x n] |
Matrix [m x n] | Matrixmultiplikation |
Der Vorgang führt die Summe der Produkte über die zweite Dimension von lhs
(oder die erste Dimension, wenn sie Rang 1 hat) und die erste Dimension von rhs
aus. Dies sind die vertraglich vereinbarten Dimensionen. Die vertraglich vereinbarten Abmessungen von lhs
und rhs
müssen dieselbe Größe haben. In der Praxis können damit Punktprodukte zwischen Vektoren, Vektor-/Matrixmultiplikationen oder Matrix-/Matrixmultiplikationen ausgeführt werden.
DotGeneral
Weitere Informationen finden Sie unter XlaBuilder::DotGeneral
.
DotGeneral(lhs, rhs, dimension_numbers)
Argumente | Typ | Semantik |
---|---|---|
lhs |
XlaOp |
Array vom Typ T |
rhs |
XlaOp |
Array vom Typ T |
dimension_numbers |
DotDimensionNumbers |
Vertrags- und Batchdimensionsnummern |
Ähnlich wie bei Dot, allerdings können für lhs
und rhs
Vertrags- und Batch-Dimensionsnummern angegeben werden.
DotDimensionNumbers-Felder | Typ | Semantik |
---|---|---|
lhs_contracting_dimensions
|
wiederholter int64-Wert | lhs Zahlen der Vertragsmaße |
rhs_contracting_dimensions
|
wiederholter int64-Wert | rhs Zahlen der Vertragsmaße |
lhs_batch_dimensions
|
wiederholter int64-Wert | lhs Batch-Dimensionsnummern |
rhs_batch_dimensions
|
wiederholter int64-Wert | rhs Batch-Dimensionsnummern |
DotGeneral führt die Summe der Produkte über die in dimension_numbers
angegebenen Reduzierungsdimensionen aus.
Die zugehörigen vertraglich vereinbarten Dimensionsnummern von lhs
und rhs
müssen nicht identisch sein, müssen aber dieselben Dimensionsgrößen haben.
Beispiel für das Minimieren von Dimensionszahlen:
lhs = { {1.0, 2.0, 3.0},
{4.0, 5.0, 6.0} }
rhs = { {1.0, 1.0, 1.0},
{2.0, 2.0, 2.0} }
DotDimensionNumbers dnums;
dnums.add_lhs_contracting_dimensions(1);
dnums.add_rhs_contracting_dimensions(1);
DotGeneral(lhs, rhs, dnums) -> { {6.0, 12.0},
{15.0, 30.0} }
Verknüpfte Batch-Dimensionsnummern aus lhs
und rhs
müssen dieselben Dimensionsgrößen haben.
Beispiel mit Batch-Dimensionsnummern (Batch-Größe 2, 2 x 2-Matrizen):
lhs = { { {1.0, 2.0},
{3.0, 4.0} },
{ {5.0, 6.0},
{7.0, 8.0} } }
rhs = { { {1.0, 0.0},
{0.0, 1.0} },
{ {1.0, 0.0},
{0.0, 1.0} } }
DotDimensionNumbers dnums;
dnums.add_lhs_contracting_dimensions(2);
dnums.add_rhs_contracting_dimensions(1);
dnums.add_lhs_batch_dimensions(0);
dnums.add_rhs_batch_dimensions(0);
DotGeneral(lhs, rhs, dnums) -> { { {1.0, 2.0},
{3.0, 4.0} },
{ {5.0, 6.0},
{7.0, 8.0} } }
Eingabe | Ausgabe | Semantik |
---|---|---|
[b0, m, k] dot [b0, k, n] |
[b0, m, n] | Batch Matmul |
[b0, b1, m, k] dot [b0, b1, k, n] |
[b0, b1, m, n] | Batch Matmul |
Daraus folgt, dass die resultierende Dimensionsnummer mit der Batchdimension beginnt, dann mit der Dimension lhs
ohne Vertragsbindung/ohne Batch und schließlich mit der Dimension rhs
ohne Vertragsbindung/kein Batch.
DynamicSlice
Weitere Informationen finden Sie unter XlaBuilder::DynamicSlice
.
DynamicSlice extrahiert ein Teilarray aus dem Eingabearray bei dynamischem start_indices
. Die Größe des Segments in jeder Dimension wird in size_indices
übergeben, das den Endpunkt exklusiver Segmentintervalle in jeder Dimension angibt: [Start, Start + Größe). Die Form von start_indices
muss ein Rang == 1 sein, wobei die Dimensionsgröße dem Rang von operand
entspricht.
DynamicSlice(operand, start_indices, size_indices)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
N-dimensionales Array des Typs T |
start_indices |
Sequenz von N XlaOp |
Liste von n skalaren Ganzzahlen, die die Startindexe des Segments für jede Dimension enthalten. Wert muss größer oder gleich null sein. |
size_indices |
ArraySlice<int64> |
Liste von N Ganzzahlen, die die Segmentgröße für jede Dimension enthalten. Jeder Wert muss größer als null sein und „start“ und „size“ müssen kleiner oder gleich der Größe der Dimension sein, damit die Modulo-Dimensionsgröße nicht umbrochen wird. |
Die effektiven Slice-Indizes werden berechnet, indem die folgende Transformation auf jeden Index i
in [1, N)
angewendet wird, bevor das Slice ausgeführt wird:
start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - size_indices[i])
Dadurch wird sichergestellt, dass das extrahierte Segment immer innerhalb der Grenzen des Operandenarrays liegt. Wenn das Slice innerhalb der Grenzen liegt, bevor die Transformation angewendet wird, hat die Transformation keine Auswirkungen.
Beispiel für eindimensionales Format:
let a = {0.0, 1.0, 2.0, 3.0, 4.0}
let s = {2}
DynamicSlice(a, s, {2}) produces:
{2.0, 3.0}
Beispiel für zweidimensionale Darstellung:
let b =
{ {0.0, 1.0, 2.0},
{3.0, 4.0, 5.0},
{6.0, 7.0, 8.0},
{9.0, 10.0, 11.0} }
let s = {2, 1}
DynamicSlice(b, s, {2, 2}) produces:
{ { 7.0, 8.0},
{10.0, 11.0} }
DynamicUpdateSlice
Weitere Informationen finden Sie unter XlaBuilder::DynamicUpdateSlice
.
DynamicUpdateSlice generiert ein Ergebnis, das dem Wert des Eingabearrays operand
entspricht, wobei das Segment update
bei start_indices
überschrieben wird.
Die Form von update
bestimmt die Form des Teilarrays des Ergebnisses, das aktualisiert wird.
Die Form von start_indices
muss ein Rang == 1 sein, wobei die Dimensionsgröße dem Rang von operand
entspricht.
DynamicUpdateSlice(operand, update, start_indices)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
N-dimensionales Array des Typs T |
update |
XlaOp |
N-dimensionales Array des Typs T, das die Segmentaktualisierung enthält Jede Dimension der Aktualisierungsform muss größer als null sein und die Größe von Start und Aktualisierung muss kleiner oder gleich der Operandengröße für jede Dimension sein, um zu vermeiden, dass die Aktualisierungsindexe außerhalb des gültigen Bereichs liegen. |
start_indices |
Sequenz von N XlaOp |
Liste von n skalaren Ganzzahlen, die die Startindexe des Segments für jede Dimension enthalten. Wert muss größer oder gleich null sein. |
Die effektiven Slice-Indizes werden berechnet, indem die folgende Transformation auf jeden Index i
in [1, N)
angewendet wird, bevor das Slice ausgeführt wird:
start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - update.dimension_size[i])
Dadurch wird sichergestellt, dass das aktualisierte Segment immer innerhalb der Grenzen des Operandenarrays liegt. Wenn das Slice innerhalb der Grenzen liegt, bevor die Transformation angewendet wird, hat die Transformation keine Auswirkungen.
Beispiel für eindimensionales Format:
let a = {0.0, 1.0, 2.0, 3.0, 4.0}
let u = {5.0, 6.0}
let s = {2}
DynamicUpdateSlice(a, u, s) produces:
{0.0, 1.0, 5.0, 6.0, 4.0}
Beispiel für zweidimensionale Darstellung:
let b =
{ {0.0, 1.0, 2.0},
{3.0, 4.0, 5.0},
{6.0, 7.0, 8.0},
{9.0, 10.0, 11.0} }
let u =
{ {12.0, 13.0},
{14.0, 15.0},
{16.0, 17.0} }
let s = {1, 1}
DynamicUpdateSlice(b, u, s) produces:
{ {0.0, 1.0, 2.0},
{3.0, 12.0, 13.0},
{6.0, 14.0, 15.0},
{9.0, 16.0, 17.0} }
Elementweise binäre Arithmetische Operationen
Weitere Informationen finden Sie unter XlaBuilder::Add
.
Es wird eine Reihe von elementweisen binären arithmetischen Operationen unterstützt.
Op(lhs, rhs)
Dabei ist Op
entweder Add
(Addition), Sub
(Subtraktion), Mul
(Multiplikation), Div
(Division), Rem
(Rest), Max
(Maximum), Min
(Minimum), LogicalAnd
(logisches UND) oder LogicalOr
(logisches ODER).
Argumente | Typ | Semantik |
---|---|---|
lhs |
XlaOp |
Linker Operand: Array vom Typ „T“ |
rhs |
XlaOp |
right-hand-side operand: Array vom Typ T |
Die Formen der Argumente müssen entweder ähnlich oder kompatibel sein. In der Broadcasting-Dokumentation erfahren Sie, was es bedeutet, wenn Formen kompatibel sind. Das Ergebnis einer Operation hat eine Form, die das Ergebnis der Übertragung der beiden Eingabearrays ist. Bei dieser Variante werden Operationen zwischen Arrays unterschiedlicher Ränge nicht unterstützt, es sei denn, einer der Operanden ist ein Skalar.
Wenn Op
den Wert Rem
hat, wird das Vorzeichen des Ergebnisses vom Dividenden entnommen. Der absolute Wert des Ergebnisses ist immer kleiner als der absolute Wert des Divisors.
Der Ganzzahl-Divisionsüberlauf (vorzeichenbehaftete/unsignierte Division/Rest durch null oder vorzeichenbehaftete Division/Restung von INT_SMIN
mit -1
) erzeugt einen Implementierungswert.
Für diese Vorgänge gibt es eine alternative Variante, die Broadcasting unterschiedlicher Ränge unterstützt:
Op(lhs, rhs, broadcast_dimensions)
Dabei gilt: Op
ist der gleiche wie oben. Diese Variante der Operation sollte für arithmetische Operationen zwischen Arrays mit verschiedenen Rängen verwendet werden (z. B. zum Hinzufügen einer Matrix zu einem Vektor).
Der zusätzliche broadcast_dimensions
-Operand ist ein Segment von Ganzzahlen, mit dem der Rang des Operanden mit niedrigerem Rang bis zum Rang des Operanden mit höherem Rang erweitert wird. broadcast_dimensions
ordnet die Abmessungen der Form mit niedrigerem Rang den Abmessungen der Form mit höherem Rang zu. Die nicht zugeordneten Abmessungen der
maximierten Form werden mit Maßen der Größe 1 ausgefüllt. Beim Broadcasting nach degenerierenden Dimensionen werden dann die Formen entlang dieser Degenerationsdimensionen gesendet, um die Formen beider Operanden anzugleichen. Die Semantik wird auf der Broadcasting-Seite ausführlich beschrieben.
Elementweise Vergleichsvorgänge
Weitere Informationen finden Sie unter XlaBuilder::Eq
.
Es wird eine Reihe standardmäßiger elementweiser binärer Vergleichsvorgänge unterstützt. Beachten Sie, dass beim Vergleich von Gleitkommatypen nach IEEE 754 die Standardsemantik für den Gleitkommavergleich gilt.
Op(lhs, rhs)
Dabei ist Op
einer der folgenden Werte: Eq
(gleich), Ne
(ungleich), Ge
(größer oder gleich), Gt
(größer als), Le
(kleiner oder gleich), Lt
(kleiner als). Ein anderer Satz von Operatoren, EqTotalOrder, NeTotalOrder, GeTotalOrder, GtTotalOrder, LeTotalOrder und LtTotalOrder, bietet dieselben Funktionen, mit der Ausnahme, dass sie zusätzlich eine Gesamtreihenfolge über den Gleitkommazahlen unterstützen, indem sie -NaN < -Inf < -Finite < -0 < +0 < +InaN.
Argumente | Typ | Semantik |
---|---|---|
lhs |
XlaOp |
Linker Operand: Array vom Typ „T“ |
rhs |
XlaOp |
right-hand-side operand: Array vom Typ T |
Die Formen der Argumente müssen entweder ähnlich oder kompatibel sein. In der Broadcasting-Dokumentation erfahren Sie, was es bedeutet, wenn Formen kompatibel sind. Das Ergebnis einer Operation hat eine Form, die das Ergebnis der Übertragung der beiden Eingabearrays mit dem Elementtyp PRED
ist. Bei dieser Variante werden Vorgänge zwischen Arrays unterschiedlicher Ränge nicht unterstützt, es sei denn, einer der Operanden ist ein Skalar.
Für diese Vorgänge gibt es eine alternative Variante, die Broadcasting unterschiedlicher Ränge unterstützt:
Op(lhs, rhs, broadcast_dimensions)
Dabei gilt: Op
ist der gleiche wie oben. Diese Variante des Vorgangs sollte für Vergleichsvorgänge zwischen Arrays mit unterschiedlichen Rängen verwendet werden, z. B. zum Hinzufügen einer Matrix zu einem Vektor.
Der zusätzliche broadcast_dimensions
-Operand ist ein Segment von Ganzzahlen, die die Dimensionen für die Übertragung der Operanden angeben. Die Semantik wird auf der Broadcasting-Seite ausführlich beschrieben.
Elementweise unäre Funktionen
XlaBuilder unterstützt die folgenden elementweisen unären Funktionen:
Abs(operand)
Elementweise abs. x -> |x|
.
Ceil(operand)
Elementweises Maximum x -> ⌈x⌉
.
Cos(operand)
Elementweiser Kosinus x -> cos(x)
.
Exp(operand)
Elementweises natürliches Exponential x -> e^x
.
Floor(operand)
Elementweise Etage x -> ⌊x⌋
.
Imag(operand)
Elementweise imaginärer Teil einer komplexen (oder reellen) Form. x -> imag(x)
. Wenn der Operand ein Gleitkommatyp ist, wird 0 zurückgegeben.
IsFinite(operand)
prüft, ob jedes Element von operand
endlich ist, d.h. nicht positiv oder negativ unendlich und nicht NaN
ist. Gibt ein Array von PRED
-Werten mit derselben Form wie die Eingabe zurück, wobei jedes Element nur dann true
ist, wenn das entsprechende Eingabeelement endlich ist.
Log(operand)
Elementweiser natürlicher Logarithmus x -> ln(x)
.
LogicalNot(operand)
Elementweise logisch nicht x -> !(x)
.
Logistic(operand)
Berechnung elementweise logistische Funktionen x ->
logistic(x)
.
PopulationCount(operand)
Berechnet die Anzahl der Bits, die in jedem Element von operand
festgelegt sind.
Neg(operand)
Elementweise Negation x -> -x
.
Real(operand)
Elementweiser reeller Teil einer komplexen (oder reellen) Form.
x -> real(x)
. Wenn der Operand ein Gleitkommatyp ist, wird derselbe Wert zurückgegeben.
Rsqrt(operand)
Elementweise Kehrzahl der Quadratwurzeloperation x -> 1.0 / sqrt(x)
.
Sign(operand)
Vorgang für elementweises Signieren x -> sgn(x)
, wobei
\[\text{sgn}(x) = \begin{cases} -1 & x < 0\\ -0 & x = -0\\ NaN & x = NaN\\ +0 & x = +0\\ 1 & x > 0 \end{cases}\]
mithilfe des Vergleichsoperators des Elementtyps operand
.
Sqrt(operand)
Elementweiser Quadratwurzelvorgang x -> sqrt(x)
.
Cbrt(operand)
Elementweiser kubischer Stammvorgang x -> cbrt(x)
.
Tanh(operand)
Elementweiser hyperbolischer Tangens x -> tanh(x)
.
Round(operand)
Elementweise Rundung, gleicht Null.
RoundNearestEven(operand)
Elementweise Rundung, wird auf die nächste Gerade ausgerichtet.
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Der Operand der Funktion |
Die Funktion wird auf jedes Element im Array operand
angewendet, sodass ein Array mit derselben Form entsteht. operand
darf ein Skalar sein (Rang 0).
FFT
Der XLA-FFT-Vorgang implementiert die Vorwärts- und Inverse Fourier-Transformationen für echte und komplexe Ein-/Ausgaben. Mehrdimensionale FFTs auf bis zu drei Achsen werden unterstützt.
Weitere Informationen finden Sie unter XlaBuilder::Fft
.
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Das Array, das mit der Fourier-Transformation transformiert wird. |
fft_type |
FftType |
Siehe die folgende Tabelle. |
fft_length |
ArraySlice<int64> |
Die Zeitdomainlängen der transformierten Achsen. Dies ist insbesondere erforderlich, damit IRFFT die Größe der innersten Achse korrekt anpasst, da RFFT(fft_length=[16]) dieselbe Ausgabeform wie RFFT(fft_length=[17]) hat. |
FftType |
Semantik |
---|---|
FFT |
Komplexe bis komplexe FFT weiterleiten. Die Form bleibt unverändert. |
IFFT |
Umgekehrte komplex-komplexe FFT. Die Form bleibt unverändert. |
RFFT |
Leitet reelle bis komplexe FFT-Werte weiter. Die Form der innersten Achse wird auf fft_length[-1] // 2 + 1 reduziert, wenn fft_length[-1] ein Wert ungleich null ist, wobei der umgekehrte konjugierte Teil des transformierten Signals jenseits der Nyquist-Frequenz weggelassen wird. |
IRFFT |
Umgekehrte reelle zu komplexe FFT (nimmt komplex an, gibt reelle Zahlen zurück). Die Form der innersten Achse wird zu fft_length[-1] erweitert, wenn fft_length[-1] ein Wert ungleich null ist. Dabei wird der Teil des transformierten Signals jenseits der Nyquist-Frequenz von der umgekehrten Konjugation der 1 - zu fft_length[-1] // 2 + 1 -Einträge abgeleitet. |
Mehrdimensionale FFT
Wenn mehr als 1 fft_length
angegeben ist, entspricht dies der Anwendung einer Kaskade von FFT-Vorgängen auf jede der innersten Achsen. Beachten Sie, dass für die reellen und komplexesten Achsen die Transformation der innersten Achse (effektiv) an erster Stelle (RFFT; zuletzt für IRFFT) ausgeführt wird. Deshalb ändert sich die Größe der innersten Achse. Andere Achsentransformationen sind dann komplex->komplex.
Implementierungsdetails
CPU-FFT wird von Eigens TensorFFT unterstützt. GPU-FFT verwendet cuFFT.
Zusammentragen
Der XLA-Abrufvorgang fügt mehrere Segmente eines Eingabearrays zusammen (jedes Segment mit einem möglicherweise anderen Laufzeitversatz).
Allgemeine Semantik
Weitere Informationen finden Sie unter XlaBuilder::Gather
.
Eine intuitivere Beschreibung finden Sie unten im Abschnitt „Informelle Beschreibung“.
gather(operand, start_indices, offset_dims, collapsed_slice_dims, slice_sizes, start_index_map)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Das Array, aus dem wir uns sammeln. |
start_indices |
XlaOp |
Array mit den Startindexen der erfassten Segmente. |
index_vector_dim |
int64 |
Die Dimension in start_indices , die die Startindexe „enthält“. Unten finden Sie eine detaillierte Beschreibung. |
offset_dims |
ArraySlice<int64> |
Die Dimensionen in der Ausgabeform, die in ein vom Operanden unterteiltes Array verschoben werden. |
slice_sizes |
ArraySlice<int64> |
slice_sizes[i] sind die Grenzen für das Segment bei der Dimension i . |
collapsed_slice_dims |
ArraySlice<int64> |
Die Abmessungen in jedem Segment, die weggeklappt sind. Diese Abmessungen müssen Größe 1 haben. |
start_index_map |
ArraySlice<int64> |
Eine Map, die beschreibt, wie Indexe in start_indices den rechtlichen Indizes in Operanden zugeordnet werden. |
indices_are_sorted |
bool |
Gibt an, ob die Indizes garantiert vom Aufrufer sortiert werden. |
unique_indices |
bool |
Gibt an, ob die Indizes für den Aufrufer garantiert eindeutig sind. |
Der Einfachheit halber benennen wir Dimensionen im Ausgabearray, die nicht in offset_dims
enthalten sind, als batch_dims
.
Die Ausgabe ist ein Array mit dem Rang batch_dims.size
+ offset_dims.size
.
operand.rank
muss der Summe von offset_dims.size
und collapsed_slice_dims.size
entsprechen. Außerdem muss slice_sizes.size
gleich operand.rank
sein.
Wenn index_vector_dim
gleich start_indices.rank
ist, betrachten wir implizit start_indices
als eine nachgestellte 1
-Dimension (d.h., wenn start_indices
die Form [6,7]
hat und index_vector_dim
2
ist, wird die Form von start_indices
implizit als [6,7,1]
betrachtet.
Die Grenzen des Ausgabearrays entlang der Dimension i
werden so berechnet:
Wenn
i
inbatch_dims
vorhanden ist (d. h. für einigek
gleichbatch_dims[k]
ist), wählen wir die entsprechenden Dimensionsgrenzen ausstart_indices.shape
aus und überspringenindex_vector_dim
. Wählen Sie alsostart_indices.shape.dims
[k
] aus, wennk
<index_vector_dim
und ansonstenstart_indices.shape.dims
[k
+1
].Wenn
i
inoffset_dims
vorhanden ist (d. h.offset_dims
[k
] für einigek
), wird die entsprechende Grenze vonslice_sizes
nach Berücksichtigung voncollapsed_slice_dims
ausgewählt. Das heißt, wir wählenadjusted_slice_sizes
[k
] aus, wobeiadjusted_slice_sizes
gleichslice_sizes
ist und die Grenzen bei den Indexencollapsed_slice_dims
entfernt wurden.
Formal wird der Operandenindex In
, der einem bestimmten Ausgabeindex Out
entspricht, so berechnet:
Lassen Sie
G
= {Out
[k
] fürk
inbatch_dims
}. Verwenden SieG
, um einen VektorS
so herauszuschneiden, dassS
[i
] =start_indices
[Zusammenführen(G
,i
)], wobei Kombinieren(A, b) b an Positionindex_vector_dim
in A einfügt. Dies ist auch dann gut definiert, wennG
leer ist: WennG
leer ist, giltS
=start_indices
.Erstellen Sie einen Startindex
S
in
inoperand
mitS
. Dazu verteilen SieS
mitstart_index_map
. Genauer gesagt:S
in
[start_index_map
[k
]] =S
[k
], wennk
<start_index_map.size
.S
Andernfallsin
[_
] =0
.
Erstellen Sie einen Index
O
in
inoperand
, indem Sie die Indexe in den Versatzdimensionen inOut
gemäß dem Satz voncollapsed_slice_dims
verteilen. Genauer gesagt:O
in
[remapped_offset_dims
(k
)] =Out
[offset_dims
[k
]], wennk
<offset_dims.size
(remapped_offset_dims
ist unten definiert).O
Andernfallsin
[_
] =0
.
In
istO
in
+S
in
, wobei + für die elementweise Addition steht.
remapped_offset_dims
ist eine monotone Funktion mit der Domain [0
, offset_dims.size
) und dem Bereich [0
, operand.rank
) \ collapsed_slice_dims
. Wenn z.B. offset_dims.size
ist 4
, operand.rank
ist 6
und collapsed_slice_dims
ist {0
, 2
}. remapped_offset_dims
ist dann {0
→1
,
1
→3
, 2
→4
, 3
→5
}.
Wenn indices_are_sorted
auf „true“ gesetzt ist, kann XLA davon ausgehen, dass start_indices
(in aufsteigender start_index_map
-Reihenfolge) nach Nutzer sortiert sind. Andernfalls ist die Semantik die Implementierung definiert.
Wenn unique_indices
auf „true“ gesetzt ist, kann XLA davon ausgehen, dass alle verstreuten Elemente eindeutig sind. XLA könnte also nicht-atomare Vorgänge verwenden. Wenn unique_indices
auf „true“ gesetzt ist und die Indizes, auf die verstreut sind, nicht eindeutig sind, ist die Semantik die Implementierung definiert.
Informelle Beschreibung und Beispiele
Inoffiziell entspricht jeder Index-Out
im Ausgabearray einem Element-E
im Operandenarray, das so berechnet wird:
Wir verwenden die Batchdimensionen in
Out
, um einen Startindex ausstart_indices
abzurufen.Wir verwenden
start_index_map
, um den Startindex, dessen Größe kleiner als operand.rank sein kann, einem "vollständigen" Startindex imoperand
zuzuordnen.Ein Segment mit der Größe
slice_sizes
wird mithilfe des vollständigen Startindex dynamisch segmentiert.Wir passen die Form des Segments an, indem wir die
collapsed_slice_dims
-Dimensionen minimieren. Da alle Abmessungen des minimierten Segments eine Grenze von 1 haben müssen, ist diese Umformung immer zulässig.Wir verwenden die Offset-Dimensionen in
Out
, um in dieses Segment zu indexieren und das EingabeelementE
zu erhalten, das dem AusgabeindexOut
entspricht.
index_vector_dim
ist in allen folgenden Beispielen auf start_indices.rank
– 1
gesetzt. Interessantere Werte für index_vector_dim
ändern den Vorgang nicht grundlegend, machen die visuelle Darstellung jedoch umständlicher.
Sehen wir uns ein Beispiel an, bei dem fünf Segmente der Form [8,6]
aus einem [16,11]
-Array gesammelt werden, um einen Eindruck davon zu bekommen, wie all dies zusammenpasst. Die Position eines Segments im Array [16,11]
kann als Indexvektor der Form S64[2]
dargestellt werden, sodass der Satz von fünf Positionen als S64[5,2]
-Array dargestellt werden kann.
Das Verhalten des Collect-Vorgangs kann dann als Indextransformation dargestellt werden, die [G
,O
0
,O
1
], einen Index in der Ausgabeform, und diesen folgendermaßen einem Element im Eingabearray zuordnet:
Zunächst wählen wir mit G
einen (X
,Y
) Vektor aus dem Array zum Erfassen der Indexe aus.
Das Element im Ausgabearray bei Index [G
,O
0
,O
1
] ist dann das Element im Eingabearray bei Index [X
+O
0
,Y
+O
1
].
slice_sizes
ist [8,6]
, wodurch der Bereich von O0
und O1
bestimmt wird, was wiederum die Grenzen des Segments bestimmt.
Dieser Erfassungsvorgang fungiert als dynamisches Batchsegment mit G
als Batchdimension.
Die Indizes zum Zusammenführen können mehrdimensional sein. In einer allgemeineren Version des obigen Beispiels mit einem Array zum Sammeln von Indexen mit der Form [4,5,2]
werden Indexe beispielsweise so übersetzt:
Auch dies fungiert als ein dynamisches Batch-Segment G
0
und G
1
als Batchdimensionen. Die Segmentgröße beträgt immer noch [8,6]
.
Der Collect-Vorgang in XLA verallgemeinert die oben beschriebene informelle Semantik auf folgende Weise:
Wir können konfigurieren, welche Dimensionen in der Ausgabeform die Versatzabmessungen sind (Dimensionen, die im letzten Beispiel
O
0
,O
1
enthalten). Die Ausgabe-Batch-Dimensionen (Dimensionen, die im letzten BeispielG
0
undG
1
enthalten) sind als Ausgabedimensionen definiert, die keine Offset-Dimensionen sind.Die Anzahl der Ausgabe-Offset-Dimensionen, die explizit in der Ausgabeform vorhanden sind, kann kleiner als der Eingaberang sein. Diese „fehlenden“ Dimensionen, die explizit als
collapsed_slice_dims
aufgeführt sind, müssen eine Segmentgröße von1
haben. Da sie eine Segmentgröße von1
haben, ist der einzige gültige Index für sie0
. Das Ausschließen dieser Segmente führt nicht zu Mehrdeutigkeit.Das aus dem Array "Sammel-Indexe" (im letzten Beispiel
X
,Y
) extrahierte Segment hat möglicherweise weniger Elemente als der Eingabearrayrang. Eine explizite Zuordnung legt fest, wie der Index erweitert werden soll, damit er den gleichen Rang wie die Eingabe hat.
Als letztes Beispiel verwenden wir (2) und (3) zur Implementierung von tf.gather_nd
:
G
0
und G
1
werden verwendet, um einen Startindex wie gewohnt aus dem Array für die Erfassung der Indexe herauszuschneiden, mit der Ausnahme, dass der Startindex nur ein Element, X
, enthält. Analog dazu gibt es nur einen Ausgabe-Offset-Index mit dem Wert O
0
. Bevor sie jedoch als Indizes im Eingabearray verwendet werden, werden sie gemäß „Sammel-Indexzuordnung“ (start_index_map
in der formalen Beschreibung) und „Offset Mapping“ (remapped_offset_dims
in der formalen Beschreibung) zu [X
,0
] bzw. [0
,O
0
] erweitert, was bis zu [X
,O
0
] ergibt. Mit anderen Worten: Der Ausgabeindex [O
1
], der den Ausgabeindex [G
1
], der den Ausgabeindex [G
1
] zu [G
1
], der den Ausgabeindex [G
1
GatherIndices
aus [G
1
], {2, der Eingabe-Index [G
1
GatherIndices
ergibt, {2, der Eingabe-Index [G
1
], der den Ausgabeindex [G
1
GatherIndices
aus [G
1
], der Ausgabeindex [G
1
GatherIndices
ergibt, bevor sie als Indexe im Eingabearray verwendet werden, erweitert.0
0
0
0
0
O
tf.gather_nd
slice_sizes
ist für diesen Fall [1,11]
. Dies bedeutet, dass jeder Index-X
im Array der Indexe abrufen eine ganze Zeile auswählt und das Ergebnis die Verkettung all dieser Zeilen ist.
GetDimensionSize
Weitere Informationen finden Sie unter XlaBuilder::GetDimensionSize
.
Gibt die Größe der angegebenen Dimension des Operanden zurück. Der Operand muss Array-Form sein.
GetDimensionSize(operand, dimension)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
n-dimensionales Eingabearray |
dimension |
int64 |
Ein Wert im Intervall [0, n) , der die Dimension angibt |
SetDimensionSize
Weitere Informationen finden Sie unter XlaBuilder::SetDimensionSize
.
Legt die dynamische Größe der gegebenen XlaOp-Dimension fest. Der Operand muss Array-Form sein.
SetDimensionSize(operand, size, dimension)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
n-dimensionales Eingabearray. |
size |
XlaOp |
int32 steht für die dynamische Laufzeitgröße. |
dimension |
int64 |
Ein Wert im Intervall [0, n) , der die Dimension angibt. |
Sie leiten den Operanden als Ergebnis weiter, wobei die dynamische Dimension vom Compiler erfasst wird.
Aufgefüllte Werte werden von nachgelagerten Reduktionsvorgängen ignoriert.
let v: f32[10] = f32[10]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
let five: s32 = 5;
let six: s32 = 6;
// Setting dynamic dimension size doesn't change the upper bound of the static
// shape.
let padded_v_five: f32[10] = set_dimension_size(v, five, /*dimension=*/0);
let padded_v_six: f32[10] = set_dimension_size(v, six, /*dimension=*/0);
// sum == 1 + 2 + 3 + 4 + 5
let sum:f32[] = reduce_sum(padded_v_five);
// product == 1 * 2 * 3 * 4 * 5
let product:f32[] = reduce_product(padded_v_five);
// Changing padding size will yield different result.
// sum == 1 + 2 + 3 + 4 + 5 + 6
let sum:f32[] = reduce_sum(padded_v_six);
GetTupleElement
Weitere Informationen finden Sie unter XlaBuilder::GetTupleElement
.
Indexiert in ein Tupel mit einem Wert für die Compile-Zeitkonstante.
Der Wert muss eine Compile-Zeitkonstante sein, damit die Forminferenz den Typ des Ergebniswerts bestimmen kann.
Dies entspricht std::get<int N>(t)
in C++. Konzeptionell:
let v: f32[10] = f32[10]{0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
let s: s32 = 5;
let t: (f32[10], s32) = tuple(v, s);
let element_1: s32 = gettupleelement(t, 1); // Inferred shape matches s32.
Weitere Informationen finden Sie unter tf.tuple
.
Einspeisung
Weitere Informationen finden Sie unter XlaBuilder::Infeed
.
Infeed(shape)
Argument | Typ | Semantik |
---|---|---|
shape |
Shape |
Form der aus der InFeed-Oberfläche gelesenen Daten. Das Layoutfeld der Form muss so festgelegt werden, dass es dem Layout der an das Gerät gesendeten Daten entspricht. Andernfalls ist sein Verhalten nicht definiert. |
Liest ein einzelnes Datenelement aus der impliziten Infeed-Streaming-Schnittstelle des Geräts, interpretiert die Daten als gegebene Form und ihr Layout und gibt eine XlaOp
der Daten zurück. In einer Berechnung sind mehrere Infeed-Vorgänge zulässig, unter den Infeed-Vorgängen muss jedoch eine Gesamtreihenfolge vorhanden sein. Beispiel: Zwei Infeeds im folgenden Code haben eine Gesamtreihenfolge, da zwischen den while-Schleifen eine Abhängigkeit besteht.
result1 = while (condition, init = init_value) {
Infeed(shape)
}
result2 = while (condition, init = result1) {
Infeed(shape)
}
Verschachtelte Tupelformen werden nicht unterstützt. Bei einer leeren Tupelform ist der Infeed-Vorgang praktisch ein managementfrei und wird fortgesetzt, ohne Daten aus dem Infeed des Geräts zu lesen.
Iota
Weitere Informationen finden Sie unter XlaBuilder::Iota
.
Iota(shape, iota_dimension)
Erstellt ein konstantes Literal auf dem Gerät anstelle einer potenziell umfangreichen Hostübertragung. Erstellt ein Array mit einer angegebenen Form und Werten, die bei null beginnen und entlang der angegebenen Dimension um 1 erhöht werden. Bei Gleitkommatypen entspricht das erstellte Array ConvertElementType(Iota(...))
, wobei Iota
ein ganzzahliger Typ und die Konvertierung in den Gleitkommatyp ist.
Argumente | Typ | Semantik |
---|---|---|
shape |
Shape |
Form des durch Iota() erstellten Arrays |
iota_dimension |
int64 |
Die Dimension, die erhöht werden soll. |
Beispiel: Iota(s32[4, 8], 0)
gibt Folgendes zurück:
[[0, 0, 0, 0, 0, 0, 0, 0 ],
[1, 1, 1, 1, 1, 1, 1, 1 ],
[2, 2, 2, 2, 2, 2, 2, 2 ],
[3, 3, 3, 3, 3, 3, 3, 3 ]]
Iota(s32[4, 8], 1)
für Retouren
[[0, 1, 2, 3, 4, 5, 6, 7 ],
[0, 1, 2, 3, 4, 5, 6, 7 ],
[0, 1, 2, 3, 4, 5, 6, 7 ],
[0, 1, 2, 3, 4, 5, 6, 7 ]]
Karte
Weitere Informationen finden Sie unter XlaBuilder::Map
.
Map(operands..., computation)
Argumente | Typ | Semantik |
---|---|---|
operands |
Sequenz von N XlaOp s |
N Arrays der Typen T0..T{N-1} |
computation |
XlaComputation |
Berechnung des Typs T_0, T_1, .., T_{N + M -1} -> S mit N-Parametern vom Typ T und M beliebigen Typs |
dimensions |
int64 -Array |
Array mit Kartenmaßen |
Wendet eine Skalarfunktion auf die angegebenen operands
-Arrays an und erzeugt ein Array derselben Dimensionen, wobei jedes Element das Ergebnis der zugeordneten Funktion ist, die auf die entsprechenden Elemente in den Eingabearrays angewendet wurde.
Die zugeordnete Funktion ist eine beliebige Berechnung mit der Einschränkung, dass sie N Eingaben vom skalaren Typ T
und eine einzelne Ausgabe vom Typ S
hat. Die Ausgabe hat dieselben Abmessungen wie die Operanden, mit der Ausnahme, dass der Elementtyp T durch S ersetzt wird.
Beispiel: Map(op1, op2, op3, computation, par1)
ordnet jedem (mehrdimensionalen) Index in den Eingabearrays elem_out <-
computation(elem1, elem2, elem3, par1)
zu, um das Ausgabearray zu erzeugen.
OptimizationBarrier
Blockiert alle Optimierungsdurchläufe, sodass Berechnungen über die Barriere hinweg erfolgen.
Sorgt dafür, dass alle Eingaben vor Operatoren ausgewertet werden, die von den Ausgaben der Barriere abhängig sind.
Pad
Weitere Informationen finden Sie unter XlaBuilder::Pad
.
Pad(operand, padding_value, padding_config)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Array vom Typ T |
padding_value |
XlaOp |
Skalar vom Typ T , um den hinzugefügten Abstand auszufüllen |
padding_config |
PaddingConfig |
Abstand an beiden Rändern (niedrig, hoch) und zwischen den Elementen jeder Abmessung |
Erweitert das angegebene Array operand
durch Auffüllung um das Array sowie zwischen den Elementen des Arrays mit dem angegebenen padding_value
. padding_config
gibt den Rand- und den Innenabstand für jede Dimension an.
PaddingConfig
ist ein wiederkehrendes Feld von PaddingConfigDimension
, das drei Felder für jede Dimension enthält: edge_padding_low
, edge_padding_high
und interior_padding
.
edge_padding_low
und edge_padding_high
geben den Abstand an, der am unteren Ende (neben Index 0) bzw. am oberen Rand (neben dem höchsten Index) jeder Dimension hinzugefügt wird. Der Umfang des Randabstands kann negativ sein – der absolute Wert gibt die Anzahl der Elemente an, die aus der angegebenen Dimension entfernt werden sollen.
interior_padding
gibt den Abstand zwischen zwei Elementen in jeder Dimension an. Er darf nicht negativ sein. Das Innen-Padding erfolgt logisch vor dem Rand-Padding. Bei einem negativen Rand-Padding werden also Elemente aus dem Operanden mit Auffüllung entfernt.
Dieser Vorgang ist ein Leerbefehl, wenn die Paare zum Auffüllen der Ränder alle (0, 0) und die Werte des Innenrands 0 sind. Die folgende Abbildung zeigt Beispiele für verschiedene edge_padding
- und interior_padding
-Werte für ein zweidimensionales Array.
Recv
Weitere Informationen finden Sie unter XlaBuilder::Recv
.
Recv(shape, channel_handle)
Argumente | Typ | Semantik |
---|---|---|
shape |
Shape |
Form der zu empfangenden Daten |
channel_handle |
ChannelHandle |
eindeutige Kennung für jedes Sende-/Empfangspaar |
Empfängt Daten mit der angegebenen Form von einer Send
-Anweisung in einer anderen Berechnung mit demselben Kanal-Handle. Gibt ein XlaOp für die empfangenen Daten zurück.
Die Client API des Recv
-Vorgangs stellt die synchrone Kommunikation dar.
Die Anweisung wird jedoch intern in zwei HLO-Anweisungen (Recv
und RecvDone
) aufgeteilt, um asynchrone Datenübertragungen zu ermöglichen. Weitere Informationen finden Sie unter HloInstruction::CreateRecv
und HloInstruction::CreateRecvDone
.
Recv(const Shape& shape, int64 channel_id)
Weist Ressourcen zu, die zum Empfangen von Daten aus einer Send
-Anweisung mit derselben channel_id erforderlich sind. Gibt einen Kontext für die zugewiesenen Ressourcen zurück, der von der folgenden RecvDone
-Anweisung verwendet wird, um auf den Abschluss der Datenübertragung zu warten. Der Kontext ist ein Tupel von {Empfangspuffer (Form), Anfragekennung (U32)} und kann nur von einem RecvDone
-Befehl verwendet werden.
RecvDone(HloInstruction context)
Wartet bei einem durch eine Recv
-Anweisung erstellten Kontext, bis die Datenübertragung abgeschlossen ist, und gibt die empfangenen Daten zurück.
Einschränken
Weitere Informationen finden Sie unter XlaBuilder::Reduce
.
Wendet eine Reduktionsfunktion gleichzeitig auf ein oder mehrere Arrays an.
Reduce(operands..., init_values..., computation, dimensions)
Argumente | Typ | Semantik |
---|---|---|
operands |
Sequenz von N XlaOp |
N Arrays vom Typ T_0, ..., T_{N-1} |
init_values |
Sequenz von N XlaOp |
N Skalare vom Typ T_0, ..., T_{N-1} . |
computation |
XlaComputation |
Berechnung des Typs T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}) . |
dimensions |
int64 -Array |
Ungeordnetes Array von Dimensionen, das reduziert werden soll. |
Wobei:
- N muss größer oder gleich 1 sein.
- Die Berechnung muss „ungefähr“ assoziativ sein (siehe unten).
- Alle Eingabearrays müssen dieselben Dimensionen haben.
- Alle Anfangswerte müssen unter
computation
eine Identität bilden. - Wenn
N = 1
, istCollate(T)
T
. - Bei
N > 1
istCollate(T_0, ..., T_{N-1})
ein Tupel vonN
-Elementen vom TypT
.
Durch diese Operation werden eine oder mehrere Dimensionen jedes Eingabearrays in Skalare reduziert.
Der Rang jedes zurückgegebenen Arrays ist rank(operand) - len(dimensions)
. Die Ausgabe des Vorgangs ist Collate(Q_0, ..., Q_N)
, wobei Q_i
ein Array vom Typ T_i
ist, dessen Abmessungen unten beschrieben werden.
Die Berechnung der Reduktion kann von unterschiedlichen Back-Ends neu zugeordnet werden. Dies kann zu numerischen Unterschieden führen, da einige Reduktionsfunktionen wie das Addieren nicht assoziativ für Gleitkommazahlen sind. Wenn der Datenbereich jedoch begrenzt ist, reicht das Hinzufügen von Gleitkommazahlen für die meisten praktischen Anwendungsfälle fast aus, um assoziativ zu sein.
Beispiele
Bei der Reduzierung einer Dimension in einem einzelnen 1D-Array mit den Werten [10, 11,
12, 13]
und der Reduktionsfunktion f
(dies ist computation
) kann dies wie folgt berechnet werden:
f(10, f(11, f(12, f(init_value, 13)))
Es gibt aber auch noch viele andere Möglichkeiten, z.B.
f(init_value, f(f(10, f(init_value, 11)), f(f(init_value, 12), f(init_value, 13))))
Im Folgenden finden Sie ein grobes Pseudocode-Beispiel dafür, wie die Reduktion implementiert werden kann, wobei die Addition als Reduktionsberechnung mit einem Anfangswert von 0 verwendet wird.
result_shape <- remove all dims in dimensions from operand_shape
# Iterate over all elements in result_shape. The number of r's here is equal
# to the rank of the result
for r0 in range(result_shape[0]), r1 in range(result_shape[1]), ...:
# Initialize this result element
result[r0, r1...] <- 0
# Iterate over all the reduction dimensions
for d0 in range(dimensions[0]), d1 in range(dimensions[1]), ...:
# Increment the result element with the value of the operand's element.
# The index of the operand's element is constructed from all ri's and di's
# in the right order (by construction ri's and di's together index over the
# whole operand shape).
result[r0, r1...] += operand[ri... di]
Hier ein Beispiel für die Reduzierung eines 2D-Arrays (Matrix). Die Form hat den Rang 2, Dimension 0 der Größe 2 und Dimension 1 der Größe 3:
Ergebnisse der Reduzierung von Dimensionen 0 oder 1 mit einer Funktion zum Hinzufügen:
Beide Reduktionsergebnisse sind 1D-Arrays. Das Diagramm zeigt eine als Spalte und eine als Zeile nur zur besseren Veranschaulichung.
Für ein komplexeres Beispiel sehen Sie hier ein 3D-Array. Ihr Rang ist 3, Dimension 0 von Größe 4, Dimension 1 von Größe 2 und Dimension 2 von Größe 3. Der Einfachheit halber werden die Werte 1 bis 6 über die Dimension 0 hinweg repliziert.
Ähnlich wie im 2D-Beispiel können wir nur eine Dimension reduzieren. Wenn wir beispielsweise die Dimension 0 reduzieren, erhalten wir ein Array vom Typ Rang 2, in dem alle Werte der Dimension 0 zu einem Skalar zusammengefasst wurden:
| 4 8 12 |
| 16 20 24 |
Wenn wir Dimension 2 reduzieren, erhalten wir auch ein Rang-2-Array, in dem alle Werte in Dimension 2 zu einem Skalar zusammengefasst wurden:
| 6 15 |
| 6 15 |
| 6 15 |
| 6 15 |
Die relative Reihenfolge zwischen den verbleibenden Dimensionen in der Eingabe bleibt in der Ausgabe erhalten, einigen Dimensionen werden jedoch möglicherweise neue Zahlen zugewiesen, da sich der Rang ändert.
Wir können auch mehrere Dimensionen reduzieren. Aus den Dimensionen 0 und 1, die addierend reduziert werden, ergibt sich das 1D-Array [20, 28, 36]
.
Wenn das 3D-Array über alle seine Abmessungen reduziert wird, ergibt sich der skalare 84
.
Variadische Reduzierung
Bei N > 1
ist die Reduce-Funktionsanwendung etwas komplexer, da sie gleichzeitig auf alle Eingaben angewendet wird. Die Operanden werden in der folgenden Reihenfolge an die Berechnung übergeben:
- Laufender reduzierter Wert für ersten Operanden
- ...
- Aktiver reduzierter Wert für den N-ten Operanden
- Eingabewert für den ersten Operanden
- ...
- Eingabewert für den n-ten Operanden
Betrachten Sie beispielsweise die folgende Reduktionsfunktion, mit der der Maximalwert und der Wert argmax eines 1D-Arrays parallel berechnet werden können:
f: (Float, Int, Float, Int) -> Float, Int
f(max, argmax, value, index):
if value >= max:
return (value, index)
else:
return (max, argmax)
Bei 1D-Eingabearrays V = Float[N], K = Int[N]
und Initialisierungswerten I_V = Float, I_K = Int
entspricht das Ergebnis f_(N-1)
der Reduzierung für die einzige Eingabedimension der folgenden rekursiven Anwendung:
f_0 = f(I_V, I_K, V_0, K_0)
f_1 = f(f_0.first, f_0.second, V_1, K_1)
...
f_(N-1) = f(f_(N-2).first, f_(N-2).second, V_(N-1), K_(N-1))
Wenn diese Reduzierung auf ein Array von Werten und ein Array mit sequenziellen Indizes (d.h. iota) angewendet wird, erfolgt die Zusammenarbeit über die Arrays und gibt ein Tupel zurück, das den höchsten Wert und den übereinstimmenden Index enthält.
ReducePrecision
Weitere Informationen finden Sie unter XlaBuilder::ReducePrecision
.
Modelliert die Auswirkungen der Umwandlung von Gleitkommawerten in ein Format mit geringerer Genauigkeit (z. B. IEEE-FP16) und wieder zurück in das Originalformat. Die Anzahl der Exponenten- und Mantissen-Bits im Format mit geringerer Genauigkeit kann beliebig angegeben werden, auch wenn möglicherweise nicht alle Bitgrößen in allen Hardwareimplementierungen unterstützt werden.
ReducePrecision(operand, mantissa_bits, exponent_bits)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Array des Gleitkommatyps T |
exponent_bits |
int32 |
Anzahl der Exponentenbits im Format mit geringerer Genauigkeit |
mantissa_bits |
int32 |
Anzahl der Mantissen-Bits in einem Format mit geringerer Genauigkeit |
Das Ergebnis ist ein Array vom Typ T
. Die Eingabewerte werden auf den nächsten Wert gerundet, der mit der angegebenen Anzahl von Mantissenbits dargestellt werden kann (mithilfe der Semantik "Gleichheitsgrad" verwendet wird). Alle Werte, die den durch die Anzahl der Exponentenbits angegebenen Bereich überschreiten, werden auf positive oder negative unendliche Werte gesetzt. NaN
-Werte bleiben erhalten, können aber in kanonische NaN
-Werte konvertiert werden.
Das Format mit niedrigerer Genauigkeit muss mindestens ein Exponentenbit enthalten und eine nicht negative Anzahl von Mantissenbits haben, um einen Nullwert von einem unendlichen Wert zu unterscheiden. Die Anzahl der Exponenten- oder Mantissen-Bits kann den entsprechenden Wert für den Typ T
überschreiten. Der entsprechende Teil der Konvertierung ist dann eine Nulloperation.
ReduceScatter
Weitere Informationen finden Sie unter XlaBuilder::ReduceScatter
.
ReduceScatter ist ein kollektiver Vorgang, bei dem AllReduce effektiv ausgeführt und das Ergebnis dann in shard_count
-Blöcke entlang des scatter_dimension
aufgeteilt wird. Das Replikat i
in der Replikatgruppe erhält den ith
-Shard.
ReduceScatter(operand, computation, scatter_dim, shard_count,
replica_group_ids, channel_id)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Array oder ein nicht leeres Tupel von Arrays, die in allen Replikaten reduziert werden sollen. |
computation |
XlaComputation |
Reduktionsberechnung |
scatter_dimension |
int64 |
Zu streuende Dimension. |
shard_count |
int64 |
Anzahl der zu teilenden Blöcke: scatter_dimension |
replica_groups |
Vektor der Vektoren von int64 |
Gruppen, zwischen denen die Reduzierungen durchgeführt werden |
channel_id |
int64 optional |
Optionale Kanal-ID für die modulübergreifende Kommunikation |
- Wenn
operand
ein Tupel von Arrays ist, wird die Streuung zum Reduzieren für jedes Element des Tupels ausgeführt. replica_groups
ist eine Liste von Replikatgruppen, zwischen denen die Reduzierung durchgeführt wird (Replikat-ID für das aktuelle Replikat kann mitReplicaId
abgerufen werden). Die Reihenfolge der Replikate in den einzelnen Gruppen bestimmt die Reihenfolge, in der das Ergebnis der Gesamtreduzierung verteilt wird.replica_groups
muss entweder leer sein (in diesem Fall gehören alle Replikate zu einer einzigen Gruppe) oder die Anzahl der Elemente muss mit der Anzahl der Replikate übereinstimmen. Wenn mehrere Replikatgruppen vorhanden sind, müssen alle dieselbe Größe haben. Beispiel:replica_groups = {0, 2}, {1, 3}
führt eine Reduzierung zwischen den Replikaten0
und2
sowie1
und3
durch und verteilt dann das Ergebnis.shard_count
ist die Größe der einzelnen Replikatgruppen. Wir benötigen dies, wennreplica_groups
leer ist. Wennreplica_groups
nicht leer ist, mussshard_count
der Größe der einzelnen Replikatgruppen entsprechen.channel_id
wird für die modulübergreifende Kommunikation verwendet: Nurreduce-scatter
-Vorgänge mit derselbenchannel_id
können miteinander kommunizieren.
Die Ausgabeform ist die Eingabeform mit scatter_dimension
shard_count
-mal kleiner. Wenn beispielsweise zwei Replikate vorhanden sind und der Operand auf den beiden Replikaten den Wert [1.0, 2.25]
bzw. [3.0, 5.25]
hat, ist der Ausgabewert dieses Vorgangs mit 0
für das erste Replikat [4.0]
und für das zweite Replikat [7.5]
.scatter_dim
ReduceWindow
Weitere Informationen finden Sie unter XlaBuilder::ReduceWindow
.
Wendet eine Reduktionsfunktion auf alle Elemente in jedem Fenster einer Sequenz von n mehrdimensionalen Arrays an, wodurch ein einzelnes oder ein Tupel von n mehrdimensionalen Arrays ausgegeben wird. Jedes Ausgabearray hat dieselbe Anzahl von Elementen wie die Anzahl der gültigen Positionen des Fensters. Eine Pooling-Ebene kann als ReduceWindow
ausgedrückt werden. Ähnlich wie bei Reduce
wird dem angewendeten computation
immer die init_values
auf der linken Seite übergeben.
ReduceWindow(operands..., init_values..., computation, window_dimensions,
window_strides, padding)
Argumente | Typ | Semantik |
---|---|---|
operands |
N XlaOps |
Eine Sequenz aus n mehrdimensionalen Arrays vom Typ T_0,..., T_{N-1} , von denen jedes die Grundfläche darstellt, auf der sich das Fenster befindet. |
init_values |
N XlaOps |
Die N Startwerte für die Reduktion, einer für jeden der N Operanden. Weitere Informationen finden Sie unter Reduzieren. |
computation |
XlaComputation |
Reduktionsfunktion vom Typ T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}) , die auf Elemente in jedem Fenster aller Eingabeoperanden angewendet wird. |
window_dimensions |
ArraySlice<int64> |
Array von Ganzzahlen für Fensterdimensionswerte |
window_strides |
ArraySlice<int64> |
Array von Ganzzahlen für Fensterschrittwerte |
base_dilations |
ArraySlice<int64> |
Array von Ganzzahlen für Werte der Basiserweiterung |
window_dilations |
ArraySlice<int64> |
Array von Ganzzahlen für Fensterdilatationswerte |
padding |
Padding |
Padding-Typ für Fenster (Padding::kSame, das so auffüllt, dass es dieselbe Ausgabeform wie die Eingabe hat, wenn der Schritt 1 ist, oder Padding::kValid, das kein Auffüllen verwendet und das Fenster „stoppt“, sobald es nicht mehr passt) |
Wobei:
- N muss größer oder gleich 1 sein.
- Alle Eingabearrays müssen dieselben Dimensionen haben.
- Wenn
N = 1
, istCollate(T)
T
. - Bei
N > 1
istCollate(T_0, ..., T_{N-1})
ein Tupel vonN
-Elementen vom Typ(T0,...T{N-1})
.
Der folgende Code und die Abbildung zeigen ein Beispiel für die Verwendung von ReduceWindow
. Die Eingabe ist eine Matrix der Größe [4x6] und sowohl window_dimensions als auch window_stride_dimensions sind [2x3].
// Create a computation for the reduction (maximum).
XlaComputation max;
{
XlaBuilder builder(client_, "max");
auto y = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "y");
auto x = builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "x");
builder.Max(y, x);
max = builder.Build().value();
}
// Create a ReduceWindow computation with the max reduction computation.
XlaBuilder builder(client_, "reduce_window_2x3");
auto shape = ShapeUtil::MakeShape(F32, {4, 6});
auto input = builder.Parameter(0, shape, "input");
builder.ReduceWindow(
input,
/*init_val=*/builder.ConstantLiteral(LiteralUtil::MinValue(F32)),
*max,
/*window_dimensions=*/{2, 3},
/*window_stride_dimensions=*/{2, 3},
Padding::kValid);
Der Schritt "1" in einer Dimension gibt an, dass die Position eines Fensters in der Dimension 1 Element vom benachbarten Fenster entfernt ist. Um anzugeben, dass sich keine Fenster überschneiden, sollte „window_stride_dimensions“ gleich „window_dimensions“ sein. Die folgende Abbildung zeigt die Verwendung von zwei verschiedenen Schrittwerten. Das Padding wird auf jede Dimension der Eingabe angewendet und die Berechnungen sind die gleichen, als ob die Eingabe mit den Dimensionen, die sie nach dem Padding hat, enthalten wäre.
Für ein nicht triviales Padding-Beispiel könnten Sie den Mindestwert für das Reduzieren des Fensters berechnen (der Anfangswert ist MAX_FLOAT
) mit der Dimension 3
und dem Schritt 2
über das Eingabearray [10000, 1000, 100, 10, 1]
. Das Auffüllen von kValid
berechnet Mindestwerte über zwei gültige Fenster: [10000, 1000, 100]
und [100, 10, 1]
, was zur Ausgabe [100, 1]
führt. Durch das Auffüllen von kSame
wird das Array zuerst so aufgefüllt, dass die Form nach dem Reduce-Window dieselbe ist wie die Eingabe für Schritt 1, indem anfängliche Elemente auf beiden Seiten hinzugefügt werden und [MAX_VALUE, 10000, 1000, 100, 10, 1,
MAX_VALUE]
abgerufen wird. Wenn Sie das Reduce-Window über das aufgefüllte Array hinweg ausführen, werden die drei Fenster [MAX_VALUE, 10000, 1000]
, [1000, 100, 10]
, [10, 1, MAX_VALUE]
und Renditen [1000, 10, 1]
ausgeführt.
Die Auswertungsreihenfolge der Reduktionsfunktion ist beliebig und kann nicht deterministisch sein. Daher sollte die Reduktionsfunktion nicht allzu empfindlich auf Neuzuordnungen reagieren. Weitere Informationen finden Sie in der Diskussion über Assoziativität im Kontext von Reduce
.
ReplicaId
Weitere Informationen finden Sie unter XlaBuilder::ReplicaId
.
Gibt die eindeutige ID (U32-Skalar) des Replikats zurück.
ReplicaId()
Die eindeutige ID jedes Replikats ist eine vorzeichenlose Ganzzahl im Intervall [0, N)
, wobei N
die Anzahl der Replikate ist. Da alle Replikate dasselbe Programm ausführen, gibt ein ReplicaId()
-Aufruf im Programm einen anderen Wert für jedes Replikat zurück.
Form ändern
Weitere Informationen finden Sie unter XlaBuilder::Reshape
und Collapse
.
Formt die Abmessungen eines Arrays in eine neue Konfiguration um.
Reshape(operand, new_sizes)
Reshape(operand, dimensions, new_sizes)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Array vom Typ T |
dimensions |
Vektor für int64 |
Reihenfolge, in der Dimensionen minimiert werden |
new_sizes |
Vektor für int64 |
Vektor mit Größen neuer Dimensionen |
Konzeptionell wird ein Array bei der Umformung zuerst zu einem eindimensionalen Vektor von Datenwerten vereinfacht und dann in einer neuen Form verfeinert. Die Eingabeargumente sind ein beliebiges Array vom Typ T, ein Vektor der Compile-Zeitkonstante von Dimensionsindexen und ein Compile-Zeitkonstanten-Vektor der Dimensionsgrößen für das Ergebnis.
Die Werte im Vektor dimension
, falls angegeben, müssen eine Permutation aller T-Dimensionen sein. Wenn nicht angegeben, muss der Standardwert {0, ..., rank - 1}
sein. Die Dimensionen in dimensions
werden in der Schleifenverschachtelung von der langsamsten Dimension (größte große) bis zur am schnellsten variierenden Dimension (größtenteils) in der Schleifenverschachtelung sortiert, wodurch das Eingabearray in eine einzige Dimension minimiert wird. Der Vektor new_sizes
bestimmt die Größe des Ausgabearrays. Der Wert bei Index 0 in new_sizes
ist die Größe der Dimension 0, der Wert bei Index 1 ist die Größe von Dimension 1 usw. Das Produkt der Dimensionen new_size
muss dem Produkt der Dimensionsgrößen des Operanden entsprechen. Beim Verfeinern des minimierten Arrays in das mehrdimensionale Array, das durch new_sizes
definiert ist, werden die Dimensionen in new_sizes
von der langsamsten (größten) bis zur schnellsten Variante (sehr geringfügig) sortiert.
Nehmen wir z. B. „v“ als Array mit 24 Elementen:
let v = f32[4x2x3] { { {10, 11, 12}, {15, 16, 17} },
{ {20, 21, 22}, {25, 26, 27} },
{ {30, 31, 32}, {35, 36, 37} },
{ {40, 41, 42}, {45, 46, 47} } };
In-order collapse:
let v012_24 = Reshape(v, {0,1,2}, {24});
then v012_24 == f32[24] {10, 11, 12, 15, 16, 17, 20, 21, 22, 25, 26, 27,
30, 31, 32, 35, 36, 37, 40, 41, 42, 45, 46, 47};
let v012_83 = Reshape(v, {0,1,2}, {8,3});
then v012_83 == f32[8x3] { {10, 11, 12}, {15, 16, 17},
{20, 21, 22}, {25, 26, 27},
{30, 31, 32}, {35, 36, 37},
{40, 41, 42}, {45, 46, 47} };
Out-of-order collapse:
let v021_24 = Reshape(v, {1,2,0}, {24});
then v012_24 == f32[24] {10, 20, 30, 40, 11, 21, 31, 41, 12, 22, 32, 42,
15, 25, 35, 45, 16, 26, 36, 46, 17, 27, 37, 47};
let v021_83 = Reshape(v, {1,2,0}, {8,3});
then v021_83 == f32[8x3] { {10, 20, 30}, {40, 11, 21},
{31, 41, 12}, {22, 32, 42},
{15, 25, 35}, {45, 16, 26},
{36, 46, 17}, {27, 37, 47} };
let v021_262 = Reshape(v, {1,2,0}, {2,6,2});
then v021_262 == f32[2x6x2] { { {10, 20}, {30, 40},
{11, 21}, {31, 41},
{12, 22}, {32, 42} },
{ {15, 25}, {35, 45},
{16, 26}, {36, 46},
{17, 27}, {37, 47} } };
In einem Sonderfall kann die Umformung ein Array mit einem einzelnen Element in einen Skalar umwandeln und umgekehrt. Beispiel:
Reshape(f32[1x1] { {5} }, {0,1}, {}) == 5;
Reshape(5, {}, {1,1}) == f32[1x1] { {5} };
Umsatz (umgekehrt)
Weitere Informationen finden Sie unter XlaBuilder::Rev
.
Rev(operand, dimensions)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Array vom Typ T |
dimensions |
ArraySlice<int64> |
Umzukehrende Dimensionen |
Kehrt die Reihenfolge der Elemente im Array operand
entlang des angegebenen dimensions
um und generiert ein Ausgabearray derselben Form. Jedes Element des Operandenarrays bei einem mehrdimensionalen Index wird im Ausgabearray in einem transformierten Index gespeichert. Der mehrdimensionale Index wird transformiert, indem der Index in jeder Dimension umgedreht wird. Wenn also eine Dimension der Größe N eine der Umkehrdimensionen ist, wird ihr Index i in N - 1 - i transformiert.
Eine Verwendung der Rev
-Operation besteht darin, das Faltungsgewichtsarray entlang der beiden Fensterdimensionen während der Gradientenberechnung in neuronalen Netzwerken umzukehren.
RngNormal
Weitere Informationen finden Sie unter XlaBuilder::RngNormal
.
Erstellt eine Ausgabe einer bestimmten Form mit Zufallszahlen, die gemäß der \(N(\mu, \sigma)\) -Normalverteilung generiert werden. Die Parameter \(\mu\) und \(\sigma\)sowie die Ausgabeform müssen einen Elementtyp „Gleitkommazahl“ haben. Die Parameter müssen außerdem skalarwert sein.
RngNormal(mu, sigma, shape)
Argumente | Typ | Semantik |
---|---|---|
mu |
XlaOp |
Skalar des Typs T, der den Mittelwert der generierten Zahlen angibt |
sigma |
XlaOp |
Skalar vom Typ T, der die Standardabweichung der generierten |
shape |
Shape |
Ausgabeform vom Typ T |
RngUniform
Weitere Informationen finden Sie unter XlaBuilder::RngUniform
.
Erstellt eine Ausgabe einer bestimmten Form mit Zufallszahlen, die anhand der gleichmäßigen Verteilung über das Intervall \([a,b)\)generiert werden. Die Parameter und der Ausgabeelementtyp müssen boolesche, ganzzahlige oder Gleitkommatypen sein und die Typen müssen einheitlich sein. Die CPU- und GPU-Back-Ends unterstützen derzeit nur F64, F32, F16, BF16, S64, U64, S32 und U32. Außerdem müssen die Parameter skalarwert sein. Wenn \(b <= a\) , ist das Ergebnis implementierungsdefiniert.
RngUniform(a, b, shape)
Argumente | Typ | Semantik |
---|---|---|
a |
XlaOp |
Skalar vom Typ T, der untere Intervallgrenze angibt |
b |
XlaOp |
Skalar des Typs T, der den oberen Grenzwert des Intervalls angibt |
shape |
Shape |
Ausgabeform vom Typ T |
RngBitGenerator
Erzeugt unter Verwendung des angegebenen Algorithmus (oder Back-End-Standards) eine Ausgabe mit einer bestimmten Form, die mit einheitlichen Zufallsbits gefüllt ist, und gibt einen aktualisierten Zustand (mit der gleichen Form wie der Anfangszustand) und die generierten Zufallsdaten zurück.
Der Anfangszustand ist der Anfangszustand der aktuellen Generierung von Zufallszahlen. Er, die erforderliche Form und gültige Werte hängen vom verwendeten Algorithmus ab.
Die Ausgabe ist garantiert eine deterministische Funktion des Anfangszustands, aber es ist nicht garantiert, dass sie zwischen Back-Ends und verschiedenen Compiler-Versionen deterministisch ist.
RngBitGenerator(algorithm, key, shape)
Argumente | Typ | Semantik |
---|---|---|
algorithm |
RandomAlgorithm |
Zu verwendender PRNG-Algorithmus. |
initial_state |
XlaOp |
Anfangsstatus für den PRNG-Algorithmus. |
shape |
Shape |
Ausgabeform für generierte Daten. |
Verfügbare Werte für algorithm
:
rng_default
: Back-End-spezifischer Algorithmus mit Back-End-spezifischen Formanforderungen.rng_three_fry
: Zählerbasierter PRNG-Algorithmus von ThreeFry. Dieinitial_state
-Form istu64[2]
mit beliebigen Werten. Salmon et al. SC: 2011. Parallele Zufallszahlen – so einfach wie 1, 2, 3.rng_philox
: Philox-Algorithmus zur parallelen Erzeugung von Zufallszahlen. Die Forminitial_state
istu64[3]
mit beliebigen Werten. Salmon et al. SC: 2011. Parallele Zufallszahlen – so einfach wie 1, 2, 3.
Streudiagramm
Der XLA-Streuvorgang generiert eine Folge von Ergebnissen, die den Werten des Eingabearrays operands
entsprechen, wobei mehrere Segmente (bei durch scatter_indices
angegebenen Indexen) mit der Reihenfolge der Werte in updates
mithilfe von update_computation
aktualisiert werden.
Weitere Informationen finden Sie unter XlaBuilder::Scatter
.
scatter(operands..., scatter_indices, updates..., update_computation,
index_vector_dim, update_window_dims, inserted_window_dims,
scatter_dims_to_operand_dims)
Argumente | Typ | Semantik |
---|---|---|
operands |
Sequenz von N XlaOp |
N Arrays vom Typ T_0, ..., T_N , in die gestreut werden sollen. |
scatter_indices |
XlaOp |
Array mit den Anfangsindexen der Segmente, auf die gestreut werden muss. |
updates |
Sequenz von N XlaOp |
N Arrays vom Typ T_0, ..., T_N updates[i] enthält die Werte, die für die Streuung von operands[i] verwendet werden müssen. |
update_computation |
XlaComputation |
Berechnung zum Kombinieren der vorhandenen Werte im Eingabearray und der Aktualisierungen während der Streuung. Diese Berechnung sollte vom Typ T_0, ..., T_N, T_0, ..., T_N -> Collate(T_0, ..., T_N) sein. |
index_vector_dim |
int64 |
Die Dimension in scatter_indices , die die Startindexe enthält. |
update_window_dims |
ArraySlice<int64> |
Die Abmessungen in der updates -Form, die Fensterabmessungen sind. |
inserted_window_dims |
ArraySlice<int64> |
Die Fensterabmessungen, die in die Form updates eingefügt werden müssen. |
scatter_dims_to_operand_dims |
ArraySlice<int64> |
Eine Dimensionszuordnung zwischen den Streuindizes und dem Operandenindexbereich Dieses Array wird als Zuordnung von i zu scatter_dims_to_operand_dims[i] interpretiert . Sie muss eine Eins-zu-eins-eins-zu-eins und eine Gesamtsumme darstellen. |
indices_are_sorted |
bool |
Gibt an, ob die Indizes garantiert vom Aufrufer sortiert werden. |
Wobei:
- N muss größer oder gleich 1 sein.
operands
[0
], ...,operands
[N-1
] müssen alle dieselbe Größe haben.updates
[0
], ...,updates
[N-1
] müssen alle dieselbe Größe haben.- Wenn
N = 1
, istCollate(T)
T
. - Bei
N > 1
istCollate(T_0, ..., T_N)
ein Tupel vonN
-Elementen vom TypT
.
Wenn index_vector_dim
gleich scatter_indices.rank
ist, betrachten wir scatter_indices
implizit als eine nachgestellte 1
-Dimension.
Wir definieren update_scatter_dims
vom Typ ArraySlice<int64>
in aufsteigender Reihenfolge als den Satz von Dimensionen in der Form updates
, die nicht update_window_dims
sind.
Für die Argumente der Streuung gelten folgende Einschränkungen:
Jedes
updates
-Array muss den Rangupdate_window_dims.size + scatter_indices.rank - 1
haben.Die Grenzen der Dimension
i
in jedemupdates
-Array müssen so aussehen:- Wenn
i
inupdate_window_dims
vorhanden ist (d. h. für einenk
-Wertupdate_window_dims
[k
]), darf die Grenze der Dimensioni
inupdates
die entsprechende Grenze vonoperand
nach Berücksichtigung voninserted_window_dims
nicht überschreiten. Beispiel:adjusted_window_bounds
[k
] enthält die Grenzen vonoperand
, wobei die Grenzen bei den Indexeninserted_window_dims
entfernt wurden.adjusted_window_bounds
- Wenn
i
inupdate_scatter_dims
vorhanden ist (d.h. für einenk
-Wertupdate_scatter_dims
[k
]), muss die Grenze der Dimensioni
inupdates
der entsprechenden Grenze vonscatter_indices
entsprechen, wobeiindex_vector_dim
übersprungen wird (d.h.scatter_indices.shape.dims
[k
], wennk
<index_vector_dim
und ansonstenscatter_indices.shape.dims
[k+1
]).
- Wenn
update_window_dims
muss in aufsteigender Reihenfolge angegeben werden und darf keine sich wiederholenden Dimensionsnummern enthalten. Außerdem muss der Wert im Bereich[0, updates.rank)
liegen.inserted_window_dims
muss in aufsteigender Reihenfolge angegeben werden und darf keine sich wiederholenden Dimensionsnummern enthalten. Außerdem muss der Wert im Bereich[0, operand.rank)
liegen.operand.rank
muss der Summe vonupdate_window_dims.size
undinserted_window_dims.size
entsprechen.scatter_dims_to_operand_dims.size
muss gleichscatter_indices.shape.dims
[index_vector_dim
] sein und seine Werte müssen im Bereich[0, operand.rank)
liegen.
Für einen bestimmten Index U
in jedem updates
-Array wird der entsprechende Index I
im entsprechenden operands
-Array, auf das diese Aktualisierung angewendet werden muss, so berechnet:
- Lassen Sie
G
= {U
[k
] fürk
inupdate_scatter_dims
}. Verwenden SieG
, um einen IndexvektorS
im Arrayscatter_indices
nachzuschlagen, sodassS
[i
] =scatter_indices
[Kombination(G
,i
)], wobei Kombinieren(A, b) b an den Positionenindex_vector_dim
in A einfügt. - Erstellen Sie mit
S
einen IndexS
in
inoperand
. Dazu teilen SieS
mit derscatter_dims_to_operand_dims
-Karte auf. Formeller:S
in
[scatter_dims_to_operand_dims
[k
]] =S
[k
], wennk
<scatter_dims_to_operand_dims.size
.S
Andernfallsin
[_
] =0
.
- Erstellen Sie einen Index
W
in
in jedemoperands
-Array, indem Sie die Indexe beiupdate_window_dims
inU
gemäßinserted_window_dims
verteilen. Formeller:W
in
[window_dims_to_operand_dims
(k
)] =U
[k
], wennk
inupdate_window_dims
enthalten ist, wobeiwindow_dims_to_operand_dims
die monotone Funktion mit der Domain [0
,update_window_dims.size
) und dem Bereich [0
,operand.rank
) \inserted_window_dims
ist. Wennupdate_window_dims.size
beispielsweise4
ist,operand.rank
6
undinserted_window_dims
{0
,2
} ist, dann istwindow_dims_to_operand_dims
{0
→1
,1
→3
,2
→4
,3
→5
}.W
Andernfallsin
[_
] =0
.
I
istW
in
+S
in
, wobei + für die elementweise Addition steht.
Zusammenfassend lässt sich die Streuoperation wie folgt definieren.
- Initialisieren Sie
output
mitoperands
, d.h. für alle IndexeJ
und für alle IndexeO
im Arrayoperands
[J
]:
output
[J
][O
] =operands
[J
][O
] - Für jeden Index
U
im Arrayupdates
[J
] und den entsprechenden IndexO
im Arrayoperand
[J
], wennO
ein gültiger Index füroutput
ist:
(output
[0
][O
], ...,output
[N-1
][O
]) =update_computation
(output
[0
][O
], ...,output
[N-1
][O
],updates
[0
][U
], ...,updates
[N-1
][U
])
Die Reihenfolge, in der Aktualisierungen angewendet werden, ist nicht deterministisch. Wenn also mehrere Indexe in updates
auf denselben Index in operands
verweisen, ist der entsprechende Wert in output
nicht deterministisch.
Der erste Parameter, der an update_computation
übergeben wird, ist immer der aktuelle Wert aus dem Array output
. Der zweite Parameter ist immer der Wert aus dem Array updates
. Dies ist vor allem dann wichtig, wenn update_computation
nicht kommutativ ist.
Wenn indices_are_sorted
auf „true“ gesetzt ist, kann XLA davon ausgehen, dass start_indices
(in aufsteigender start_index_map
-Reihenfolge) nach Nutzer sortiert sind. Andernfalls ist die Semantik die Implementierung definiert.
Informellerweise kann der Streuvorgang als Inverse des Sammelvorgangs angesehen werden, d.h., er aktualisiert die Elemente in der Eingabe, die vom entsprechenden Collect-Vorgang extrahiert werden.
Eine detaillierte informelle Beschreibung und Beispiele finden Sie im Abschnitt „Informelle Description“ unter Gather
.
Auswählen
Weitere Informationen finden Sie unter XlaBuilder::Select
.
Erstellt ein Ausgabearray aus Elementen zweier Eingabearrays, basierend auf den Werten eines Prädikatarrays.
Select(pred, on_true, on_false)
Argumente | Typ | Semantik |
---|---|---|
pred |
XlaOp |
Array vom Typ PRED |
on_true |
XlaOp |
Array vom Typ T |
on_false |
XlaOp |
Array vom Typ T |
Die Arrays on_true
und on_false
müssen dieselbe Form haben. Dies ist auch die Form
des Ausgabearrays. Das Array pred
muss dieselbe Dimensionalität wie on_true
und on_false
mit dem Elementtyp PRED
haben.
Für jedes Element P
von pred
wird das entsprechende Element des Ausgabearrays aus on_true
übernommen, wenn der Wert von P
true
ist, und aus on_false
, wenn der Wert von P
false
ist. Als eingeschränkte Form der Übertragung kann pred
ein Skalar vom Typ PRED
sein. In diesem Fall wird das Ausgabearray vollständig aus on_true
verwendet, wenn pred
gleich true
ist, und aus on_false
, wenn pred
gleich false
ist.
Beispiel mit einem nicht skalaren pred
:
let pred: PRED[4] = {true, false, false, true};
let v1: s32[4] = {1, 2, 3, 4};
let v2: s32[4] = {100, 200, 300, 400};
==>
Select(pred, v1, v2) = s32[4]{1, 200, 300, 4};
Beispiel mit dem skalaren pred
:
let pred: PRED = true;
let v1: s32[4] = {1, 2, 3, 4};
let v2: s32[4] = {100, 200, 300, 400};
==>
Select(pred, v1, v2) = s32[4]{1, 2, 3, 4};
Die Auswahl zwischen Tupeln wird unterstützt. Tupel gelten für diesen Zweck
als skalare Typen. Wenn on_true
und on_false
Tupel sind (die dieselbe Form haben müssen), muss pred
ein Skalar vom Typ PRED
sein.
SelectAndScatter
Weitere Informationen finden Sie unter XlaBuilder::SelectAndScatter
.
Diese Operation kann als zusammengesetzte Operation betrachtet werden, bei der zuerst ReduceWindow
für das operand
-Array berechnet wird, um ein Element aus jedem Fenster auszuwählen, und dann das source
-Array auf die Indizes der ausgewählten Elemente verteilt wird, um ein Ausgabearray mit derselben Form wie das Operandenarray zu erstellen. Die binäre select
-Funktion wird verwendet, um ein Element aus jedem Fenster auszuwählen, indem es auf jedes Fenster angewendet wird. Sie wird mit der Eigenschaft aufgerufen, dass der Indexvektor des ersten Parameters lexikografisch kleiner als der Indexvektor des zweiten Parameters ist. Die select
-Funktion gibt true
zurück, wenn der erste Parameter ausgewählt ist, und gibt false
zurück, wenn der zweite Parameter ausgewählt ist. Die Funktion muss die Transitivität enthalten, d. h., wenn select(a, b)
und select(b, c)
true
sind, dann ist select(a, c)
ebenfalls true
, sodass das ausgewählte Element nicht von der Reihenfolge der Elemente abhängt, die in einem bestimmten Fenster durchlaufen werden.
Die Funktion scatter
wird auf jeden ausgewählten Index im Ausgabearray angewendet. Sie benötigt zwei Skalarparameter:
- Aktueller Wert am ausgewählten Index im Ausgabearray
- Streuwert aus
source
, der für den ausgewählten Index gilt
Die Funktion kombiniert die beiden Parameter und gibt einen skalaren Wert zurück, mit dem der Wert am ausgewählten Index im Ausgabearray aktualisiert wird. Zu Beginn sind alle Indizes des Ausgabearrays auf init_value
gesetzt.
Das Ausgabearray hat dieselbe Form wie das Array operand
und das Array source
muss dieselbe Form haben, wie das Ergebnis nach Anwendung eines ReduceWindow
-Vorgangs auf das Array operand
. SelectAndScatter
kann verwendet werden, um die Gradientenwerte für eine Pooling-Ebene in einem neuronalen Netzwerk rückzupropagieren.
SelectAndScatter(operand, select, window_dimensions, window_strides,
padding, source, init_value, scatter)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Array vom Typ T, über das die Fenster gleiten |
select |
XlaComputation |
Binäre Berechnung vom Typ T, T -> PRED , die auf alle Elemente in jedem Fenster angewendet wird. Gibt true zurück, wenn der erste Parameter ausgewählt ist, und false , wenn der zweite Parameter ausgewählt ist. |
window_dimensions |
ArraySlice<int64> |
Array von Ganzzahlen für Fensterdimensionswerte |
window_strides |
ArraySlice<int64> |
Array von Ganzzahlen für Fensterschrittwerte |
padding |
Padding |
Padding-Typ für Fenster (Padding::kSame oder Padding::kValid) |
source |
XlaOp |
Array vom Typ T mit den zu streuenden Werten |
init_value |
XlaOp |
Skalarwert vom Typ T für den Anfangswert des Ausgabearrays |
scatter |
XlaComputation |
Binäre Berechnung vom Typ T, T -> T , um jedes Streuquellenelement mit seinem Zielelement anzuwenden |
Die folgende Abbildung zeigt Beispiele für die Verwendung von SelectAndScatter
, wobei die Funktion select
den maximalen Wert aus ihren Parametern berechnet. Wenn sich die Fenster überschneiden, wie in Abbildung (2) unten gezeigt, kann ein Index des Arrays operand
von verschiedenen Fenstern mehrmals ausgewählt werden. In der Abbildung wird das Element mit dem Wert 9 durch die beiden oberen Fenster (blau und rot) ausgewählt und die binäre Additionsfunktion scatter
erzeugt das Ausgabeelement des Werts 8 (2 + 6).
Die Auswertungsreihenfolge der scatter
-Funktion ist beliebig und kann nicht deterministisch sein. Daher sollte die Funktion scatter
nicht allzu empfindlich auf Neuverknüpfungen reagieren. Weitere Informationen finden Sie in der Diskussion über Assoziativität im Kontext von Reduce
.
Senden
Weitere Informationen finden Sie unter XlaBuilder::Send
.
Send(operand, channel_handle)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Zu sendende Daten (Array vom Typ T) |
channel_handle |
ChannelHandle |
eindeutige Kennung für jedes Sende-/Empfangspaar |
Sendet die angegebenen Operandendaten in einer anderen Berechnung mit demselben Kanal-Handle an einen Recv
-Befehl. Es werden keine Daten zurückgegeben.
Ähnlich wie beim Recv
-Vorgang stellt die Client-API des Send
-Vorgangs die synchrone Kommunikation dar und wird intern in zwei HLO-Anweisungen (Send
und SendDone
) aufgeteilt, um asynchrone Datenübertragungen zu ermöglichen. Weitere Informationen finden Sie unter HloInstruction::CreateSend
und HloInstruction::CreateSendDone
.
Send(HloInstruction operand, int64 channel_id)
Initiiert eine asynchrone Übertragung des Operanden an die Ressourcen, die durch den Befehl Recv
mit derselben Kanal-ID zugewiesen wurden. Gibt einen Kontext zurück, der von der folgenden SendDone
-Anweisung verwendet wird, um auf den Abschluss der Datenübertragung zu warten. Der Kontext ist ein Tupel von {operand (shape), Request Identifier (U32)} und kann nur von einem SendDone
-Befehl verwendet werden.
SendDone(HloInstruction context)
Bei einem mit einer Send
-Anweisung erstellten Kontext wird auf den Abschluss der Datenübertragung gewartet. Die Anweisung gibt keine Daten zurück.
Planung der Kanalanleitung
Die Ausführungsreihenfolge der 4 Anweisungen für jeden Kanal (Recv
, RecvDone
, Send
, SendDone
) sieht so aus:
Recv
findet vor demSend
stattSend
findet vor demRecvDone
stattRecv
findet vor demRecvDone
stattSend
findet vor demSendDone
statt
Wenn die Back-End-Compiler einen linearen Zeitplan für jede Berechnung generieren, die über Kanalanweisungen kommuniziert, darf es keine Zyklen zwischen den Berechnungen geben. Die folgenden Zeitpläne führen beispielsweise zu Deadlocks.
Slice
Weitere Informationen finden Sie unter XlaBuilder::Slice
.
Beim Aufteilen wird ein Teilarray aus dem Eingabearray extrahiert. Das Subarray hat den gleichen Rang wie die Eingabe und enthält die Werte innerhalb eines Begrenzungsrahmens innerhalb des Eingabearrays, wobei die Dimensionen und Indexe des Begrenzungsrahmens als Argumente für den Aufteilungsvorgang angegeben werden.
Slice(operand, start_indices, limit_indices, strides)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
N-dimensionales Array des Typs T |
start_indices |
ArraySlice<int64> |
Liste von N Ganzzahlen, die die Startindexe des Segments für jede Dimension enthalten. Werte müssen größer oder gleich null sein. |
limit_indices |
ArraySlice<int64> |
Liste von N Ganzzahlen mit den Endindexen (ausschließlich) für das Segment für jede Dimension. Jeder Wert muss größer oder gleich dem jeweiligen start_indices -Wert für die Dimension und kleiner oder gleich der Größe der Dimension sein. |
strides |
ArraySlice<int64> |
Liste von N Ganzzahlen, die den Eingabeschritt des Segments bestimmt. Im Segment wird jedes strides[d] -Element der Dimension d ausgewählt. |
Beispiel für eindimensionales Format:
let a = {0.0, 1.0, 2.0, 3.0, 4.0}
Slice(a, {2}, {4}) produces:
{2.0, 3.0}
Beispiel für zweidimensionale Darstellung:
let b =
{ {0.0, 1.0, 2.0},
{3.0, 4.0, 5.0},
{6.0, 7.0, 8.0},
{9.0, 10.0, 11.0} }
Slice(b, {2, 1}, {4, 3}) produces:
{ { 7.0, 8.0},
{10.0, 11.0} }
Sortieren
Weitere Informationen finden Sie unter XlaBuilder::Sort
.
Sort(operands, comparator, dimension, is_stable)
Argumente | Typ | Semantik |
---|---|---|
operands |
ArraySlice<XlaOp> |
Die zu sortierenden Operanden |
comparator |
XlaComputation |
Die zu verwendende Vergleichsberechnung. |
dimension |
int64 |
Die Dimension, nach der sortiert werden soll. |
is_stable |
bool |
Gibt an, ob eine stabile Sortierung verwendet werden soll. |
Wenn nur ein Operand angegeben ist:
Wenn der Operand ein Rang-1-Tensor (ein Array) ist, ist das Ergebnis ein sortiertes Array. Wenn Sie das Array in aufsteigender Reihenfolge sortieren möchten, sollte der Vergleichsoperator einen Kleiner-als-Vergleich durchführen. Formell, nachdem das Array sortiert wurde, gilt es für alle Indexpositionen
i, j
miti < j
, die entwedercomparator(value[i], value[j]) = comparator(value[j], value[i]) = false
odercomparator(value[i], value[j]) = true
sind.Wenn der Operand einen höheren Rang hat, wird der Operand entlang der angegebenen Dimension sortiert. Bei einem Rang-2-Tensor (einer Matrix) wird beispielsweise mit dem Dimensionswert
0
jede Spalte unabhängig sortiert und bei einem Dimensionswert von1
jede Zeile unabhängig voneinander. Wenn keine Dimensionsnummer angegeben ist, wird standardmäßig die letzte Dimension ausgewählt. Für die sortierte Dimension gilt dieselbe Sortierreihenfolge wie für den Rang 1.
Wenn n > 1
-Operanden angegeben sind:
Alle
n
-Operanden müssen Tensoren mit denselben Dimensionen sein. Die Elementtypen der Tensoren können unterschiedlich sein.Alle Operanden werden zusammen, nicht einzeln sortiert. Die Operanden werden konzeptionell als Tupel behandelt. Bei der Prüfung, ob die Elemente jedes Operanden an den Indexpositionen
i
undj
vertauscht werden müssen, wird der Vergleichsoperator mit skalaren2 * n
-Parametern aufgerufen, wobei der Parameter2 * k
dem Wert an Positioni
aus demk-th
-Operanden und der Parameter2 * k + 1
dem Wert an Positionj
vomk-th
-Operanden entspricht. Normalerweise würde der Vergleichsoperator also die Parameter2 * k
und2 * k + 1
miteinander vergleichen und möglicherweise andere Parameterpaare als Stichpunkte verwenden.Das Ergebnis ist ein Tupel, das aus den Operanden in sortierter Reihenfolge (entlang der angegebenen Dimension, wie oben) besteht. Der
i-th
-Operand des Tupels entspricht demi-th
-Operanden von Sort.
Wenn es beispielsweise die drei Operanden operand0 = [3, 1]
, operand1 = [42, 50]
und operand2 = [-3.0, 1.1]
gibt und der Vergleichsoperator nur die Werte von operand0
mit einem Kleiner-als-Wert vergleicht, ist die Ausgabe der Sortierung das Tupel ([1, 3], [50, 42], [1.1, -3.0])
.
Wenn is_stable
auf „true“ gesetzt ist, ist die Sortierung garantiert stabil. Das heißt, wenn Elemente vorhanden sind, die vom Vergleichsoperator als gleich angesehen werden, wird die relative Reihenfolge der gleichen Werte beibehalten. Die zwei Elemente e1
und e2
sind genau dann gleich, wenn comparator(e1, e2) = comparator(e2, e1) = false
. Standardmäßig ist is_stable
auf "false" gesetzt.
Transponieren
Weitere Informationen finden Sie unter tf.reshape
.
Transpose(operand)
Argumente | Typ | Semantik |
---|---|---|
operand |
XlaOp |
Der Operand, der transponiert werden soll. |
permutation |
ArraySlice<int64> |
So passen Sie die Abmessungen an. |
Aktiviert die Operandendimensionen mit der angegebenen Permutation, sodass ∀ i . 0 ≤ i < rank ⇒ input_dimensions[permutation[i]] = output_dimensions[i]
.
Dies entspricht Reshape(operand, Permutation, Permute(permutation, operand.shape.dimensions)).
TriangularSolve
Weitere Informationen finden Sie unter XlaBuilder::TriangularSolve
.
Löst lineare Gleichungssysteme mit unteren oder oberen dreieckigen Koeffizientenmatrizen durch Vorwärts- oder RückwärtsSubstitution. Diese Routine wird entlang führender Dimensionen übertragen und löst eines der Matrixsysteme op(a) * x =
b
oder x * op(a) = b
für die Variable x
unter der Voraussetzung von a
und b
, wobei op(a)
entweder op(a) = a
, op(a) = Transpose(a)
oder op(a) = Conj(Transpose(a))
ist.
TriangularSolve(a, b, left_side, lower, unit_diagonal, transpose_a)
Argumente | Typ | Semantik |
---|---|---|
a |
XlaOp |
Ein Array aus dem Rang > 2 eines komplexen oder Gleitkommatyps mit der Form [..., M, M] . |
b |
XlaOp |
ein Array mit Rang > 2 desselben Typs mit der Form [..., M, K] , wenn left_side wahr ist, sonst [..., K, M] . |
left_side |
bool |
gibt an, ob ein System der Form op(a) * x = b (true ) oder x * op(a) = b (false ) gelöst werden soll. |
lower |
bool |
Gibt an, ob das obere oder untere Dreieck von a verwendet werden soll. |
unit_diagonal |
bool |
bei true wird angenommen, dass die diagonalen Elemente von a 1 sind und es nicht darauf zugegriffen wird. |
transpose_a |
Transpose |
Ob a in der vorliegenden Form verwendet werden soll oder ob das Transponieren oder Transponieren verwendet werden soll. |
Eingabedaten werden je nach Wert von lower
nur aus dem unteren/oberen Dreieck von a
gelesen. Werte aus dem anderen Dreieck werden ignoriert. Ausgabedaten werden im selben Dreieck zurückgegeben. Die Werte im anderen Dreieck sind implementierungsdefiniert und können beliebig sein.
Wenn der Rang von a
und b
größer als 2 ist, werden sie als Batches von Matrizen behandelt, wobei alle mit Ausnahme der beiden Dimensionen Batchdimensionen sind. a
und b
müssen dieselben Batchdimensionen haben.
Tupel
Weitere Informationen finden Sie unter XlaBuilder::Tuple
.
Ein Tupel mit einer variablen Anzahl von Daten-Handles, von denen jedes eine eigene Form hat.
Dies entspricht std::tuple
in C++. Konzeptionell:
let v: f32[10] = f32[10]{0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
let s: s32 = 5;
let t: (f32[10], s32) = tuple(v, s);
Tupel können über den GetTupleElement
-Vorgang zerlegt (aufgerufen) werden.
Während
Weitere Informationen finden Sie unter XlaBuilder::While
.
While(condition, body, init)
Argumente | Typ | Semantik |
---|---|---|
condition |
XlaComputation |
XlaComputation vom Typ T -> PRED , die die Beendigungsbedingung der Schleife definiert. |
body |
XlaComputation |
XlaComputation vom Typ T -> T , die den Hauptteil der Schleife definiert. |
init |
T |
Anfangswert für den Parameter condition und body . |
Führt body
sequenziell aus, bis condition
fehlschlägt. Dies ähnelt einer typischen „Während“-Schleife in vielen anderen Sprachen, mit Ausnahme der unten aufgeführten Unterschiede und Einschränkungen.
- Ein
While
-Knoten gibt einen Wert vom TypT
zurück, der das Ergebnis der letzten Ausführung vonbody
ist. - Die Form des Typs
T
wird statisch bestimmt und muss in allen Iterationen gleich sein.
Die T-Parameter der Berechnungen werden im ersten Durchlauf mit dem Wert init
initialisiert und bei jedem nachfolgenden Durchlauf automatisch auf das neue Ergebnis von body
aktualisiert.
Ein Hauptanwendungsfall des While
-Knotens besteht darin, die wiederholte Ausführung des Trainings in neuronalen Netzwerken zu implementieren. Unten ist ein vereinfachter Pseudocode mit einem Diagramm dargestellt, das die Berechnung darstellt. Den Code finden Sie in while_test.cc
.
Der Typ T
in diesem Beispiel ist eine Tuple
, die aus einem int32
für die Iterationsanzahl und einem vector[10]
für den Akkumulator besteht. Für 1.000 Iterationen fügt die Schleife dem Akkumulator weiterhin einen konstanten Vektor hinzu.
// Pseudocode for the computation.
init = {0, zero_vector[10]} // Tuple of int32 and float[10].
result = init;
while (result(0) < 1000) {
iteration = result(0) + 1;
new_vector = result(1) + constant_vector[10];
result = {iteration, new_vector};
}