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