Vous trouverez ci-dessous la description de la sémantique des opérations définies dans l'interface XlaBuilder
. En règle générale, ces opérations sont mappées une à une aux opérations définies dans l'interface RPC de la tâche xla_data.proto
.
Remarque sur la nomenclature: le type de données généralisé que XLA traite est un tableau à N dimensions contenant des éléments d'un type uniforme (par exemple, float 32 bits). Dans la documentation, array est utilisé pour désigner un tableau à dimensions arbitraires. Pour plus de commodité, les cas particuliers portent des noms plus spécifiques et familiers. Par exemple, un vecteur est un tableau unidimensionnel et une matrice est un tableau bidimensionnel.
AfterAll
Consultez également XlaBuilder::AfterAll
.
AfterAll prend un nombre variable de jetons et produit un seul jeton. Les jetons sont des types primitifs qui peuvent être regroupés entre des opérations ayant des effets secondaires pour appliquer l'ordonnancement. AfterAll
peut être utilisé comme jointure de jetons pour ordonner une opération après un ensemble d'opérations.
AfterAll(operands)
Arguments | Type | Sémantique |
---|---|---|
operands |
XlaOp |
nombre variable de jetons |
AllGather
Consultez également XlaBuilder::AllGather
.
Effectue une concaténation entre les instances répliquées.
AllGather(operand, all_gather_dim, shard_count, replica_group_ids,
channel_id)
Arguments | Type | Sémantique |
---|---|---|
operand
|
XlaOp
|
Tableau à concaténer entre plusieurs instances répliquées |
all_gather_dim |
int64 |
Dimension de concaténation |
replica_groups
|
vecteur de vecteurs de
int64 |
Groupes entre lesquels la concaténation est effectuée |
channel_id
|
(Facultatif) int64
|
ID de canal facultatif pour la communication entre les modules |
replica_groups
est une liste de groupes d'instances répliquées entre lesquels la concaténation est effectuée (l'ID de l'instance répliquée actuelle peut être récupéré à l'aide deReplicaId
). L'ordre des instances répliquées dans chaque groupe détermine l'ordre dans lequel leurs entrées se trouvent dans le résultat.replica_groups
doit être vide (auquel cas toutes les instances répliquées appartiennent à un seul groupe, classées de0
àN - 1
), ou contiennent le même nombre d'éléments que le nombre d'instances répliquées. Par exemple,replica_groups = {0, 2}, {1, 3}
effectue une concaténation entre les instances répliquées0
et2
, ainsi qu'entre1
et3
.shard_count
est la taille de chaque groupe d'instances répliquées. Nous en avons besoin dans les cas oùreplica_groups
est vide.channel_id
est utilisé pour la communication entre les modules: seules les opérationsall-gather
ayant le mêmechannel_id
peuvent communiquer entre elles.
La forme de sortie est celle de l'entrée, et all_gather_dim
agrandit shard_count
fois. Par exemple, s'il y a deux instances répliquées et que l'opérande a les valeurs [1.0, 2.5]
et [3.0, 5.25]
respectivement sur les deux instances répliquées, la valeur de sortie de cette opération, où all_gather_dim
est 0
, sera [1.0, 2.5, 3.0,
5.25]
sur les deux instances répliquées.
AllReduce
Consultez également XlaBuilder::AllReduce
.
Effectue un calcul personnalisé sur plusieurs instances répliquées.
AllReduce(operand, computation, replica_group_ids, channel_id)
Arguments | Type | Sémantique |
---|---|---|
operand
|
XlaOp
|
ou un tuple de tableaux non vide à réduire pour le nombre d'instances répliquées |
computation |
XlaComputation |
Calcul de la réduction |
replica_groups
|
vecteur de vecteurs de
int64 |
Groupes entre lesquels les réductions sont effectuées |
channel_id
|
(Facultatif) int64
|
ID de canal facultatif pour la communication entre les modules |
- Lorsque
operand
est un tuple de tableaux, la réduction globale est effectuée sur chaque élément du tuple. replica_groups
est une liste de groupes d'instances répliquées entre lesquels la réduction s'effectue (l'ID de l'instance répliquée actuelle peut être récupéré à l'aide deReplicaId
).replica_groups
doit être vide (auquel cas toutes les instances répliquées appartiennent à un seul groupe) ou contenir le même nombre d'éléments que le nombre d'instances répliquées. Par exemple,replica_groups = {0, 2}, {1, 3}
effectue une réduction entre les instances répliquées0
et2
, et1
et3
.channel_id
est utilisé pour la communication entre les modules: seules les opérationsall-reduce
ayant le mêmechannel_id
peuvent communiquer entre elles.
La forme de sortie est identique à la forme d'entrée. Par exemple, s'il y a deux instances répliquées et que l'opérande a respectivement les valeurs [1.0, 2.5]
et [3.0, 5.25]
sur les deux instances répliquées, la valeur de sortie de ce calcul d'opération et de somme sera [4.0, 7.75]
sur les deux instances répliquées. Si l'entrée est un tuple, la sortie est également un tuple.
Pour calculer le résultat de AllReduce
, vous devez disposer d'une entrée de chaque instance répliquée. Par conséquent, si une instance répliquée exécute un nœud AllReduce
plus de fois qu'une autre, l'ancienne instance répliquée attendra indéfiniment. Étant donné que les instances répliquées exécutent toutes le même programme, il n'y a pas beaucoup de façons de le faire. Toutefois, cela est possible lorsque l'état d'une boucle "ben" dépend des données du flux d'entrée et que les données fournies provoquent plus d'itérations sur une instance répliquée que sur une autre.
AllToAll
Consultez également XlaBuilder::AllToAll
.
AllToAll est une opération collective qui envoie des données de tous les cœurs à tous les cœurs. Elle comporte deux phases:
- La phase de dispersion. Sur chaque cœur, l'opérande est divisé en
split_count
blocs le long desplit_dimensions
, et les blocs sont dispersés sur tous les cœurs. Par exemple, le n-ième bloc est envoyé au i-ième cœur. - La phase de collecte. Chaque cœur concatène les blocs reçus le long de
concat_dimension
.
Les cœurs participants peuvent être configurés comme suit:
replica_groups
: chaque ReplicaGroup contient une liste des ID d'instances dupliquées participant au calcul (l'ID d'instance répliquée de l'instance répliquée actuelle peut être récupéré à l'aide deReplicaId
). AllToAll est appliqué dans les sous-groupes dans l'ordre spécifié. Par exemple,replica_groups = { {1,2,3}, {4,5,0} }
signifie qu'un élément AllToAll sera appliqué dans les instances répliquées{1, 2, 3}
et lors de la phase de collecte, et que les blocs reçus seront concaténés dans le même ordre de 1, 2, 3. Ensuite, un autre élément AllToAll est appliqué entre les instances répliquées 4, 5 et 0, et l'ordre de concaténation est également 4, 5, 0. Sireplica_groups
est vide, toutes les instances répliquées appartiennent à un groupe, dans l'ordre de concaténation de leur apparence.
Conditions préalables :
- La taille de la dimension de l'opérande sur
split_dimension
est divisible parsplit_count
. - La forme de l'opérande n'est pas tuple.
AllToAll(operand, split_dimension, concat_dimension, split_count,
replica_groups)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau d'entrée à n dimensions |
split_dimension
|
int64
|
Une valeur de l'intervalle [0,
n) qui nomme la dimension le long de laquelle l'opérande est fractionné. |
concat_dimension
|
int64
|
Une valeur dans l'intervalle [0,
n) qui nomme la dimension le long de laquelle les blocs de fractionnement sont concaténés. |
split_count
|
int64
|
Nombre de cœurs qui participent à cette opération. Si replica_groups est vide, il doit s'agir du nombre d'instances répliquées. Sinon, il doit être égal au nombre d'instances répliquées dans chaque groupe. |
replica_groups
|
ReplicaGroup vecteur
|
Chaque groupe contient une liste d'ID d'instances répliquées. |
Vous trouverez ci-dessous un exemple d'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);
Dans cet exemple, quatre cœurs participent à l'alltoall. Sur chaque cœur, l'opérande est divisé en quatre parties le long de la dimension 0, de sorte que chaque partie ait la forme f32[4,4]. Les quatre parties sont dispersées sur tous les cœurs. Chaque cœur concatène ensuite les parties reçues selon la dimension 1, dans l'ordre des cœurs 0 à 4. Ainsi, la sortie sur chaque cœur a la forme f32[16,4].
BatchNormGrad
Consultez également XlaBuilder::BatchNormGrad
et l'article sur la normalisation des lots d'origine pour obtenir une description détaillée de l'algorithme.
Calcule les gradients d'une norme de lot.
BatchNormGrad(operand, scale, mean, variance, grad_output, epsilon, feature_index)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau à n dimensions à normaliser (x) |
scale |
XlaOp |
Tableau à une dimension (\(\gamma\)) |
mean |
XlaOp |
Tableau à une dimension (\(\mu\)) |
variance |
XlaOp |
Tableau à une dimension (\(\sigma^2\)) |
grad_output |
XlaOp |
Gradients transmis à BatchNormTraining (\(\nabla y\)) |
epsilon |
float |
Valeur Epsilon (\(\epsilon\)) |
feature_index |
int64 |
Index de la dimension de caractéristique dans operand |
Pour chaque caractéristique de la dimension de caractéristique (feature_index
est l'index de la dimension d'élément dans operand
), l'opération calcule les dégradés par rapport à operand
, offset
et scale
dans toutes les autres dimensions. feature_index
doit être un index valide pour la dimension de caractéristique dans operand
.
Les trois gradients sont définis par les formules suivantes (en supposant un tableau à quatre dimensions operand
avec l'indice de dimension de caractéristique l
, la taille de lot m
et les tailles spatiales w
et 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} \]
Les entrées mean
et variance
représentent les valeurs de moments dans les dimensions de lot et spatiales.
Le type de sortie est un tuple de trois poignées:
Sorties | Type | Sémantique |
---|---|---|
grad_operand
|
XlaOp
|
dégradé par rapport à l'entrée operand ($\nabla
x$) |
grad_scale
|
XlaOp
|
gradient par rapport à la valeur scale d'entrée ($\nabla
\gamma$) |
grad_offset
|
XlaOp
|
dégradé par rapport à l'entrée offset ($\nabla
\beta$) |
BatchNormInference
Consultez également XlaBuilder::BatchNormInference
et l'article sur la normalisation des lots d'origine pour obtenir une description détaillée de l'algorithme.
Normalise un tableau en fonction de ses dimensions de lot et spatiales.
BatchNormInference(operand, scale, offset, mean, variance, epsilon, feature_index)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau à n dimensions à normaliser |
scale |
XlaOp |
Tableau à une dimension |
offset |
XlaOp |
Tableau à une dimension |
mean |
XlaOp |
Tableau à une dimension |
variance |
XlaOp |
Tableau à une dimension |
epsilon |
float |
Valeur Epsilon |
feature_index |
int64 |
Index de la dimension de caractéristique dans operand |
Pour chaque caractéristique de la dimension de caractéristique (feature_index
est l'indice de la dimension de caractéristique dans operand
), l'opération calcule la moyenne et la variance entre toutes les autres dimensions, et utilise la moyenne et la variance pour normaliser chaque élément dans operand
. feature_index
doit être un index valide pour la dimension de caractéristique dans operand
.
BatchNormInference
équivaut à appeler BatchNormTraining
sans calculer mean
et variance
pour chaque lot. Elle utilise à la place les valeurs d'entrée mean
et variance
en tant que valeurs estimées. L'objectif de cette opération est de réduire la latence d'inférence, d'où le nom BatchNormInference
.
La sortie est un tableau normalisé à N dimensions ayant la même forme que la valeur operand
d'entrée.
BatchNormTraining
Consultez également XlaBuilder::BatchNormTraining
et the original batch normalization paper
pour obtenir une description détaillée de l'algorithme.
Normalise un tableau en fonction de ses dimensions de lot et spatiales.
BatchNormTraining(operand, scale, offset, epsilon, feature_index)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau à n dimensions à normaliser (x) |
scale |
XlaOp |
Tableau à une dimension (\(\gamma\)) |
offset |
XlaOp |
Tableau à une dimension (\(\beta\)) |
epsilon |
float |
Valeur Epsilon (\(\epsilon\)) |
feature_index |
int64 |
Index de la dimension de caractéristique dans operand |
Pour chaque caractéristique de la dimension de caractéristique (feature_index
est l'indice de la dimension de caractéristique dans operand
), l'opération calcule la moyenne et la variance entre toutes les autres dimensions, et utilise la moyenne et la variance pour normaliser chaque élément dans operand
. feature_index
doit être un index valide pour la dimension de caractéristique dans operand
.
L'algorithme fonctionne comme suit pour chaque lot dans operand
\(x\) contenant des éléments m
avec w
et h
comme taille des dimensions spatiales (en supposant que operand
est un tableau à 4 dimensions):
Calcule la moyenne par lot \(\mu_l\) pour chaque caractéristique
l
dans la dimension de caractéristique : \(\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}\)Calcule la variance par lot \(\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$
Normalisation, adaptation et décale : \(y_{ijkl}=\frac{\gamma_l(x_{ijkl}-\mu_l)}{\sqrt[2]{\sigma^2_l+\epsilon} }+\beta_l\)
La valeur epsilon, généralement un petit nombre, est ajoutée pour éviter les erreurs de division par zéro.
Le type de sortie est un tuple de trois valeurs XlaOp
:
Sorties | Type | Sémantique |
---|---|---|
output
|
XlaOp
|
Tableau à n dimensions ayant la même forme que l'entrée operand (y) |
batch_mean |
XlaOp |
Tableau à une dimension (\(\mu\)) |
batch_var |
XlaOp |
Tableau à une dimension (\(\sigma^2\)) |
batch_mean
et batch_var
sont des moments calculés pour les dimensions de lot et spatiales à l'aide des formules ci-dessus.
BitcastConvertType
Consultez également XlaBuilder::BitcastConvertType
.
À l'instar d'une tf.bitcast
dans TensorFlow, cette méthode effectue une opération de diffusion de bits par élément d'une forme de données à une forme cible. Les tailles d'entrée et de sortie doivent correspondre: par exemple, les éléments s32
deviennent des éléments f32
via une routine Bitcast, et un élément s32
devient quatre éléments s8
. Bitcast est implémenté en tant que diffusion de bas niveau. Par conséquent, les machines avec différentes représentations à virgule flottante génèrent des résultats différents.
BitcastConvertType(operand, new_element_type)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau de type T avec les valeurs D |
new_element_type |
PrimitiveType |
type U |
Les dimensions de l'opérande et de la forme cible doivent correspondre, à l'exception de la dernière dimension, qui sera modifiée par le ratio de la taille de la primitive avant et après la conversion.
Les types d'éléments source et de destination ne doivent pas être des tuples.
Conversion de bitmap en type primitif de largeur différente
L'instruction HLO BitcastConvert
est compatible avec les cas où la taille du type d'élément de sortie T'
n'est pas égale à la taille de l'élément d'entrée T
. Étant donné que l'ensemble de l'opération est conceptuellement un bitcast et ne modifie pas les octets sous-jacents, la forme de l'élément de sortie doit changer. Pour B = sizeof(T), B' =
sizeof(T')
, deux cas de figure sont possibles.
Tout d'abord, lorsque la valeur est B > B'
, la forme de sortie se voit attribuer une nouvelle dimension de plus petite taille, de taille B/B'
. Exemple :
f16[10,2]{1,0} %output = f16[10,2]{1,0} bitcast-convert(f32[10]{0} %input)
La règle reste la même pour les scalaires efficaces:
f16[2]{0} %output = f16[2]{0} bitcast-convert(f32[] %input)
Pour B' > B
, l'instruction nécessite que la dernière dimension logique de la forme d'entrée soit égale à B'/B
. Cette dimension est supprimée lors de la conversion:
f32[10]{0} %output = f32[10]{0} bitcast-convert(f16[10,2]{1,0} %input)
Notez que les conversions entre différentes largeurs de bits ne sont pas définies par élément.
Diffuser
Consultez également XlaBuilder::Broadcast
.
Ajoute des dimensions à un tableau en dupliquant les données qu'il contient.
Broadcast(operand, broadcast_sizes)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
Tableau à dupliquer |
broadcast_sizes |
ArraySlice<int64> |
Les tailles des nouvelles dimensions |
Les nouvelles dimensions sont insérées à gauche. Par exemple, si broadcast_sizes
a les valeurs {a0, ..., aN}
et que la forme de l'opérande a les dimensions {b0, ..., bM}
, la forme de la sortie a les dimensions {a0, ..., aN, b0, ..., bM}
.
Le nouvel index de dimensions se compose de copies de l'opérande.
output[i0, ..., iN, j0, ..., jM] = operand[j0, ..., jM]
Par exemple, si operand
est une f32
scalaire avec la valeur 2.0f
et que broadcast_sizes
est {2, 3}
, le résultat sera un tableau avec la forme f32[2, 3]
et toutes les valeurs dans le résultat seront 2.0f
.
BroadcastInDim
Consultez également XlaBuilder::BroadcastInDim
.
Développe la taille et le rang d'un tableau en dupliquant les données qu'il contient.
BroadcastInDim(operand, out_dim_size, broadcast_dimensions)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
Tableau à dupliquer |
out_dim_size |
ArraySlice<int64> |
Les tailles des dimensions de la forme cible |
broadcast_dimensions |
ArraySlice<int64> |
À quelle dimension de la forme cible correspond chaque dimension de la forme de l'opérande |
Semblable à Broadcast, mais permet d'ajouter des dimensions n'importe où et de développer des dimensions existantes de taille 1.
operand
est diffusé sur la forme décrite par out_dim_size
.
broadcast_dimensions
mappe les dimensions de operand
aux dimensions de la forme cible, c'est-à-dire que la dimension i'e de l'opérande est mappée à la dimension
broadcast_dimension[i]'e de la forme de sortie. Les dimensions de operand
doivent avoir la taille 1 ou être de la même taille que la dimension dans la forme de sortie à laquelle elles sont mappées. Les dimensions restantes sont remplies avec des dimensions de taille 1. La diffusion de dimensions dégénérées est ensuite diffusée le long de ces dimensions dégénérées pour atteindre la forme de sortie. La sémantique est décrite en détail sur la page de diffusion.
Call
Consultez également XlaBuilder::Call
.
Appelle un calcul avec les arguments indiqués.
Call(computation, args...)
Arguments | Type | Sémantique |
---|---|---|
computation |
XlaComputation |
un calcul de type T_0, T_1, ..., T_{N-1} -> S avec N paramètres de type arbitraire |
args |
séquence de N XlaOp |
N arguments de type arbitraire |
L'arité et les types de args
doivent correspondre aux paramètres de computation
. Aucun args
n'est autorisé.
Cholésky
Consultez également XlaBuilder::Cholesky
.
Calcule la décomposition de Cholésky d'un lot de matrices définies positives symétriques (hermitiennes).
Cholesky(a, lower)
Arguments | Type | Sémantique |
---|---|---|
a |
XlaOp |
un tableau rang > 2 d'un type complexe ou à virgule flottante. |
lower |
bool |
utiliser le triangle supérieur ou inférieur de a . |
Si lower
est true
, calcule les matrices triangulaires inférieures l
de sorte que $a = l .
l^T$. Si lower
est false
, calcule les matrices triangulaires supérieures u
de sorte que\(a = u^T . u\).
Les données d'entrée sont lues uniquement à partir du triangle inférieur/supérieur de a
, en fonction de la valeur de lower
. Les valeurs de l'autre triangle sont ignorées. Les données de sortie sont renvoyées dans le même triangle. Les valeurs de l'autre triangle sont définies par l'implémentation et peuvent être n'importe quelle valeur.
Si le rang de a
est supérieur à 2, a
est traité comme un lot de matrices, où toutes les dimensions, à l'exception des deux mineures, sont des dimensions de lot.
Si a
n'est pas symétrique (hermitien) positif, le résultat est défini par l'implémentation.
Pince
Consultez également XlaBuilder::Clamp
.
Fixe un opérande dans la plage comprise entre une valeur minimale et une valeur maximale.
Clamp(min, operand, max)
Arguments | Type | Sémantique |
---|---|---|
min |
XlaOp |
tableau de type T |
operand |
XlaOp |
tableau de type T |
max |
XlaOp |
tableau de type T |
En fonction d'un opérande et des valeurs minimale et maximale, renvoie l'opérande s'il se trouve dans la plage comprise entre les valeurs minimale et maximale, et la valeur minimale si l'opérande est inférieur à cette plage, ou la valeur maximale si l'opérande se trouve au-dessus de cette plage. Par exemple, clamp(a, x, b) = min(max(a, x), b)
.
Les trois tableaux doivent avoir la même forme. Par ailleurs, en tant que forme restreinte de diffusion, min
et/ou max
peuvent être une valeur scalaire de type T
.
Exemple avec min
et max
scalaires:
let operand: s32[3] = {-1, 5, 9};
let min: s32 = 0;
let max: s32 = 6;
==>
Clamp(min, operand, max) = s32[3]{0, 5, 6};
Réduire
Consultez également XlaBuilder::Collapse
et l'opération tf.reshape
.
Réduire les dimensions d'un tableau en une seule dimension.
Collapse(operand, dimensions)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau de type T |
dimensions |
int64 vecteur |
dans l'ordre et consécutif des dimensions de T. |
Le repli remplace le sous-ensemble donné de dimensions de l'opérande par une seule dimension. Les arguments d'entrée sont un tableau arbitraire de type T et un vecteur constant d'indices de dimension au moment de la compilation. Les index de dimension doivent être
un sous-ensemble consécutif de dimensions de T (dans l'ordre croissant). Ainsi, {0, 1, 2}, {0, 1} ou {1, 2} sont tous des ensembles de dimensions valides, mais {1, 0} ou {0, 2} ne le sont pas. Elles sont remplacées par une nouvelle dimension unique, située dans la même position dans l'ordre des dimensions que celles qu'elles remplacent, et dont la nouvelle dimension est égale au produit des tailles d'origine. Le nombre de dimension le plus bas dans dimensions
est la dimension la plus lente (la plus majeure) dans l'imbrication de boucle, qui réduit ces dimensions, et le nombre de dimension le plus élevé varie le plus rapidement (la plus mineure). Reportez-vous à l'opérateur tf.reshape
si un ordre de réduction plus général est nécessaire.
Par exemple, imaginons que "v" représente un tableau de 24 éléments:
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
Consultez également XlaBuilder::CollectivePermute
.
CollectivePermute est une opération collective qui envoie et reçoit des données répliquées entre elles.
CollectivePermute(operand, source_target_pairs)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau d'entrée à n dimensions |
source_target_pairs |
<int64, int64> vecteur |
Une liste des paires (source_replica_id, target_replica_id). Pour chaque paire, l'opérande est envoyé de l'instance répliquée source à l'instance répliquée cible. |
Notez que les restrictions suivantes s'appliquent à source_target_pair
:
- Deux paires ne doivent pas avoir le même ID de l'instance répliquée cible ni le même ID de l'instance répliquée source.
- Si un ID d'instance répliquée n'est une cible dans aucune paire, la sortie de cette instance répliquée est un Tensor composé de zéro(s) ayant la même forme que l'entrée.
Concatenate
Consultez également XlaBuilder::ConcatInDim
.
"Concaténer" compose un tableau à partir de plusieurs opérandes de tableau. Le tableau a le même rang que chacun des opérandes du tableau d'entrée (qui doivent avoir le même rang) et contient les arguments dans l'ordre dans lequel ils ont été spécifiés.
Concatenate(operands..., dimension)
Arguments | Type | Sémantique |
---|---|---|
operands |
séquence de N XlaOp |
N tableaux de type T ayant les dimensions [L0, L1, ...]. Nécessite N >= 1. |
dimension |
int64 |
Valeur de l'intervalle [0, N) qui nomme la dimension à concaténer entre l'élément operands . |
À l'exception de dimension
, toutes les dimensions doivent être identiques. En effet, XLA n'accepte pas les tableaux "irréguliers". Notez également que les valeurs RANK-0 ne peuvent pas être concaténées (car il est impossible de nommer la dimension en fonction de laquelle la concaténation se produit).
Exemple unidimensionnel:
Concat({ {2, 3}, {4, 5}, {6, 7} }, 0)
>>> {2, 3, 4, 5, 6, 7}
Exemple en deux dimensions:
let a = {
{1, 2},
{3, 4},
{5, 6},
};
let b = {
{7, 8},
};
Concat({a, b}, 0)
>>> {
{1, 2},
{3, 4},
{5, 6},
{7, 8},
}
Diagramme:
Conditional
Consultez également XlaBuilder::Conditional
.
Conditional(pred, true_operand, true_computation, false_operand,
false_computation)
Arguments | Type | Sémantique |
---|---|---|
pred |
XlaOp |
Scalaire de type PRED |
true_operand |
XlaOp |
Type d'argument \(T_0\) |
true_computation |
XlaComputation |
XlaComputation de type \(T_0 \to S\) |
false_operand |
XlaOp |
Type d'argument \(T_1\) |
false_computation |
XlaComputation |
XlaComputation de type \(T_1 \to S\) |
Exécute true_computation
si pred
est true
, false_computation
si pred
est false
, et renvoie le résultat.
true_computation
doit accepter un seul argument de type \(T_0\) et sera appelé avec true_operand
, qui doit être du même type. false_computation
doit accepter un seul argument de type \(T_1\) et sera appelé avec false_operand
, qui doit être du même type. Le type de la valeur renvoyée pour true_computation
et false_computation
doit être identique.
Notez que seul true_computation
et false_computation
seront exécutés en fonction de la valeur de pred
.
Conditional(branch_index, branch_computations, branch_operands)
Arguments | Type | Sémantique |
---|---|---|
branch_index |
XlaOp |
Scalaire de type S32 |
branch_computations |
séquence de N XlaComputation |
XlaComputations de type \(T_0 \to S , T_1 \to S , ..., T_{N-1} \to S\) |
branch_operands |
séquence de N XlaOp |
Arguments de type \(T_0 , T_1 , ..., T_{N-1}\) |
Exécute branch_computations[branch_index]
et renvoie le résultat. Si branch_index
est une S32
inférieure à 0 ou supérieure à N, branch_computations[N-1]
est exécuté en tant que branche par défaut.
Chaque branch_computations[b]
doit accepter un seul argument de type \(T_b\) et sera appelé avec branch_operands[b]
, qui doit être du même type. Le type de la valeur renvoyée pour chaque élément branch_computations[b]
doit être identique.
Notez qu'un seul des branch_computations
sera exécuté en fonction de la valeur de branch_index
.
Conv (convolution)
Consultez également XlaBuilder::Conv
.
En tant que ConvWithGeneralPadding, mais la marge intérieure est spécifiée de manière abrégée comme étant MÊME ou VALIDE. La MÊME marge intérieure remplit l'entrée (lhs
) avec des zéros afin que la sortie ait la même forme que l'entrée lorsqu'elle ne tient pas compte des pas. Un remplissage VALIDE signifie simplement qu'il n'y a pas de marge intérieure.
ConvWithGeneralPadding (convolution)
Consultez également XlaBuilder::ConvWithGeneralPadding
.
Calcule une convolution du type utilisé dans les réseaux de neurones. Ici, une convolution peut être considérée comme une fenêtre à N dimensions se déplaçant sur une zone de base à N dimensions, et un calcul est effectué pour chaque position possible de la fenêtre.
Arguments | Type | Sémantique |
---|---|---|
lhs |
XlaOp |
tableau d'entrées n+2 de rang |
rhs |
XlaOp |
rang n+2 du tableau des pondérations du noyau |
window_strides |
ArraySlice<int64> |
Tableau n-d de progressions du noyau |
padding |
ArraySlice< pair<int64,int64>> |
Tableau n-d de marge intérieure (faible ou élevée) |
lhs_dilation |
ArraySlice<int64> |
Tableau des facteurs de dilatation de n-d lhs |
rhs_dilation |
ArraySlice<int64> |
Tableau des facteurs de dilatation de n-d rhs |
feature_group_count |
int64 | le nombre de groupes de caractéristiques |
batch_group_count |
int64 | le nombre de groupes de lots |
Soit n le nombre de dimensions spatiales. L'argument lhs
est un tableau de rang n+2 décrivant l'aire de la base. C'est ce qu'on appelle l'entrée, même si
bien sûr le rhs est aussi une entrée. Dans un réseau de neurones, il s'agit des activations d'entrée.
Les dimensions n+2 sont, dans cet ordre:
batch
: chaque coordonnée de cette dimension représente une entrée indépendante pour laquelle une convolution est effectuée.z/depth/features
: chaque position (y,x) dans la zone de base est associée à un vecteur, qui entre dans cette dimension.spatial_dims
: décrit les dimensions spatialesn
qui définissent la zone de base sur laquelle se déplace la fenêtre.
L'argument rhs
est un tableau de rang n+2 décrivant le filtre, le noyau ou la fenêtre convolutifs. Les dimensions sont les suivantes, dans cet ordre:
output-z
: dimensionz
du résultat.input-z
: la taille de cette dimension multipliée parfeature_group_count
doit être égale à la taille de la dimensionz
en lhs.spatial_dims
: décrit les dimensions spatialesn
qui définissent la fenêtre n-d qui se déplace dans la zone de base.
L'argument window_strides
spécifie la foulée de la fenêtre convolutive dans les dimensions spatiales. Par exemple, si la foulée de la première dimension spatiale est de 3, la fenêtre ne peut être placée qu'aux coordonnées où le premier index spatial est divisible par 3.
L'argument padding
spécifie la quantité de marge intérieure nulle à appliquer à la zone de base. La quantité de marge intérieure peut être négative. La valeur absolue de la marge intérieure négative indique le nombre d'éléments à supprimer de la dimension spécifiée avant d'effectuer la convolution. padding[0]
spécifie la marge intérieure pour la dimension y
et padding[1]
spécifie la marge intérieure pour la dimension x
. Chaque paire utilise la marge intérieure faible comme premier élément et la marge intérieure élevée comme deuxième élément. La marge intérieure faible est appliquée dans la direction des index inférieurs, tandis que la marge intérieure élevée est appliquée dans la direction des index plus élevés. Par exemple, si padding[1]
est défini sur (2,3)
, il y aura une marge intérieure de 2 zéros à gauche et de 3 zéros à droite dans la deuxième dimension spatiale. L'utilisation du remplissage équivaut à insérer ces mêmes valeurs nulles dans l'entrée (lhs
) avant d'effectuer la convolution.
Les arguments lhs_dilation
et rhs_dilation
spécifient le facteur de dilation à appliquer respectivement aux valeurs lhs et rhs dans chaque dimension spatiale. Si le facteur de dilatation d'une dimension spatiale est "d", des trous d-1 sont implicitement placés entre chacune des entrées de cette dimension, ce qui augmente la taille du tableau. Les trous sont remplis avec une valeur no-op, ce qui pour la convolution correspond à des zéros.
La dilatation des rhs est également appelée convolution atrouille. Pour en savoir plus, consultez tf.nn.atrous_conv2d
. La dilatation des lhs est également appelée convolution transposée. Pour en savoir plus, consultez tf.nn.conv2d_transpose
.
L'argument feature_group_count
(valeur par défaut 1) peut être utilisé pour les convolutions groupées. feature_group_count
doit être un diviseur de la dimension de caractéristique d'entrée et de sortie. Si feature_group_count
est supérieur à 1, cela signifie que, conceptuellement, les dimensions de caractéristique d'entrée et de sortie et la dimension de caractéristique de sortie rhs
sont divisées de manière égale en plusieurs groupes feature_group_count
, chacun étant constitué d'une sous-séquence consécutive de caractéristiques. La dimension de caractéristique d'entrée de rhs
doit être égale à la dimension de caractéristique d'entrée lhs
divisée par feature_group_count
(elle a donc déjà la taille d'un groupe de caractéristiques d'entrée). Les i-ièmes groupes sont utilisés ensemble afin de calculer feature_group_count
pour de nombreuses convolutions distinctes. Les résultats de ces convolutions sont concaténés dans la dimension de caractéristique de sortie.
Pour une convolution en profondeur, l'argument feature_group_count
serait défini sur la dimension de caractéristique d'entrée, et le filtre passerait de [filter_height, filter_width, in_channels, channel_multiplier]
à [filter_height, filter_width, 1, in_channels * channel_multiplier]
. Pour en savoir plus, consultez tf.nn.depthwise_conv2d
.
L'argument batch_group_count
(valeur par défaut 1) peut être utilisé pour les filtres groupés lors de la rétropropagation. batch_group_count
doit être un diviseur de la taille de la dimension de lot lhs
(entrée). Si batch_group_count
est supérieur à 1, cela signifie que la dimension du lot de sortie doit avoir une taille de input batch
/ batch_group_count
. batch_group_count
doit être un diviseur de la taille de la caractéristique de sortie.
La forme de sortie présente les dimensions suivantes, dans cet ordre:
batch
: la taille de cette dimension multipliée parbatch_group_count
doit être égale à la taille de la dimensionbatch
en lhs.z
: même taille queoutput-z
sur le noyau (rhs
).spatial_dims
: une valeur pour chaque emplacement valide de la fenêtre convolutive.
La figure ci-dessus montre le fonctionnement du champ batch_group_count
. En effet, nous segmentons chaque lot lhs en groupes batch_group_count
et faisons de même pour les caractéristiques de sortie. Ensuite, pour chacun de ces groupes, nous effectuons des convolutions par paire et concaténons la sortie avec la dimension de caractéristique de sortie. La sémantique opérationnelle de toutes les autres dimensions (caractéristique et spatiale) reste la même.
Les emplacements valides de la fenêtre convolutive sont déterminés par les progrès et la taille de la zone de base après remplissage.
Pour décrire le rôle d'une convolution, considérez une convolution 2D et choisissez des coordonnées batch
, z
, y
et x
fixes dans la sortie. Ensuite, (y,x)
est la position d'un angle de la fenêtre dans la zone de la base (par exemple, l'angle supérieur gauche, selon la façon dont vous interprétez les dimensions spatiales). Nous avons maintenant une fenêtre 2D, issue de la zone de base, où chaque point 2D est associé à un vecteur 1D. Nous obtenons donc une case 3D. À partir du noyau convolutif, comme nous avons corrigé la coordonnée de sortie z
, nous disposons également d'une zone 3D. Les deux cases ont les mêmes dimensions. Nous pouvons donc prendre la somme des produits par élément entre les deux cases (semblable à un produit scalaire). Il s'agit de la valeur de sortie.
Notez que si output-z
est par exemple, 5, chaque position de la fenêtre génère 5 valeurs dans la sortie dans la dimension z
. Ces valeurs diffèrent selon la partie du noyau convolutif utilisée : un cadre 3D distinct de valeurs est utilisé pour chaque coordonnée output-z
. Vous pouvez donc considérer cela comme 5 convolutions distinctes
avec un filtre différent pour chacune d'elles.
Voici un pseudo-code pour une convolution 2D avec remplissage et striding:
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
Consultez également XlaBuilder::ConvertElementType
.
Semblable à un static_cast
au niveau de l'élément en C++, cette méthode effectue une opération de conversion par élément d'une forme de données à une forme cible. Les dimensions doivent correspondre, et la conversion doit être définie au niveau des éléments. Par exemple, les éléments s32
deviennent des éléments f32
via une routine de conversion s32
en f32
.
ConvertElementType(operand, new_element_type)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau de type T avec les valeurs D |
new_element_type |
PrimitiveType |
type U |
Les dimensions de l'opérande et de la forme cible doivent correspondre. Les types d'éléments source et de destination ne doivent pas être des tuples.
Une conversion telle que T=s32
en U=f32
effectue une routine de normalisation de type "int-to-float" (par exemple, "arrondi au plus proche").
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
Effectue AllReduce
avec un calcul de somme.
CustomCall
Consultez également XlaBuilder::CustomCall
.
Appeler une fonction fournie par l'utilisateur dans un calcul
CustomCall(target_name, args..., shape)
Arguments | Type | Sémantique |
---|---|---|
target_name |
string |
Nom de la fonction. Une instruction d'appel ciblant ce nom de symbole sera émise. |
args |
séquence de N XlaOp |
N arguments de type arbitraire, qui seront transmis à la fonction. |
shape |
Shape |
Forme de sortie de la fonction |
La signature de la fonction est la même, quels que soient l'arité ou le type d'arguments:
extern "C" void target_name(void* out, void** in);
Par exemple, si CustomCall est utilisé comme suit:
let x = f32[2] {1,2};
let y = f32[2x3] { {10, 20, 30}, {40, 50, 60} };
CustomCall("myfunc", {x, y}, f32[3x3])
Voici un exemple d'implémentation de 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];
// ...
}
La fonction fournie par l'utilisateur ne doit pas avoir d'effets secondaires et son exécution doit être idempotente.
Dot
Consultez également XlaBuilder::Dot
.
Dot(lhs, rhs)
Arguments | Type | Sémantique |
---|---|---|
lhs |
XlaOp |
tableau de type T |
rhs |
XlaOp |
tableau de type T |
La sémantique exacte de cette opération dépend des rangs des opérandes:
Entrée | Sortie | Sémantique |
---|---|---|
vecteur [n] dot vecteur [n] |
scalaire | produit scalaire vectoriel |
matrice [m x k] vecteur dot [k] |
vecteur [m] | multiplication matricielle-vecteur |
matrice [m x k] matrice dot [k x n] |
matrice [m x n] | multiplication matricielle-matrice |
L'opération effectue la somme des produits en fonction de la deuxième dimension de lhs
(ou de la première si elle est de rang 1) et de la première dimension de rhs
. Ce sont les dimensions "contractées". Les dimensions contractées de lhs
et rhs
doivent être de la même taille. En pratique, il peut être utilisé pour effectuer des produits scalaires entre des vecteurs, des multiplications de vecteur/matrice ou des multiplications de matrice/matrice.
DotGeneral
Consultez également XlaBuilder::DotGeneral
.
DotGeneral(lhs, rhs, dimension_numbers)
Arguments | Type | Sémantique |
---|---|---|
lhs |
XlaOp |
tableau de type T |
rhs |
XlaOp |
tableau de type T |
dimension_numbers |
DotDimensionNumbers |
la contraction et les nombres de dimensions de lot |
Semblable à Dot, mais permet de spécifier des nombres de dimensions de lot et de contraction à la fois pour lhs
et rhs
.
Champs DotDimensionNumbers | Type | Sémantique |
---|---|---|
lhs_contracting_dimensions
|
(valeur répétée de type int64) | lhs numéros de dimension contractant |
rhs_contracting_dimensions
|
(valeur répétée de type int64) | rhs numéros de dimension contractant |
lhs_batch_dimensions
|
(valeur répétée de type int64) | lhs numéros de dimensions de lot |
rhs_batch_dimensions
|
(valeur répétée de type int64) | rhs numéros de dimensions de lot |
DotGeneral effectue la somme des produits sur les dimensions contractantes spécifiées dans dimension_numbers
.
Les numéros de dimensions contractées associés à partir de lhs
et rhs
n'ont pas besoin d'être identiques, mais doivent avoir les mêmes tailles de dimension.
Exemple avec des nombres de dimensions qui se contractent:
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} }
Les numéros de dimension de lot associés à partir de lhs
et rhs
doivent avoir les mêmes tailles de dimension.
Exemple avec des nombres de dimensions de lot (taille de lot 2, matrices 2x2):
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} } }
Entrée | Sortie | Sémantique |
---|---|---|
[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 |
Il s'ensuit que le numéro de dimension résultant commence par la dimension de lot, puis par la dimension lhs
non contractuelle/non liée au lot, et enfin par la dimension rhs
non contractuelle.
DynamicSlice
Consultez également XlaBuilder::DynamicSlice
.
DynamicSlice extrait un sous-tableau du tableau d'entrée au niveau du start_indices
dynamique. La taille de la tranche dans chaque dimension est transmise dans size_indices
, qui spécifie le point de fin des intervalles de tranche exclusifs dans chaque dimension: [début, début + taille). La forme de start_indices
doit être rang ==
1, la taille de la dimension étant égale au rang de operand
.
DynamicSlice(operand, start_indices, size_indices)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
Tableau à N dimensions de type T |
start_indices |
séquence de N XlaOp |
Liste de N entiers scalaires contenant les index de départ de la tranche pour chaque dimension. La valeur doit être supérieure ou égale à zéro. |
size_indices |
ArraySlice<int64> |
Liste de N entiers contenant la taille de la tranche pour chaque dimension. Chaque valeur doit être strictement supérieure à zéro, et le début + la taille doit être inférieur ou égal à la taille de la dimension pour éviter d'encapsuler la taille de la dimension modulo. |
Les index effectifs des tranches sont calculés en appliquant la transformation suivante pour chaque indice i
dans [1, N)
avant d'effectuer la tranche:
start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - size_indices[i])
Cela garantit que la tranche extraite est toujours comprise dans les limites par rapport au tableau d'opérande. Si la tranche est comprise dans les limites avant l'application de la transformation, celle-ci n'a aucun effet.
Exemple unidimensionnel:
let a = {0.0, 1.0, 2.0, 3.0, 4.0}
let s = {2}
DynamicSlice(a, s, {2}) produces:
{2.0, 3.0}
Exemple en deux dimensions:
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
Consultez également XlaBuilder::DynamicUpdateSlice
.
DynamicUpdateSlice génère un résultat qui correspond à la valeur du tableau d'entrée operand
, avec une tranche update
remplacée par start_indices
.
La forme de update
détermine la forme du sous-tableau du résultat mis à jour.
La forme de start_indices
doit être un rang de 1, avec une taille de dimension égale au rang de operand
.
DynamicUpdateSlice(operand, update, start_indices)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
Tableau à N dimensions de type T |
update |
XlaOp |
Tableau à N dimensions de type T contenant la mise à jour des tranches. Chaque dimension de la forme de mise à jour doit être strictement supérieure à zéro, et la combinaison début + actualisation doit être inférieure ou égale à la taille de l'opérande de chaque dimension pour éviter de générer des index de mise à jour hors limites. |
start_indices |
séquence de N XlaOp |
Liste de N entiers scalaires contenant les index de départ de la tranche pour chaque dimension. La valeur doit être supérieure ou égale à zéro. |
Les index effectifs des tranches sont calculés en appliquant la transformation suivante pour chaque indice i
dans [1, N)
avant d'effectuer la tranche:
start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - update.dimension_size[i])
Cela garantit que la tranche mise à jour est toujours dans les limites par rapport au tableau d'opérandes. Si la tranche est comprise dans les limites avant l'application de la transformation, celle-ci n'a aucun effet.
Exemple unidimensionnel:
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}
Exemple en deux dimensions:
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} }
Opérations arithmétiques binaires par élément
Consultez également XlaBuilder::Add
.
Un ensemble d'opérations arithmétiques binaires par élément est pris en charge.
Op(lhs, rhs)
Où Op
est l'un des éléments suivants : Add
(addition), Sub
(soustraction), Mul
(multiplication), Div
(division), Rem
(reste), Max
(maximum), Min
(minimum), LogicalAnd
(logique AND) ou LogicalOr
(opérateur logique OR).
Arguments | Type | Sémantique |
---|---|---|
lhs |
XlaOp |
opérande de gauche: tableau de type T |
rhs |
XlaOp |
opérande de droite: tableau de type T |
Les formes des arguments doivent être similaires ou compatibles. Consultez la documentation sur la diffusion pour découvrir ce que signifie la compatibilité des formes. Le résultat d'une opération a une forme résultant de la diffusion des deux tableaux d'entrée. Dans cette variante, les opérations entre des tableaux de rangs différents ne sont pas acceptées, sauf si l'un des opérandes est scalaire.
Lorsque Op
est défini sur Rem
, le signe du résultat est tiré du dividende, et la valeur absolue du résultat est toujours inférieure à la valeur absolue du diviseur.
Le dépassement de division d'un entier (division/résidence signée/non signée/reste par zéro ou division/reste signée de INT_SMIN
avec -1
) génère une valeur définie pour l'implémentation.
Il existe une autre variante compatible avec la diffusion de rangs différents pour ces opérations:
Op(lhs, rhs, broadcast_dimensions)
Où Op
est identique à ci-dessus. Cette variante de l'opération doit être utilisée pour les opérations arithmétiques entre des tableaux de rangs différents (par exemple, l'ajout d'une matrice à un vecteur).
L'opérande broadcast_dimensions
supplémentaire est une tranche d'entiers utilisée pour étendre le rang de l'opérande de rang inférieur jusqu'à celui de l'opérande de rang supérieur. broadcast_dimensions
mappe les dimensions de la forme de rang inférieur à celles de la forme de rang supérieur. Les dimensions non mappées de la forme développée
sont remplies avec des dimensions de taille 1. La diffusion de dimensions dégénérées diffuse ensuite les formes le long de ces dimensions dégénérées pour uniformiser les formes des deux opérandes. La sémantique est décrite en détail sur la page de diffusion.
Opérations de comparaison par élément
Consultez également XlaBuilder::Eq
.
Un ensemble d'opérations de comparaison binaire par élément standard est accepté. Notez que la sémantique de comparaison à virgule flottante standard IEEE 754 s'applique lors de la comparaison de types à virgule flottante.
Op(lhs, rhs)
Où Op
est l'un des éléments suivants : Eq
(égal à), Ne
(pas égal à), Ge
(supérieur ou égal à), Gt
(supérieur à), Le
(inférieur ou égal à), Lt
(inférieur à). Un autre ensemble d'opérateurs, EqTotalOrder, NeTotalOrder, GeTotalOrder, GtTotalOrder, LeTotalOrder et LtTotalOrder, offrent les mêmes fonctionnalités, sauf qu'ils acceptent également un ordre total sur les nombres à virgule flottante, en appliquant -NaN < -Inf < -Finite < -0 < +0 < +Infnite < +Infnite <
Arguments | Type | Sémantique |
---|---|---|
lhs |
XlaOp |
opérande de gauche: tableau de type T |
rhs |
XlaOp |
opérande de droite: tableau de type T |
Les formes des arguments doivent être similaires ou compatibles. Consultez la documentation sur la diffusion pour découvrir ce que signifie la compatibilité des formes. Le résultat d'une opération a une forme résultant de la diffusion des deux tableaux d'entrée avec le type d'élément PRED
. Dans cette variante, les opérations entre des tableaux de rangs différents ne sont pas acceptées, sauf si l'un des opérandes est scalaire.
Il existe une autre variante compatible avec la diffusion de rangs différents pour ces opérations:
Op(lhs, rhs, broadcast_dimensions)
Où Op
est identique à ci-dessus. Cette variante de l'opération doit être utilisée pour les opérations de comparaison entre des tableaux de rangs différents (par exemple, pour ajouter une matrice à un vecteur).
L'opérande broadcast_dimensions
supplémentaire est une tranche d'entiers spécifiant les dimensions à utiliser pour diffuser les opérandes. La sémantique est décrite en détail sur la page de diffusion.
Fonctions unaires par élément
XlaBuilder prend en charge ces fonctions unaires par élément:
Abs(operand)
x -> |x|
Abs au niveau de l'élément.
Ceil(operand)
Plafond au niveau des éléments x -> ⌈x⌉
.
Cos(operand)
Cosinus par élément x -> cos(x)
.
Exp(operand)
x -> e^x
exponentielle naturel au niveau des éléments.
Floor(operand)
Valeur minimale de l'élément : x -> ⌊x⌋
.
Imag(operand)
Partie imaginaire par élément d'une forme complexe (ou réelle). x -> imag(x)
. Si l'opérande est de type à virgule flottante, renvoie 0.
IsFinite(operand)
Teste si chaque élément de operand
est fini, c'est-à-dire s'il ne correspond pas à l'infini positif ou négatif, et n'est pas NaN
. Renvoie un tableau de valeurs PRED
ayant la même forme que l'entrée, où chaque élément est true
si et seulement si l'élément d'entrée correspondant est fini.
Log(operand)
Logarithme naturel par élément x -> ln(x)
.
LogicalNot(operand)
Logique par élément et non x -> !(x)
.
Logistic(operand)
Calcul de la fonction logistique par élément x ->
logistic(x)
.
PopulationCount(operand)
Calcule le nombre de bits définis dans chaque élément de operand
.
Neg(operand)
Négation par élément x -> -x
.
Real(operand)
Partie réelle par élément d'une forme complexe (ou réelle).
x -> real(x)
. Si l'opérande est de type à virgule flottante, renvoie la même valeur.
Rsqrt(operand)
Réciproque élément par élément de l'opération de racine carrée x -> 1.0 / sqrt(x)
.
Sign(operand)
Opération de signe par élément x -> sgn(x)
où
\[\text{sgn}(x) = \begin{cases} -1 & x < 0\\ -0 & x = -0\\ NaN & x = NaN\\ +0 & x = +0\\ 1 & x > 0 \end{cases}\]
à l'aide de l'opérateur de comparaison du type d'élément operand
.
Sqrt(operand)
Opération de racine carrée par élément x -> sqrt(x)
.
Cbrt(operand)
Opération de racine cubique par élément x -> cbrt(x)
.
Tanh(operand)
Tangente hyperbolique par élément x -> tanh(x)
.
Round(operand)
Arrondissement au niveau des éléments, éliminant les écarts par rapport à zéro.
RoundNearestEven(operand)
Arrondi au niveau des éléments, lié au pair le plus proche.
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
L'opérande de la fonction |
La fonction est appliquée à chaque élément du tableau operand
, ce qui génère un tableau ayant la même forme. operand
peut être un scalaire (rang 0).
TF1
L'opération XLA FFT implémente les transformations de Fourier avant et inverse pour les entrées/sorties réelles et complexes. Les FFT multidimensionnels sont acceptés sur trois axes au maximum.
Consultez également XlaBuilder::Fft
.
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
Tableau que nous transformons de Fourier. |
fft_type |
FftType |
Consultez le tableau ci-dessous. |
fft_length |
ArraySlice<int64> |
Longueurs de domaine temporel des axes transformés. Cela est particulièrement nécessaire pour que IRFFT puisse dimensionner à droite l'axe le plus interne, car RFFT(fft_length=[16]) a la même forme de sortie que RFFT(fft_length=[17]) . |
FftType |
Sémantique |
---|---|
FFT |
Transférer des opérations FFT complexes à complexes. La forme reste inchangée. |
IFFT |
FFT inverse de complexe à complexe. La forme reste inchangée. |
RFFT |
Transférer les fichiers FFT réel vers complexe La forme de l'axe le plus interne est réduite à fft_length[-1] // 2 + 1 si fft_length[-1] est une valeur non nulle, en omettant la partie conjuguée inversée du signal transformé au-delà de la fréquence de Nyquist. |
IRFFT |
FFT inverse de vrai à complexe (plus complexe, renvoie un vrai). La forme de l'axe le plus interne est étendue à fft_length[-1] si fft_length[-1] est une valeur non nulle, en déduisant la partie du signal transformé au-delà de la fréquence de Nyquist à partir du conjugué inverse des entrées 1 en fft_length[-1] // 2 + 1 . |
FFT multidimensionnel
Lorsque plusieurs fft_length
sont fournis, cela équivaut à appliquer une cascade d'opérations FFT à chacun des axes les plus internes. Notez que pour les cas réels, complexes ou complexes, la transformation de l'axe le plus interne est (effectivement) effectuée en premier (RFFT ; dernière pour IRFFT). C'est pourquoi l'axe le plus interne est celui qui modifie la taille. Les autres transformations des axes seront alors complexes.
Détails de mise en œuvre
Le processeur FFT s'appuie sur TensorFFT d'Eigen. Le GPU FFT utilise cuFFT.
Rassembler
L'opération de collecte XLA assemble plusieurs tranches d'un tableau d'entrée (chaque tranche avec un décalage d'exécution potentiellement différent).
Sémantique générale
Consultez également XlaBuilder::Gather
.
Pour une description plus intuitive, consultez la section "Description informelle" ci-dessous.
gather(operand, start_indices, offset_dims, collapsed_slice_dims, slice_sizes, start_index_map)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
Tableau à partir duquel nous collectons les données. |
start_indices |
XlaOp |
Tableau contenant les index de départ des tranches que nous recueillons. |
index_vector_dim |
int64 |
Dimension de start_indices qui "contient" les index de départ. Vous trouverez une description détaillée ci-dessous. |
offset_dims |
ArraySlice<int64> |
Ensemble des dimensions de la forme de sortie qui se décalent dans un tableau divisé par rapport à l'opérande. |
slice_sizes |
ArraySlice<int64> |
slice_sizes[i] est les limites de la tranche au niveau de la dimension i . |
collapsed_slice_dims |
ArraySlice<int64> |
Ensemble des dimensions de chaque secteur réduit. Ces dimensions doivent être de taille 1. |
start_index_map |
ArraySlice<int64> |
Carte qui décrit comment mapper les index de start_indices à des index juridiques dans un opérande. |
indices_are_sorted |
bool |
Indique si le tri des index est garanti par l'appelant. |
Pour plus de commodité, nous étiquetons les dimensions dans le tableau de sortie sous la forme batch_dims
, et non dans offset_dims
.
La sortie est un tableau de rang batch_dims.size
+ offset_dims.size
.
La valeur operand.rank
doit être égale à la somme de offset_dims.size
et collapsed_slice_dims.size
. De plus, slice_sizes.size
doit être égal à operand.rank
.
Si index_vector_dim
est égal à start_indices.rank
, nous considérons implicitement que start_indices
a une dimension 1
de fin (par exemple, si start_indices
avait la forme [6,7]
et que index_vector_dim
était 2
, nous considérons implicitement que la forme de start_indices
est [6,7,1]
).
Les limites du tableau de sortie en fonction de la dimension i
sont calculées comme suit:
Si
i
est présent dansbatch_dims
(c'est-à-dire qu'il est égal àbatch_dims[k]
pour certainsk
), nous choisissons les limites de dimension correspondantes dansstart_indices.shape
, en ignorantindex_vector_dim
(par exemple, sélectionnezstart_indices.shape.dims
[k
] sik
<index_vector_dim
etstart_indices.shape.dims
[k
+1
] dans le cas contraire).Si
i
est présent dansoffset_dims
(c'est-à-dire égal àoffset_dims
[k
] pour certainsk
), nous choisissons la limite correspondante deslice_sizes
après avoir pris en comptecollapsed_slice_dims
(par exemple, nous choisissonsadjusted_slice_sizes
[k
] oùadjusted_slice_sizes
estslice_sizes
en supprimant les limites des indicescollapsed_slice_dims
).
Officiellement, l'index d'opérande In
correspondant à un index de sortie donné (Out
) est calculé comme suit:
Soit
G
= {Out
[k
] pourk
dansbatch_dims
}. UtilisezG
pour diviser un vecteurS
tel queS
[i
] =start_indices
[Combine(G
,i
)], où Combine(A, b) insère b à la positionindex_vector_dim
dans A. Notez que ce paramètre est bien défini même siG
est vide: siG
est vide,S
=start_indices
.Créez un index de départ,
S
in
, dansoperand
à l'aide deS
en dispersantS
à l'aide destart_index_map
. Plus précisément:S
in
[start_index_map
[k
]] =S
[k
] sik
<start_index_map.size
.S
in
[_
] =0
dans les autres cas.
Créez un indice
O
in
dansoperand
en dispersant les indices au niveau des dimensions de décalage dansOut
, en fonction de l'ensemblecollapsed_slice_dims
. Plus précisément:O
in
[remapped_offset_dims
(k
)] =Out
[offset_dims
[k
]] sik
<offset_dims.size
(remapped_offset_dims
est défini ci-dessous).O
in
[_
] =0
dans les autres cas.
In
correspond àO
in
+S
in
, où + correspond à l'addition au niveau de l'élément.
remapped_offset_dims
est une fonction monotone avec le domaine [0
, offset_dims.size
) et la plage [0
, operand.rank
) \ collapsed_slice_dims
. Par exemple, offset_dims.size
est 4
, operand.rank
est 6
et collapsed_slice_dims
est {0
, 2
}, puis remapped_offset_dims
est {0
→1
,
1
→3
, 2
→4
, 3
→5
}.
Si indices_are_sorted
est défini sur "true", XLA peut supposer que les start_indices
sont triées (dans l'ordre croissant des start_index_map
) par l'utilisateur. Si ce n'est pas le cas, la sémantique est définie par l'implémentation.
Description et exemples informels
De manière informelle, chaque indice Out
dans le tableau de sortie correspond à un élément E
dans le tableau d'opérandes, calculé comme suit:
Nous utilisons les dimensions de lot dans
Out
pour rechercher un index de départ à partir destart_indices
.Nous utilisons
start_index_map
pour mapper l'index de départ (dont la taille peut être inférieure à operand.rank) à un index de départ "complet" dansoperand
.Une tranche de taille
slice_sizes
est divisée de façon dynamique à l'aide de l'index de départ complet.Nous remodelons la tranche en réduisant les dimensions
collapsed_slice_dims
. Étant donné que toutes les dimensions des secteurs réduits doivent avoir une limite de 1, cette restructuration est toujours légale.Nous utilisons les dimensions de décalage dans
Out
pour indexer cette tranche et obtenir l'élément d'entrée,E
, correspondant à l'index de sortieOut
.
index_vector_dim
est défini sur start_indices.rank
- 1
dans tous les exemples suivants. Des valeurs plus intéressantes pour index_vector_dim
ne modifient pas fondamentalement l'opération, mais rendent la représentation visuelle plus fastidieuse.
Pour comprendre comment tous les éléments ci-dessus s'imbriquent, examinons un exemple qui collecte cinq tranches de forme [8,6]
à partir d'un tableau [16,11]
. La position d'une tranche dans le tableau [16,11]
peut être représentée par un vecteur d'indice de forme S64[2]
. L'ensemble des cinq positions peut donc être représenté sous la forme d'un tableau S64[5,2]
.
Le comportement de l'opération de collecte peut ensuite être représenté sous la forme d'une transformation d'index qui utilise [G
,O
0
,O
1
], un indice en sortie, et le mappe à un élément du tableau d'entrée de la manière suivante:
Nous sélectionnons d'abord un vecteur (X
,Y
) dans le tableau de collecte d'indices à l'aide de G
.
L'élément du tableau de sortie avec l'indice [G
,O
0
,O
1
] est alors l'élément du tableau d'entrée avec l'indice [X
+O
0
,Y
+O
1
].
slice_sizes
est [8,6]
, ce qui détermine la plage de O0
et O1
, ce qui détermine à son tour les limites de la tranche.
Cette opération de collecte agit comme une tranche dynamique par lot avec G
comme dimension de lot.
Les index de collecte peuvent être multidimensionnels. Par exemple, une version plus générale de l'exemple ci-dessus utilisant un tableau "gather index" de forme [4,5,2]
traduira les index comme suit:
Là encore, il s'agit d'une tranche dynamique de lot G
0
et G
1
comme dimensions de lot. La taille du segment est toujours de [8,6]
.
L'opération de collecte dans XLA généralise la sémantique informelle décrite ci-dessus de la manière suivante:
Nous pouvons configurer les dimensions de la forme de sortie qui correspondent aux dimensions du décalage (dimensions
O
0
,O
1
dans le dernier exemple). Les dimensions du lot de sortie (dimensions contenantG
0
etG
1
dans le dernier exemple) sont définies comme étant les dimensions de sortie qui ne sont pas des dimensions décalées.Le nombre de dimensions de décalage de sortie explicitement présentes dans la forme de sortie peut être inférieur au rang d'entrée. Ces dimensions "manquantes", qui sont explicitement listées comme
collapsed_slice_dims
, doivent avoir une taille de tranche de1
. Étant donné qu'ils ont une taille de tranche de1
, le seul indice valide pour eux est0
et leur suppression n'introduit pas d'ambiguïté.La tranche extraite du tableau "Compiler des indices" ((
X
,Y
) dans le dernier exemple) peut comporter moins d'éléments que le rang du tableau d'entrée. En outre, un mappage explicite indique comment l'index doit être développé pour qu'il ait le même rang que l'entrée.
Comme dernier exemple, nous utilisons (2) et (3) pour implémenter tf.gather_nd
:
G
0
et G
1
sont utilisés pour séparer un index de départ du tableau de collecte d'indices comme d'habitude, sauf que l'index de départ ne contient qu'un seul élément, X
. De même, il n'existe qu'un seul index de décalage de sortie avec la valeur O
0
. Toutefois, avant d'être utilisés en tant qu'indices dans le tableau d'entrée, ceux-ci sont développés conformément aux règles "Collecter les mappages d'index" (start_index_map
dans la description formelle) et "Mappage de décalage" (remapped_offset_dims
dans la description formelle) dans [X
,0
] et [0
,O
0
] respectivement, en ajoutant jusqu'à [X
,O
0
]. En d'autres termes, l'index de sortie est mappé sur 0
/1
, l'index de sortie.0
0
0
0
O
O
G
G
G
G
1
GatherIndices
tf.gather_nd
Dans ce cas, slice_sizes
est [1,11]
. Intuitivement, cela signifie que chaque index X
du tableau de collecte des index sélectionne une ligne entière et que le résultat est la concaténation de toutes ces lignes.
GetDimensionSize
Consultez également XlaBuilder::GetDimensionSize
.
Renvoie la taille de la dimension donnée de l'opérande. L'opérande doit être en forme de tableau.
GetDimensionSize(operand, dimension)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau d'entrée à n dimensions |
dimension |
int64 |
Une valeur de l'intervalle [0, n) qui spécifie la dimension |
SetDimensionSize
Consultez également XlaBuilder::SetDimensionSize
.
Définit la taille dynamique de la dimension donnée de XlaOp. L'opérande doit être en forme de tableau.
SetDimensionSize(operand, size, dimension)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau d'entrée à n dimensions. |
size |
XlaOp |
"int32" représentant la taille dynamique de l'environnement d'exécution. |
dimension |
int64 |
Valeur de l'intervalle [0, n) qui spécifie la dimension. |
Transmettez l'opérande comme résultat, avec la dimension dynamique suivie par le compilateur.
Les valeurs remplies seront ignorées par les opérations de réduction en aval.
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
Consultez également XlaBuilder::GetTupleElement
.
Indexe dans un tuple avec une valeur de constante au moment de la compilation.
La valeur doit être une constante de temps de compilation afin que l'inférence de forme puisse déterminer le type de la valeur obtenue.
Cette méthode est semblable à std::get<int N>(t)
en C++. D'un point de vue conceptuel:
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.
Voir aussi tf.tuple
.
Flux d'entrée
Consultez également XlaBuilder::Infeed
.
Infeed(shape)
Argument | Type | Sémantique |
---|---|---|
shape |
Shape |
Forme des données lues à partir de l'interface Infeed. Le champ de mise en page de la forme doit être défini pour correspondre à la mise en page des données envoyées à l'appareil. Sinon, son comportement n'est pas défini. |
Lit un seul élément de données à partir de l'interface de flux Infeed implicite de l'appareil, en interprétant les données comme la forme et la mise en page données, puis renvoie un XlaOp
des données. Plusieurs opérations Infeed sont autorisées dans un calcul, mais il doit y avoir un ordre total parmi les opérations Infeed. Par exemple, deux flux Infeed dans le code ci-dessous ont un ordre total, car il existe une dépendance entre les boucles "when".
result1 = while (condition, init = init_value) {
Infeed(shape)
}
result2 = while (condition, init = result1) {
Infeed(shape)
}
Les formes tuple imbriquées ne sont pas acceptées. Pour une forme de tuple vide, l'opération Infeed est effectivement une opération no-op et se poursuit sans lire les données du flux d'entrée de l'appareil.
Iota
Consultez également XlaBuilder::Iota
.
Iota(shape, iota_dimension)
Crée un littéral constant sur l'appareil plutôt qu'un transfert d'hôte potentiellement volumineux. Crée un tableau avec une forme spécifiée et contient des valeurs commençant à zéro et incrémentées d'une unité selon la dimension spécifiée. Pour les types à virgule flottante, le tableau produit est équivalent à ConvertElementType(Iota(...))
, où Iota
est de type intégral et que la conversion est de type à virgule flottante.
Arguments | Type | Sémantique |
---|---|---|
shape |
Shape |
Forme du tableau créé par Iota() |
iota_dimension |
int64 |
Dimension à incrémenter. |
Par exemple, Iota(s32[4, 8], 0)
renvoie
[[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 ]]
Retours pour Iota(s32[4, 8], 1)
[[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 ]]
Mappage
Consultez également XlaBuilder::Map
.
Map(operands..., computation)
Arguments | Type | Sémantique |
---|---|---|
operands |
séquence de N XlaOp |
N tableaux de types T0..T{N-1} |
computation |
XlaComputation |
un calcul de type T_0, T_1, .., T_{N + M -1} -> S avec N paramètres de type T et M de type arbitraire |
dimensions |
Tableau int64 |
tableau de dimensions de carte |
Applique une fonction scalaire aux tableaux operands
donnés, produisant ainsi un tableau aux mêmes dimensions, où chaque élément est le résultat de la fonction mappée appliquée aux éléments correspondants dans les tableaux d'entrée.
La fonction mappée est un calcul arbitraire avec la restriction selon laquelle elle comporte N entrées de type scalaire T
et une seule sortie de type S
. La sortie a les mêmes dimensions que les opérandes, à la différence que le type d'élément T est remplacé par S.
Par exemple, Map(op1, op2, op3, computation, par1)
mappe elem_out <-
computation(elem1, elem2, elem3, par1)
à chaque indice (multidimensionnel) des tableaux d'entrée pour produire le tableau de sortie.
OptimizationBarrier
Empêche toute passe d'optimisation de déplacer les calculs au-delà de la barrière.
Il garantit que toutes les entrées sont évaluées avant les opérateurs qui dépendent des sorties de la barrière.
Pad
Consultez également XlaBuilder::Pad
.
Pad(operand, padding_value, padding_config)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau de type T |
padding_value |
XlaOp |
scalaire de type T pour remplir la marge intérieure ajoutée |
padding_config |
PaddingConfig |
quantité de marge intérieure sur les deux bords (bas, haut) et entre les éléments de chaque dimension |
Développe le tableau operand
donné en appliquant une marge intérieure autour du tableau et entre les éléments du tableau avec la valeur padding_value
donnée. padding_config
spécifie la quantité de marge intérieure des bords et la marge intérieure intérieure pour chaque dimension.
PaddingConfig
est un champ répété de PaddingConfigDimension
, qui contient trois champs pour chaque dimension: edge_padding_low
, edge_padding_high
et interior_padding
.
edge_padding_low
et edge_padding_high
spécifient la quantité de marge intérieure ajoutée au bas de la page (à côté de l'index 0) et à la marge supérieure (à côté de l'indice le plus élevé) de chaque dimension. La valeur de la marge intérieure négative peut être négative. La valeur absolue de la marge intérieure négative indique le nombre d'éléments à supprimer de la dimension spécifiée.
interior_padding
spécifie la quantité de marge intérieure ajoutée entre deux éléments de chaque dimension. Elle ne peut pas être négative. Le remplissage intérieur se produit logiquement avant le remplissage des bords. Par conséquent, dans le cas d'une marge intérieure négative, des éléments sont supprimés de l'opérande de remplissage intérieur.
Il s'agit d'une opération no-op si toutes les paires de marges intérieures du bord sont (0, 0) et que les valeurs de marge intérieure intérieure sont toutes égales à 0. La figure ci-dessous montre des exemples de différentes valeurs edge_padding
et interior_padding
pour un tableau à deux dimensions.
Recv
Consultez également XlaBuilder::Recv
.
Recv(shape, channel_handle)
Arguments | Type | Sémantique |
---|---|---|
shape |
Shape |
forme des données à recevoir |
channel_handle |
ChannelHandle |
un identifiant unique pour chaque paire envoi/réception |
Reçoit les données d'une forme donnée à partir d'une instruction Send
dans un autre calcul partageant le même handle de canal. Renvoie un XlaOp pour les données reçues.
L'API cliente de l'opération Recv
représente une communication synchrone.
Cependant, l'instruction est décomposée en interne en deux instructions HLO (Recv
et RecvDone
) pour permettre les transferts de données asynchrones. Consultez également HloInstruction::CreateRecv
et HloInstruction::CreateRecvDone
.
Recv(const Shape& shape, int64 channel_id)
Alloue les ressources requises pour recevoir les données d'une instruction Send
avec le même channel_id. Renvoie un contexte pour les ressources allouées, qui est utilisé par une instruction RecvDone
suivante pour attendre la fin du transfert de données. Le contexte est un tuple de {tampon de réception (forme), identifiant de requête (U32)} et ne peut être utilisé que par une instruction RecvDone
.
RecvDone(HloInstruction context)
À partir d'un contexte créé par une instruction Recv
, attend la fin du transfert de données et renvoie les données reçues.
Restreindre
Consultez également XlaBuilder::Reduce
.
Applique une fonction de réduction à un ou plusieurs tableaux en parallèle.
Reduce(operands..., init_values..., computation, dimensions)
Arguments | Type | Sémantique |
---|---|---|
operands |
Séquence de N XlaOp |
N tableaux de types T_0, ..., T_{N-1} . |
init_values |
Séquence de N XlaOp |
N scalaires de types T_0, ..., T_{N-1} . |
computation |
XlaComputation |
de type T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}) . |
dimensions |
Tableau int64 |
tableau de dimensions non ordonnées à réduire. |
Où :
- N doit être supérieur ou égal à 1.
- Le calcul doit être "grossièrement" associatif (voir ci-dessous).
- Tous les tableaux d'entrée doivent avoir les mêmes dimensions.
- Toutes les valeurs initiales doivent former une identité sous
computation
. - Si
N = 1
,Collate(T)
estT
. - Si la valeur est
N > 1
,Collate(T_0, ..., T_{N-1})
est un tuple d'élémentsN
de typeT
.
Cette opération réduit une ou plusieurs dimensions de chaque tableau d'entrée en scalaires.
Le rang de chaque tableau renvoyé est rank(operand) - len(dimensions)
. Le résultat de l'opération est Collate(Q_0, ..., Q_N)
, où Q_i
est un tableau de type T_i
, dont les dimensions sont décrites ci-dessous.
Différents backends sont autorisés à réassocier le calcul de réduction. Cela peut entraîner des différences numériques, car certaines fonctions de réduction telles que l'addition ne sont pas associatives pour les flottants. Toutefois, si la plage des données est limitée, l'addition à virgule flottante est suffisamment associative pour la plupart des utilisations pratiques.
Exemples
Lors de la réduction d'une dimension dans un seul tableau unidimensionnel avec des valeurs [10, 11,
12, 13]
, avec la fonction de réduction f
(computation
), cela peut être calculé comme suit :
f(10, f(11, f(12, f(init_value, 13)))
mais il existe bien d'autres possibilités. Par exemple,
f(init_value, f(f(10, f(init_value, 11)), f(f(init_value, 12), f(init_value, 13))))
Voici un exemple approximatif de pseudo-code de la façon dont la réduction peut être mise en œuvre, en utilisant la somme comme calcul de réduction avec une valeur initiale de 0.
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]
Voici un exemple de réduction d'un tableau bidimensionnel (matrice). La forme a le rang 2, la dimension 0 de la taille 2 et la dimension 1 de la taille 3:
Réduire les dimensions de 0 ou 1 à l'aide d'une fonction "add" peut s'afficher:
Notez que les deux résultats de réduction sont des tableaux unidimensionnels. Pour plus de commodité, l'une est représentée par une colonne et l'autre par ligne.
Voici un exemple plus complexe : Son rang est 3, dimension 0 de taille 4, dimension 1 de taille 2 et dimension 2 de taille 3. Pour plus de simplicité, les valeurs de 1 à 6 sont répliquées sur la dimension 0.
Comme dans l'exemple en 2D, nous pouvons réduire une seule dimension. Si nous réduisons la dimension 0, par exemple, nous obtenons un tableau de rang 2 dans lequel toutes les valeurs de la dimension 0 ont été pliées dans une valeur scalaire:
| 4 8 12 |
| 16 20 24 |
Si nous réduisons la dimension 2, nous obtenons également un tableau rang-2 dans lequel toutes les valeurs de la dimension 2 ont été pliées dans une valeur scalaire:
| 6 15 |
| 6 15 |
| 6 15 |
| 6 15 |
Notez que l'ordre relatif entre les dimensions restantes de l'entrée est conservé dans la sortie, mais que de nouveaux nombres peuvent être attribués à certaines dimensions (puisque le rang change).
Nous pouvons également réduire plusieurs dimensions. La réduction des dimensions 0 et 1 génère le tableau unidimensionnel [20, 28, 36]
.
La réduction du tableau 3D dans toutes ses dimensions produit la 84
scalaire.
Réduit variadique
Lorsque N > 1
, l'application de la fonction de réduction est légèrement plus complexe, car elle est appliquée simultanément à toutes les entrées. Les opérandes sont fournis au calcul dans l'ordre suivant:
- Exécuter une valeur réduite pour le premier opérande
- ...
- Exécution de la valeur réduite pour le n-ième opérande
- Valeur d'entrée pour le premier opérande
- ...
- Valeur d'entrée pour le n-ième opérande
Prenons l'exemple de la fonction de réduction suivante, qui permet de calculer la valeur maximale et la valeur argmax d'un tableau 1D en parallèle:
f: (Float, Int, Float, Int) -> Float, Int
f(max, argmax, value, index):
if value >= max:
return (value, index)
else:
return (max, argmax)
Pour les tableaux d'entrées unidimensionnels V = Float[N], K = Int[N]
et les valeurs d'initialisation I_V = Float, I_K = Int
, le résultat f_(N-1)
de la réduction dans la seule dimension d'entrée est équivalent à l'application récursive suivante:
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))
L'application de cette réduction à un tableau de valeurs et à un tableau d'indices séquentiels (iota, par exemple) permet de co-itérer les tableaux et de renvoyer un tuple contenant la valeur maximale et l'index correspondant.
ReducePrecision
Consultez également XlaBuilder::ReducePrecision
.
Modélise l'effet de la conversion des valeurs à virgule flottante dans un format de précision inférieure (par exemple, IEEE-FP16) et revient au format d'origine. Le nombre de bits d'exposant et de mantissa au format de précision inférieure peut être spécifié de manière arbitraire, bien que toutes les tailles de bits puissent ne pas être compatibles avec toutes les implémentations matérielles.
ReducePrecision(operand, mantissa_bits, exponent_bits)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
Tableau de type T à virgule flottante. |
exponent_bits |
int32 |
nombre de bits d'exposant dans un format de précision inférieure |
mantissa_bits |
int32 |
nombre de bits de mantissa dans un format de précision inférieure |
Le résultat est un tableau de type T
. Les valeurs d'entrée sont arrondies à la valeur la plus proche pouvant être représentée par le nombre donné de bits de mantisse (à l'aide de la sémantique "liaisons à pair"), et toutes les valeurs qui dépassent la plage spécifiée par le nombre de bits d'exposant sont limitées à l'infini positif ou négatif. Les valeurs NaN
sont conservées, bien qu'elles puissent être converties en valeurs NaN
canoniques.
Le format de précision inférieure doit comporter au moins un bit d'exposant (afin de distinguer une valeur nulle d'un infini, car ils ont tous deux une mantisse nulle) et un nombre non négatif de bits de mantisse. Le nombre de bits d'exposant ou de mantissa peut dépasser la valeur correspondante pour le type T
. La partie correspondante de la conversion est alors simplement une "no-op".
ReduceScatter
Consultez également XlaBuilder::ReduceScatter
.
La stratégie de dispersion est une opération collective qui effectue efficacement une opération AllReduce, puis disperse le résultat en le divisant en blocs shard_count
le long de l'élément scatter_dimension
. L'instance répliquée i
du groupe d'instances répliquées reçoit le segment ith
.
ReduceScatter(operand, computation, scatter_dim, shard_count,
replica_group_ids, channel_id)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
Tableau ou tuple de tableaux non vide à réduire pour le nombre d'instances répliquées. |
computation |
XlaComputation |
Calcul de la réduction |
scatter_dimension |
int64 |
Dimension à disperser. |
shard_count |
int64 |
Nombre de blocs à diviser scatter_dimension |
replica_groups |
vecteur de vecteurs de int64 |
Groupes entre lesquels les réductions sont effectuées |
channel_id |
(Facultatif) int64 |
ID de canal facultatif pour la communication entre les modules |
- Lorsque
operand
est un tuple de tableaux, la réduction de la diffusion est effectuée sur chaque élément du tuple. replica_groups
est une liste de groupes d'instances répliquées entre lesquels la réduction est effectuée (l'ID de l'instance répliquée de l'instance répliquée actuelle peut être récupéré à l'aide deReplicaId
). L'ordre des instances répliquées dans chaque groupe détermine l'ordre dans lequel le résultat de la réduction complète sera dispersé.replica_groups
doit être vide (auquel cas toutes les instances répliquées appartiennent à un seul groupe) ou contenir le même nombre d'éléments que le nombre d'instances répliquées. Lorsque plusieurs groupes d'instances répliquées sont présents, ils doivent tous avoir la même taille. Par exemple,replica_groups = {0, 2}, {1, 3}
effectue une réduction entre les instances répliquées0
et2
, et1
et3
, puis disperse le résultat.shard_count
est la taille de chaque groupe d'instances répliquées. Nous en avons besoin dans les cas oùreplica_groups
est vide. Sireplica_groups
n'est pas vide,shard_count
doit être égal à la taille de chaque groupe d'instances répliquées.channel_id
est utilisé pour la communication entre les modules: seules les opérationsreduce-scatter
ayant le mêmechannel_id
peuvent communiquer entre elles.
La forme de sortie est celle d'entrée, la scatter_dimension
rendue shard_count
fois plus petite. Par exemple, s'il y a deux instances répliquées et que l'opérande a respectivement les valeurs [1.0, 2.25]
et [3.0, 5.25]
sur les deux instances répliquées, la valeur de sortie de cette opération, où scatter_dim
est 0
, sera [4.0]
pour la première instance répliquée et [7.5]
pour la deuxième instance répliquée.
ReduceWindow
Consultez également XlaBuilder::ReduceWindow
.
Applique une fonction de réduction à tous les éléments de chaque fenêtre d'une séquence de N tableaux multidimensionnels, produisant ainsi un seul ou un tuple de N tableaux multidimensionnels en sortie. Chaque tableau de sortie contient le même nombre d'éléments que le nombre de positions valides de la fenêtre. Une couche de pooling peut être exprimée par ReduceWindow
. Comme pour Reduce
, le computation
appliqué reçoit toujours le init_values
à gauche.
ReduceWindow(operands..., init_values..., computation, window_dimensions,
window_strides, padding)
Arguments | Type | Sémantique |
---|---|---|
operands |
N XlaOps |
Séquence de N tableaux multidimensionnels de types T_0,..., T_{N-1} , chacun représentant la zone de base sur laquelle la fenêtre est placée. |
init_values |
N XlaOps |
N valeurs de départ pour la réduction, une pour chacun des N opérandes. Pour en savoir plus, consultez Réduire. |
computation |
XlaComputation |
Fonction de réduction de type T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}) , à appliquer aux éléments de chaque fenêtre de tous les opérandes d'entrée. |
window_dimensions |
ArraySlice<int64> |
tableau d'entiers pour les valeurs des dimensions de fenêtre |
window_strides |
ArraySlice<int64> |
tableau d'entiers pour les valeurs de pas de fenêtre |
base_dilations |
ArraySlice<int64> |
tableau d'entiers pour les valeurs de dilatation de base |
window_dilations |
ArraySlice<int64> |
tableau d'entiers pour les valeurs de dilatation de fenêtre |
padding |
Padding |
Type de marge intérieure pour la fenêtre (Padding::kSame, qui crée une marge intérieure afin d'avoir la même forme de sortie que l'entrée si le pas a la valeur 1, ou Padding::kValid, qui n'utilise pas de marge intérieure et "arrête" la fenêtre lorsqu'elle ne rentre plus) |
Où :
- N doit être supérieur ou égal à 1.
- Tous les tableaux d'entrée doivent avoir les mêmes dimensions.
- Si
N = 1
,Collate(T)
estT
. - Si la valeur est
N > 1
,Collate(T_0, ..., T_{N-1})
est un tuple d'élémentsN
de type(T0,...T{N-1})
.
Le code et la figure ci-dessous présentent un exemple d'utilisation de ReduceWindow
. L'entrée est une matrice de taille [4x6], et window_dimensions et window_stride_dimensions correspondent à [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);
Un pas de 1 dans une dimension indique que la position d'une fenêtre dans la dimension est à 1 élément de sa fenêtre adjacente. Pour qu'aucune fenêtre ne se chevauche, la valeur window_stride_dimensions doit être égale à window_dimensions. La figure ci-dessous illustre l'utilisation de deux valeurs de pas différentes. La marge intérieure est appliquée à chaque dimension de l'entrée, et les calculs sont les mêmes que si l'entrée arrivait avec les dimensions qu'elle contient après la marge intérieure.
Pour un exemple de marge intérieure non négligeable, envisagez de calculer une valeur minimale de la fenêtre de réduction (la valeur initiale est MAX_FLOAT
) avec une dimension 3
et de fouler 2
sur le tableau d'entrée [10000, 1000, 100, 10, 1]
. La marge intérieure kValid
calcule les valeurs minimales sur deux fenêtres valides: [10000, 1000, 100]
et [100, 10, 1]
, ce qui donne la sortie [100, 1]
. La marge intérieure kSame
remplit d'abord le tableau de sorte que la forme après la fenêtre de réduction soit identique à celle de l'entrée pour le premier pas en ajoutant des éléments initiaux des deux côtés, ce qui donne [MAX_VALUE, 10000, 1000, 100, 10, 1,
MAX_VALUE]
. L'exécution de la fonction de réduction de la fenêtre sur le tableau rempli fonctionne sur trois fenêtres [MAX_VALUE, 10000, 1000]
, [1000, 100, 10]
, [10, 1, MAX_VALUE]
et les rendements [1000, 10, 1]
.
L'ordre d'évaluation de la fonction de réduction est arbitraire et peut être non déterministe. Par conséquent, la fonction de réduction ne doit pas être trop sensible à la réassociation. Pour en savoir plus, consultez la discussion sur l'associativité dans le contexte de Reduce
.
ReplicaId
Consultez également XlaBuilder::ReplicaId
.
Renvoie l'ID unique (scalaire U32) de l'instance répliquée.
ReplicaId()
L'ID unique de chaque instance répliquée est un entier non signé dans l'intervalle [0, N)
, où N
correspond au nombre d'instances répliquées. Étant donné que toutes les instances répliquées exécutent le même programme, un appel ReplicaId()
dans le programme renvoie une valeur différente sur chaque instance répliquée.
Remodeler
Consultez également XlaBuilder::Reshape
et l'opération Collapse
.
Remodele les dimensions d'un tableau dans une nouvelle configuration.
Reshape(operand, new_sizes)
Reshape(operand, dimensions, new_sizes)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau de type T |
dimensions |
int64 vecteur |
ordre dans lequel les dimensions sont réduites |
new_sizes |
int64 vecteur |
vecteur de tailles de nouvelles dimensions |
Sur le plan conceptuel, cette opération aplatit d'abord un tableau en un vecteur unidimensionnel de valeurs de données, puis affine ce vecteur dans une nouvelle forme. Les arguments d'entrée sont un tableau arbitraire de type T, un vecteur constante d'indices de dimension de temps de compilation, et un vecteur de constante de temps de compilation de tailles de dimension.
Les valeurs du vecteur dimension
, si elles sont fournies, doivent être une permutation de toutes les dimensions de T. Si elles ne sont pas fournies, la valeur par défaut est {0, ..., rank - 1}
. Dans dimensions
, les dimensions vont de la dimension la plus variable la plus lente (la plus majeure) à la dimension la plus variable (la plus mineure) dans l'imbrication de boucle, qui réduit le tableau d'entrée en une seule dimension. Le vecteur new_sizes
détermine la taille du tableau de sortie. La valeur à l'index 0 dans new_sizes
correspond à la taille de la dimension 0, la valeur à l'index 1 correspond à la taille de la dimension 1, et ainsi de suite. Le produit des dimensions new_size
doit être égal au produit des tailles des dimensions de l'opérande. Lorsque vous affinez le tableau réduit pour obtenir le tableau multidimensionnel défini par new_sizes
, les dimensions de new_sizes
sont classées de la variation la plus lente (les plus importantes) et la plus rapide (les plus mineures).
Par exemple, imaginons que "v" représente un tableau de 24 éléments:
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} } };
En tant que cas particulier, reshape peut transformer un tableau à un seul élément en scalaire et inversement. Par exemple :
Reshape(f32[1x1] { {5} }, {0,1}, {}) == 5;
Reshape(5, {}, {1,1}) == f32[1x1] { {5} };
Retour (inverse)
Consultez également XlaBuilder::Rev
.
Rev(operand, dimensions)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau de type T |
dimensions |
ArraySlice<int64> |
de dimensions à inverser |
Inverse l'ordre des éléments dans le tableau operand
selon le dimensions
spécifié, ce qui génère un tableau de sortie de même forme. Chaque élément du tableau d'opérandes à un indice multidimensionnel est stocké dans le tableau de sortie à un indice transformé. L'index multidimensionnel est transformé en inversant l'index dans chaque dimension à inverser (par exemple, si une dimension de taille N est l'une des dimensions inverses, son index i est transformé en N - 1 - i).
L'opération Rev
permet, par exemple, d'inverser le tableau des pondérations de convolution selon les deux dimensions de fenêtre lors du calcul du gradient dans les réseaux de neurones.
RngNormal
Consultez également XlaBuilder::RngNormal
.
Construit une sortie d'une forme donnée avec des nombres aléatoires générés selon la \(N(\mu, \sigma)\) distribution normale. Les paramètres \(\mu\) et \(\sigma\), ainsi que la forme de sortie, doivent être de type élémentaire à virgule flottante. De plus, les paramètres doivent avoir une valeur scalaire.
RngNormal(mu, sigma, shape)
Arguments | Type | Sémantique |
---|---|---|
mu |
XlaOp |
Scalaire de type T spécifiant la moyenne des nombres générés |
sigma |
XlaOp |
Scalaire de type T spécifiant l'écart type des valeurs générées |
shape |
Shape |
Forme de sortie de type T |
RngUniform
Consultez également XlaBuilder::RngUniform
.
Construit une sortie d'une forme donnée avec des nombres aléatoires générés après la distribution uniforme sur l'intervalle \([a,b)\). Les paramètres et le type d'élément de sortie doivent être de type booléen, intégral ou à virgule flottante, et les types doivent être cohérents. Les backends de processeur et de GPU ne sont actuellement compatibles qu'avec F64, F32, F16, BF16, S64, U64, S32 et U32. De plus, les paramètres doivent avoir une valeur scalaire. Si la valeur est \(b <= a\) , le résultat est défini par l'implémentation.
RngUniform(a, b, shape)
Arguments | Type | Sémantique |
---|---|---|
a |
XlaOp |
Scalaire de type T spécifiant la limite inférieure de l'intervalle |
b |
XlaOp |
Scalaire de type T spécifiant la limite supérieure de l'intervalle |
shape |
Shape |
Forme de sortie de type T |
RngBitGenerator
Génère une sortie avec une forme donnée remplie de bits aléatoires uniformes à l'aide de l'algorithme spécifié (ou du backend par défaut), puis renvoie un état mis à jour (avec la même forme que l'état initial) et les données aléatoires générées.
L'état initial est l'état initial de la génération actuelle de nombres aléatoires. Cette valeur, ainsi que la forme et les valeurs valides requises, dépendent de l'algorithme utilisé.
Il est garanti que le résultat sera une fonction déterministe de l'état initial, mais il n'est pas garanti qu'il soit déterministe entre les backends et les différentes versions du compilateur.
RngBitGenerator(algorithm, key, shape)
Arguments | Type | Sémantique |
---|---|---|
algorithm |
RandomAlgorithm |
Algorithme PRNG à utiliser. |
initial_state |
XlaOp |
État initial de l'algorithme PRNG. |
shape |
Shape |
Forme de sortie pour les données générées. |
Valeurs disponibles pour algorithm
:
rng_default
: algorithme spécifique au backend avec des exigences de forme spécifiques au backend.rng_three_fry
: algorithme de PRNG basé sur un compteur ThreeFry. La formeinitial_state
estu64[2]
avec des valeurs arbitraires. Salmon et al. SC 2011. Des nombres aléatoires parallèles: c'est facile.rng_philox
: algorithme de Philox permettant de générer des nombres aléatoires en parallèle La formeinitial_state
estu64[3]
avec des valeurs arbitraires. Salmon et al. SC 2011. Des nombres aléatoires parallèles: c'est facile.
Nuage de points
L'opération de dispersion XLA génère une séquence de résultats qui sont les valeurs du tableau d'entrée operands
, avec plusieurs tranches (aux indices spécifiés par scatter_indices
) mises à jour avec la séquence de valeurs dans updates
à l'aide de update_computation
.
Consultez également XlaBuilder::Scatter
.
scatter(operands..., scatter_indices, updates..., update_computation,
index_vector_dim, update_window_dims, inserted_window_dims,
scatter_dims_to_operand_dims)
Arguments | Type | Sémantique |
---|---|---|
operands |
Séquence de N XlaOp |
N tableaux de types T_0, ..., T_N dans lesquels la diffusion doit être effectuée. |
scatter_indices |
XlaOp |
Tableau contenant les index de départ des tranches sur lesquelles la distribution doit être effectuée. |
updates |
Séquence de N XlaOp |
N tableaux de types T_0, ..., T_N . updates[i] contient les valeurs à utiliser pour diffuser operands[i] . |
update_computation |
XlaComputation |
Calcul à utiliser pour combiner les valeurs existantes dans le tableau d'entrée et les mises à jour lors de la diffusion. Ce calcul doit être de type T_0, ..., T_N, T_0, ..., T_N -> Collate(T_0, ..., T_N) . |
index_vector_dim |
int64 |
Dimension de scatter_indices contenant les index de départ. |
update_window_dims |
ArraySlice<int64> |
Ensemble des dimensions de la forme updates correspondant aux dimensions de la fenêtre. |
inserted_window_dims |
ArraySlice<int64> |
Ensemble des dimensions de la fenêtre à insérer dans la forme updates . |
scatter_dims_to_operand_dims |
ArraySlice<int64> |
Carte de dimensions entre les indices de dispersion et l'espace des indices d'opérande. Ce tableau est interprété comme si vous mappiez i à scatter_dims_to_operand_dims[i] . Il doit s'agir d'un échange individuel et total. |
indices_are_sorted |
bool |
Indique si le tri des index est garanti par l'appelant. |
unique_indices |
bool |
Indique si l'appelant garantit l'unicité des index. |
Où :
- N doit être supérieur ou égal à 1.
operands
[0
], ...,operands
[N-1
] doivent tous avoir les mêmes dimensions.updates
[0
], ...,updates
[N-1
] doivent tous avoir les mêmes dimensions.- Si
N = 1
,Collate(T)
estT
. - Si la valeur est
N > 1
,Collate(T_0, ..., T_N)
est un tuple d'élémentsN
de typeT
.
Si index_vector_dim
est égal à scatter_indices.rank
, nous considérons implicitement que scatter_indices
a une dimension 1
à la fin.
Nous définissons update_scatter_dims
de type ArraySlice<int64>
comme l'ensemble des dimensions dans la forme updates
qui ne sont pas dans update_window_dims
, par ordre croissant.
Les arguments de dispersion doivent respecter les contraintes suivantes:
Chaque tableau
updates
doit être de rangupdate_window_dims.size + scatter_indices.rank - 1
.Les limites de la dimension
i
dans chaque tableauupdates
doivent respecter ce qui suit:- Si
i
est présent dansupdate_window_dims
(c'est-à-dire égal àupdate_window_dims
[k
] pour certainsk
), la limite de la dimensioni
dansupdates
ne doit pas dépasser la limite correspondante deoperand
après avoir pris en compteinserted_window_dims
(par exemple,adjusted_window_bounds
[k
], oùadjusted_window_bounds
contient les limites deoperand
avec les limites des indexinserted_window_dims
supprimées). - Si
i
est présent dansupdate_scatter_dims
(c'est-à-dire égal àupdate_scatter_dims
[k
] pour certainesk
), alors la limite de la dimensioni
dansupdates
doit être égale à la limite correspondante descatter_indices
, en ignorantindex_vector_dim
(scatter_indices.shape.dims
[k
], sik
<index_vector_dim
etscatter_indices.shape.dims
[k+1
] dans le cas contraire).
- Si
update_window_dims
doit être dans l'ordre croissant, ne pas comporter de numéros de dimensions répétés et être compris dans la plage[0, updates.rank)
.inserted_window_dims
doit être dans l'ordre croissant, ne pas comporter de numéros de dimensions répétés et être compris dans la plage[0, operand.rank)
.operand.rank
doit être égal à la somme deupdate_window_dims.size
etinserted_window_dims.size
.scatter_dims_to_operand_dims.size
doit être égal àscatter_indices.shape.dims
[index_vector_dim
], et ses valeurs doivent être comprises dans la plage[0, operand.rank)
.
Pour un indice U
donné dans chaque tableau updates
, l'index I
correspondant dans le tableau operands
correspondant auquel cette mise à jour doit être appliquée est calculé comme suit:
- Soit
G
= {U
[k
] pourk
dansupdate_scatter_dims
}. UtilisezG
pour rechercher un vecteur d'indiceS
dans le tableauscatter_indices
, de sorte queS
[i
] =scatter_indices
[Combine(G
,i
)], où Combine(A, b) insère b aux positionsindex_vector_dim
dans A. - Créez un index
S
in
dansoperand
à l'aide deS
en dispersantS
à l'aide de la cartescatter_dims_to_operand_dims
. Plus formellement :S
in
[scatter_dims_to_operand_dims
[k
]] =S
[k
] sik
<scatter_dims_to_operand_dims.size
.S
in
[_
] =0
dans les autres cas.
- Créez un indice
W
in
dans chaque tableauoperands
en dispersant les indices au niveau deupdate_window_dims
dansU
en fonction deinserted_window_dims
. Plus formellement :W
in
[window_dims_to_operand_dims
(k
)] =U
[k
] sik
se trouve dansupdate_window_dims
, oùwindow_dims_to_operand_dims
est la fonction monotone avec le domaine [0
,update_window_dims.size
) et la plage [0
,operand.rank
) \inserted_window_dims
. (Par exemple, siupdate_window_dims.size
est4
,operand.rank
est6
etinserted_window_dims
est {0
,2
},window_dims_to_operand_dims
est {0
→1
,1
→3
,2
→4
,3
→5
}.)W
in
[_
] =0
dans les autres cas.
I
correspond àW
in
+S
in
, où + correspond à l'addition au niveau de l'élément.
En résumé, l'opération de dispersion peut être définie comme suit.
- Initialiser
output
avecoperands
, c'est-à-dire pour tous les indexJ
, pour tous les indicesO
du tableauoperands
[J
] :
output
[J
][O
] =operands
[J
][O
] - Pour chaque indice
U
dans le tableauupdates
[J
] et l'index correspondantO
dans le tableauoperand
[J
], siO
est un indice valide pouroutput
:
(output
[0
][O
], ...,output
[N-1
][O
]) =update_computation
(output
[0
][O
], ..., ,output
[N-1
][O
],updates
[0
][U
], ...,updates
[N-1
][U
])
L'ordre dans lequel les mises à jour sont appliquées n'est pas déterministe. Ainsi, lorsque plusieurs indices dans updates
font référence au même indice dans operands
, la valeur correspondante dans output
sera non déterministe.
Notez que le premier paramètre transmis dans update_computation
sera toujours la valeur actuelle du tableau output
, et que le deuxième paramètre correspondra toujours à la valeur du tableau updates
. Ceci est particulièrement important dans les cas où update_computation
n'est pas commutatif.
Si indices_are_sorted
est défini sur "true", XLA peut supposer que les start_indices
sont triées (dans l'ordre croissant des start_index_map
) par l'utilisateur. Si ce n'est pas le cas, la sémantique est définie par l'implémentation.
Si unique_indices
est défini sur "true", XLA peut supposer que tous les éléments diffusés sont uniques. XLA peut donc utiliser des opérations non atomiques. Si unique_indices
est défini sur "true" et que les index dispersés ne sont pas uniques, la sémantique est définie pour l'implémentation.
De manière informelle, l'opération de dispersion peut être considérée comme l'inverse de l'opération de collecte, c'est-à-dire que l'opération de dispersion met à jour les éléments de l'entrée extraits par l'opération de collecte correspondante.
Pour obtenir une description informelle détaillée et des exemples, consultez la section "Description informelle" sous Gather
.
Sélectionnez
Consultez également XlaBuilder::Select
.
Construit un tableau de sortie à partir des éléments de deux tableaux d'entrée, à partir des valeurs d'un tableau de prédicat.
Select(pred, on_true, on_false)
Arguments | Type | Sémantique |
---|---|---|
pred |
XlaOp |
tableau de type PRED |
on_true |
XlaOp |
tableau de type T |
on_false |
XlaOp |
tableau de type T |
Les tableaux on_true
et on_false
doivent avoir la même forme. Il s'agit également de la forme du tableau de sortie. Le tableau pred
doit avoir la même dimensionnalité que on_true
et on_false
, avec le type d'élément PRED
.
Pour chaque élément P
de pred
, l'élément correspondant du tableau de sortie est extrait de on_true
si la valeur de P
est true
, et de on_false
si la valeur de P
est false
. En tant que forme limitée de diffusion, pred
peut être une valeur scalaire de type PRED
. Dans ce cas, le tableau de sortie est entièrement extrait de on_true
si pred
est true
, et de on_false
si pred
est false
.
Exemple avec un pred
non scalaire:
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};
Exemple avec une valeur scalaire de 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};
La sélection entre les tuples est acceptée. À cette fin, les tuples sont considérés comme des types scalaires. Si on_true
et on_false
sont des tuples (qui doivent avoir la même forme), pred
doit être un scalaire de type PRED
.
SelectAndScatter
Consultez également XlaBuilder::SelectAndScatter
.
Cette opération peut être considérée comme une opération composite qui calcule d'abord ReduceWindow
sur le tableau operand
pour sélectionner un élément dans chaque fenêtre, puis disperse le tableau source
aux index des éléments sélectionnés pour construire un tableau de sortie ayant la même forme que le tableau d'opérande. La fonction binaire select
permet de sélectionner un élément de chaque fenêtre en l'appliquant à chaque fenêtre. Elle est appelée avec la propriété selon laquelle le vecteur d'index du premier paramètre est lexicographiquement inférieur à celui du deuxième paramètre. La fonction select
renvoie true
si le premier paramètre est sélectionné et false
si le deuxième paramètre est sélectionné. La fonction doit maintenir la transitivité (par exemple, si select(a, b)
et select(b, c)
sont true
, select(a, c)
est également true
) afin que l'élément sélectionné ne dépend pas de l'ordre des éléments traversés pour une fenêtre donnée.
La fonction scatter
est appliquée à chaque index sélectionné dans le tableau de sortie. Elle utilise deux paramètres scalaires:
- Valeur actuelle à l'index sélectionné dans le tableau de sortie
- Valeur de dispersion de
source
qui s'applique à l'index sélectionné
Elle combine les deux paramètres et renvoie une valeur scalaire utilisée pour mettre à jour la valeur au niveau de l'index sélectionné dans le tableau de sortie. Au départ, tous les index du tableau de sortie sont définis sur init_value
.
Le tableau de sortie a la même forme que le tableau operand
, et le tableau source
doit avoir la même forme que le résultat de l'application d'une opération ReduceWindow
au tableau operand
. SelectAndScatter
peut être utilisé pour rétropropager les valeurs de gradient d'une couche de pooling dans un réseau de neurones.
SelectAndScatter(operand, select, window_dimensions, window_strides,
padding, source, init_value, scatter)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau de type T sur lequel les fenêtres glissent |
select |
XlaComputation |
Calcul binaire de type T, T -> PRED , à appliquer à tous les éléments de chaque fenêtre ; renvoie true si le premier paramètre est sélectionné et false si le deuxième paramètre est sélectionné |
window_dimensions |
ArraySlice<int64> |
tableau d'entiers pour les valeurs des dimensions de fenêtre |
window_strides |
ArraySlice<int64> |
tableau d'entiers pour les valeurs de pas de fenêtre |
padding |
Padding |
Type de marge intérieure pour la fenêtre (Padding::kSame ou Padding::kValid) |
source |
XlaOp |
tableau de type T avec les valeurs à disperser |
init_value |
XlaOp |
valeur scalaire de type T pour la valeur initiale du tableau de sortie |
scatter |
XlaComputation |
Calcul binaire de type T, T -> T , pour appliquer chaque élément source de dispersion avec son élément de destination |
La figure ci-dessous montre des exemples d'utilisation de SelectAndScatter
, la fonction select
calculant la valeur maximale parmi ses paramètres. Notez que lorsque les fenêtres se chevauchent, comme dans la figure (2) ci-dessous, un index du tableau operand
peut être sélectionné plusieurs fois par différentes fenêtres. Sur la figure, l'élément de valeur 9 est sélectionné par les deux fenêtres supérieures (bleu et rouge), et la fonction d'addition binaire scatter
génère l'élément de sortie de valeur 8 (2 + 6).
L'ordre d'évaluation de la fonction scatter
est arbitraire et peut être non déterministe. Par conséquent, la fonction scatter
ne doit pas être trop sensible à la réassociation. Pour en savoir plus, consultez la discussion sur l'associativité dans le contexte de Reduce
.
Envoyer
Consultez également XlaBuilder::Send
.
Send(operand, channel_handle)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
données à envoyer (tableau de type T) |
channel_handle |
ChannelHandle |
un identifiant unique pour chaque paire envoi/réception |
Envoie les données d'opérande données à une instruction Recv
dans un autre calcul partageant le même handle de canal. Ne renvoie aucune donnée.
Comme pour l'opération Recv
, l'API cliente de l'opération Send
représente une communication synchrone. Elle est décomposée en interne en deux instructions HLO (Send
et SendDone
) pour permettre les transferts de données asynchrones. Consultez également HloInstruction::CreateSend
et HloInstruction::CreateSendDone
.
Send(HloInstruction operand, int64 channel_id)
Lance un transfert asynchrone de l'opérande vers les ressources allouées par l'instruction Recv
avec le même ID de canal. Renvoie un contexte qui est utilisé par une instruction SendDone
suivante pour attendre la fin du transfert de données. Le contexte est un tuple d'{opérande (forme), d'identifiant de requête (U32)} et il ne peut être utilisé que par une instruction SendDone
.
SendDone(HloInstruction context)
Compte tenu d'un contexte créé par une instruction Send
, attend la fin du transfert de données. L'instruction ne renvoie aucune donnée.
Programmation des instructions pour la chaîne
L'ordre d'exécution des quatre instructions pour chaque canal (Recv
, RecvDone
, Send
, SendDone
) est le suivant.
Recv
a lieu avant leSend
Send
a lieu avant leRecvDone
Recv
a lieu avant leRecvDone
Send
a lieu avant leSendDone
Lorsque les compilateurs backend génèrent une planification linéaire pour chaque calcul qui communique via des instructions de canal, il ne doit pas y avoir de cycles entre les calculs. Par exemple, les programmations ci-dessous entraînent des interblocages.
Tranche
Consultez également XlaBuilder::Slice
.
Le découpage extrait un sous-tableau du tableau d'entrée. Le sous-tableau a le même rang que l'entrée et contient les valeurs à l'intérieur d'un cadre de délimitation au sein du tableau d'entrée, où les dimensions et les index du cadre de délimitation sont fournis en tant qu'arguments pour l'opération de tranche.
Slice(operand, start_indices, limit_indices, strides)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
Tableau à N dimensions de type T |
start_indices |
ArraySlice<int64> |
Liste de N entiers contenant les index de départ de la tranche pour chaque dimension. Les valeurs doivent être supérieures ou égales à zéro. |
limit_indices |
ArraySlice<int64> |
Liste de N entiers contenant les index de fin (exclus) de la tranche pour chaque dimension. Chaque valeur doit être supérieure ou égale à la valeur start_indices de la dimension correspondante et inférieure ou égale à la taille de la dimension. |
strides |
ArraySlice<int64> |
Liste de N entiers qui décide du pas d'entrée de la tranche. Le secteur sélectionne tous les éléments strides[d] de la dimension d . |
Exemple unidimensionnel:
let a = {0.0, 1.0, 2.0, 3.0, 4.0}
Slice(a, {2}, {4}) produces:
{2.0, 3.0}
Exemple en deux dimensions:
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} }
Tri
Consultez également XlaBuilder::Sort
.
Sort(operands, comparator, dimension, is_stable)
Arguments | Type | Sémantique |
---|---|---|
operands |
ArraySlice<XlaOp> |
Opérandes à trier. |
comparator |
XlaComputation |
Le calcul de comparaison à utiliser. |
dimension |
int64 |
Dimension selon laquelle trier. |
is_stable |
bool |
Indique si le tri stable doit être utilisé. |
Si un seul opérande est fourni:
Si l'opérande est un Tensor de rang 1 (un tableau), le résultat est un tableau trié. Si vous souhaitez trier le tableau par ordre croissant, le comparateur doit effectuer une comparaison "inférieur à". Officiellement, une fois le tableau trié, il contient toutes les positions d'index
i, j
aveci < j
comparator(value[i], value[j]) = comparator(value[j], value[i]) = false
oucomparator(value[i], value[j]) = true
.Si l'opérande a un rang plus élevé, il est trié en fonction de la dimension fournie. Par exemple, pour un Tensor de rang 2 (une matrice), une valeur de dimension de
0
trie indépendamment chaque colonne, et une valeur de dimension1
trie chaque ligne indépendamment. Si aucun numéro de dimension n'est fourni, la dernière dimension est choisie par défaut. Pour la dimension triée, le même ordre de tri s'applique que dans le cas RANK-1.
Si des opérandes n > 1
sont fournis:
Tous les opérandes
n
doivent être des Tensors de mêmes dimensions. Les types d'éléments des Tensors peuvent être différents.Tous les opérandes sont triés ensemble, et non individuellement. Conceptuellement, les opérandes sont traités comme un tuple. Lorsque vous vérifiez si les éléments de chaque opérande aux positions d'index
i
etj
doivent être échangés, le comparateur est appelé avec des paramètres scalaires2 * n
, où le paramètre2 * k
correspond à la valeur à la positioni
de l'opérandek-th
et le paramètre2 * k + 1
à la valeur à la positionj
de l'opérandek-th
. En règle générale, le comparateur doit comparer les paramètres2 * k
et2 * k + 1
entre eux, et éventuellement d'autres paires de paramètres comme éléments de départ.Le résultat est un tuple qui se compose des opérandes dans l'ordre trié (le long de la dimension fournie, comme ci-dessus). L'opérande
i-th
du tuple correspond à l'opérandei-th
de tri.
Par exemple, s'il existe trois opérandes operand0 = [3, 1]
, operand1 = [42, 50]
et operand2 = [-3.0, 1.1]
, et que le comparateur ne compare que les valeurs de operand0
avec "inférieur à", la sortie du tri est le tuple ([1, 3], [50, 42], [1.1, -3.0])
.
Si is_stable
est défini sur "true", le tri est stable, c'est-à-dire que si des éléments sont considérés comme égaux par le comparateur, l'ordre relatif des valeurs égales est conservé. Deux éléments e1
et e2
sont égaux si et seulement si comparator(e1, e2) = comparator(e2, e1) = false
. Par défaut, is_stable
est défini sur "false".
Transpose
Consultez également l'opération tf.reshape
.
Transpose(operand)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
Opérande à transposer. |
permutation |
ArraySlice<int64> |
Découvrez comment activer les dimensions. |
Permute les dimensions de l'opérande avec la permutation donnée, donc ∀ i . 0 ≤ i < rank ⇒ input_dimensions[permutation[i]] = output_dimensions[i]
.
Identique à Reshape(opérande, permutation, Permute(permutation, operand.shape.dimensions)).
TriangularSolve
Consultez également XlaBuilder::TriangularSolve
.
Résolution de systèmes d'équations linéaires comportant des matrices de coefficient triangulaire inférieur ou supérieur par substitution ou inverse. En diffusant cette fonction le long des dimensions principales, cette routine résout l'un des systèmes matriciels op(a) * x =
b
, ou x * op(a) = b
, pour la variable x
, en fonction de a
et b
, où op(a)
est soit op(a) = a
, soit op(a) = Transpose(a)
, ou op(a) = Conj(Transpose(a))
.
TriangularSolve(a, b, left_side, lower, unit_diagonal, transpose_a)
Arguments | Type | Sémantique |
---|---|---|
a |
XlaOp |
Un tableau rang > 2 de type complexe ou à virgule flottante de forme [..., M, M] |
b |
XlaOp |
un tableau rang > 2 du même type ayant la forme [..., M, K] si la valeur de left_side est "true", [..., K, M] dans le cas contraire. |
left_side |
bool |
indique s'il faut résoudre un système au format op(a) * x = b (true ) ou x * op(a) = b (false ). |
lower |
bool |
utiliser le triangle supérieur ou inférieur de a . |
unit_diagonal |
bool |
Si la valeur est true , les éléments diagonals de a sont considérés comme 1 et ne sont pas accessibles. |
transpose_a |
Transpose |
choisir d'utiliser a tel quel, de le transposer ou de transposer son conjugué. |
Les données d'entrée sont lues uniquement à partir du triangle inférieur/supérieur de a
, en fonction de la valeur de lower
. Les valeurs de l'autre triangle sont ignorées. Les données de sortie sont renvoyées dans le même triangle. Les valeurs de l'autre triangle sont définies par l'implémentation et peuvent être n'importe quelle valeur.
Si les rangs de a
et b
sont supérieurs à 2, ils sont traités comme des lots de matrices, où toutes, à l'exception des deux dimensions mineures, sont des dimensions de lot. a
et b
doivent avoir les mêmes dimensions de lot.
Tuple
Consultez également XlaBuilder::Tuple
.
Un tuple contenant un nombre variable de poignées de données, chacun ayant sa propre forme.
Cette méthode est semblable à std::tuple
en C++. D'un point de vue conceptuel:
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);
Les Tuples peuvent être déstructurés (accessibles) via l'opération GetTupleElement
.
Bien que
Consultez également XlaBuilder::While
.
While(condition, body, init)
Arguments | Type | Sémantique |
---|---|---|
condition |
XlaComputation |
XlaComputation de type T -> PRED qui définit la condition d'arrêt de la boucle. |
body |
XlaComputation |
XlaComputation de type T -> T qui définit le corps de la boucle. |
init |
T |
Valeur initiale pour le paramètre condition et body . |
Exécute de manière séquentielle body
jusqu'à ce que condition
échoue. Cette méthode est semblable à une boucle "when" classique dans de nombreux autres langages, à l'exception des différences et des restrictions listées ci-dessous.
- Un nœud
While
renvoie une valeur de typeT
, qui est le résultat de la dernière exécution debody
. - La forme du type
T
est déterminée de manière statique et doit être identique dans toutes les itérations.
Les paramètres T des calculs sont initialisés avec la valeur init
lors de la première itération et sont automatiquement mis à jour vers le nouveau résultat de body
à chaque itération suivante.
L'un des principaux cas d'utilisation du nœud While
consiste à mettre en œuvre l'exécution répétée de l'entraînement dans les réseaux de neurones. Un pseudo-code simplifié est illustré ci-dessous avec un graphique représentant le calcul. Vous trouverez ce code dans while_test.cc
.
Le type T
dans cet exemple est un Tuple
composé d'un int32
pour le nombre d'itérations et d'un vector[10]
pour l'accumulateur. Pendant 1 000 itérations, la boucle continue d'ajouter un vecteur constant à l'accumulateur.
// 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};
}