Vous trouverez ci-dessous la sémantique des opérations définies dans l'interface XlaBuilder
. En règle générale, ces opérations sont mappées de manière individuelle aux opérations définies dans l'interface RPC dans 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 (tel qu'un nombre à virgule flottante 32 bits). Dans l'ensemble de la documentation, le terme tableau est utilisé pour désigner un tableau à dimension arbitraire. Pour plus de commodité, les cas particuliers ont des noms plus spécifiques et plus familiers. Par exemple, un vecteur est un tableau à une dimension et une matrice est un tableau à deux dimensions.
AfterAll
Voir aussi XlaBuilder::AfterAll
.
AfterAll prend un nombre variable de jetons et produit un seul jeton. Les jetons sont des types primitifs qui peuvent être enfilés entre des opérations à effet secondaire pour appliquer l'ordre. AfterAll
peut être utilisé comme jointure de jetons pour organiser une opération après une opération d'ensemble.
AfterAll(operands)
Arguments | Type | Sémantique |
---|---|---|
operands |
XlaOp |
nombre variable de jetons |
AllGather
Voir aussi XlaBuilder::AllGather
.
Effectue la concaténation entre les réplicas.
AllGather(operand, all_gather_dim, shard_count, replica_group_ids,
channel_id)
Arguments | Type | Sémantique |
---|---|---|
operand
|
XlaOp
|
Tableau à concaténer entre les réplicas |
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
|
int64 facultatif
|
ID de canal facultatif pour la communication intermodule |
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 situent dans le résultat.replica_groups
doit être vide (auquel cas toutes les instances dupliquées appartiennent à un seul groupe, classées de0
àN - 1
) ou contenir le même nombre d'éléments que le nombre d'instances dupliquées. Par exemple,replica_groups = {0, 2}, {1, 3}
effectue la concaténation entre les instances dupliquées0
et2
, et1
et3
.shard_count
correspond à la taille de chaque groupe de réplication. Nous en avons besoin dans les cas oùreplica_groups
est vide.channel_id
est utilisé pour la communication inter-module : seules les opérationsall-gather
avec le mêmechannel_id
peuvent communiquer entre elles.
La forme de sortie est la forme d'entrée avec all_gather_dim
multiplié par shard_count
. Par exemple, s'il existe deux instances dupliquées et que l'opérande a respectivement les valeurs [1.0, 2.5]
et [3.0, 5.25]
sur les deux instances dupliqué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 dupliquées.
AllReduce
Voir aussi XlaBuilder::AllReduce
.
Effectue un calcul personnalisé sur plusieurs réplicas.
AllReduce(operand, computation, replica_group_ids, channel_id)
Arguments | Type | Sémantique |
---|---|---|
operand
|
XlaOp
|
Tableau ou tuple non vide de tableaux à réduire entre les réplicas |
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
|
int64 facultatif
|
ID de canal facultatif pour la communication entre 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 de réplicas entre lesquels la réduction est effectuée (l'ID de réplica pour le réplica actuel peut être récupéré à l'aide deReplicaId
).replica_groups
doit être vide (dans ce cas, tous les réplicas appartiennent à un seul groupe) ou contenir le même nombre d'éléments que le nombre de réplicas. Par exemple,replica_groups = {0, 2}, {1, 3}
effectue une réduction entre les réplicas0
et2
, et1
et3
.channel_id
est utilisé pour la communication inter-module : seules les opérationsall-reduce
avec le mêmechannel_id
peuvent communiquer entre elles.
La forme de sortie est identique à la forme d'entrée. Par exemple, s'il existe deux instances dupliquées et que l'opérande a respectivement les valeurs [1.0, 2.5]
et [3.0, 5.25]
sur les deux instances dupliquées, la valeur de sortie de cette opération et de ce calcul de somme sera [4.0, 7.75]
sur les deux instances dupliquées. Si l'entrée est un tuple, la sortie est également un tuple.
Le calcul du résultat de AllReduce
nécessite une entrée de chaque réplica, donc si un réplica exécute un nœud AllReduce
plus de fois qu'un autre, l'ancien réplica attendra indéfiniment. Étant donné que les réplicas exécutent tous le même programme, il n'existe pas beaucoup de façons de procéder, mais cela est possible lorsque la condition d'une boucle while dépend des données de l'influx et que les données infusées entraînent l'itération de la boucle while plus de fois sur un réplica que sur un autre.
AllToAll
Voir aussi XlaBuilder::AllToAll
.
AllToAll est une opération collective qui envoie des données de tous les cœurs à tous les cœurs. Il comporte deux phases :
- Phase de dispersion. Sur chaque cœur, l'opérande est divisé en un nombre de blocs
split_count
le long desplit_dimensions
, et les blocs sont dispersés sur tous les cœurs (par exemple, le bloc i est envoyé au cœur i). - 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 d'ID de réplication participant au calcul (l'ID de réplication du réplica actuel peut être récupéré à l'aide deReplicaId
). AllToAll sera appliqué dans les sous-groupes dans l'ordre spécifié. Par exemple,replica_groups = { {1,2,3}, {4,5,0} }
signifie qu'une opération AllToAll sera appliquée dans les réplicas{1, 2, 3}
et lors de la phase de collecte. Les blocs reçus seront ensuite concaténés dans l'ordre 1, 2, 3. Ensuite, un autre AllToAll est appliqué aux instances répliquées 4, 5, 0, et l'ordre de concaténation est également 4, 5, 0. Sireplica_groups
est vide, toutes les instances dupliquées appartiennent à un groupe, dans l'ordre de concaténation de leur apparence.
Conditions préalables :
- La taille de dimension de l'opérande sur
split_dimension
est divisible parsplit_count
. - La forme de l'opérande n'est pas un 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
|
Valeur de l'intervalle [0,
n) qui nomme la dimension le long de laquelle l'opérande est divisée |
concat_dimension
|
int64
|
Valeur dans l'intervalle [0,
n) qui nomme la dimension sur laquelle les blocs de fractionnement sont concaténés |
split_count
|
int64
|
Nombre de cœurs participant à cette opération. Si replica_groups est vide, il doit s'agir du nombre de réplicas. Sinon, il doit être égal au nombre de réplicas de chaque groupe. |
replica_groups
|
Vecteur ReplicaGroup
|
Chaque groupe contient une liste d'ID d'instances répliquées. |
Vous trouverez ci-dessous un exemple avec 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'opération Alltoall. Sur chaque cœur, l'opérande est divisé en quatre parties selon la dimension 0, de sorte que chaque partie soit de forme f32[4,4]. Les quatre parties sont réparties sur tous les cœurs. Ensuite, chaque cœur concatène les parties reçues le long de la dimension 1, dans l'ordre des cœurs 0 à 4. Ainsi, la sortie de chaque cœur a la forme f32[16,4].
BatchNormGrad
Consultez également XlaBuilder::BatchNormGrad
et l'article d'origine sur la normalisation des lots pour obtenir une description détaillée de l'algorithme.
Calcule les gradients de la 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 unidimensionnel (\(\gamma\)) |
mean |
XlaOp |
Tableau à une dimension (\(\mu\)) |
variance |
XlaOp |
Tableau à une dimension (\(\sigma^2\)) |
grad_output |
XlaOp |
Dégradés 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 de caractéristique dans operand
), l'opération calcule les gradients par rapport à operand
, offset
et scale
pour toutes les autres dimensions. feature_index
doit être un indice valide pour la dimension d'éléments géographiques dans operand
.
Les trois dégradés sont définis par les formules suivantes (en supposant un tableau à quatre dimensions comme operand
et avec l'index 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 des 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
|
gradient par rapport à l'entrée operand ($\nabla x$) |
grad_scale
|
XlaOp
|
gradient par rapport à l'entrée scale ($\nabla \gamma$) |
grad_offset
|
XlaOp
|
gradient par rapport à l'entrée offset ($\nabla \beta$) |
BatchNormInference
Consultez également XlaBuilder::BatchNormInference
et l'article original sur la normalisation par lots pour obtenir une description détaillée de l'algorithme.
Normalise un tableau selon les 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 de 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. Il utilise plutôt les valeurs estimées mean
et variance
comme valeurs d'entrée. L'objectif de cette opération est de réduire la latence lors de l'inférence, d'où le nom BatchNormInference
.
La sortie est un tableau normalisé n-dimensionnel de la même forme que l'entrée operand
.
BatchNormTraining
Consultez également XlaBuilder::BatchNormTraining
et the original batch normalization paper
pour obtenir une description détaillée de l'algorithme.
Normalise un tableau selon les 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 unidimensionnel (\(\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 de toutes les autres dimensions, et utilise la moyenne et la variance pour normaliser chaque élément dans operand
. feature_index
doit être un indice valide pour la dimension d'éléments géographiques dans operand
.
L'algorithme se déroule comme suit pour chaque lot dans operand
\(x\) qui contient des éléments m
avec w
et h
comme taille des dimensions spatiales (en supposant que operand
est un tableau à quatre 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$
Normalise, étale 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 XlaOp
:
Sorties | Type | Sémantique |
---|---|---|
output
|
XlaOp
|
Tableau à N dimensions de 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
Voir aussi XlaBuilder::BitcastConvertType
.
Semblable à un tf.bitcast
dans TensorFlow, effectue une opération de cast de bits par élément à partir d'une forme de données vers une forme cible. La taille d'entrée et de sortie doit correspondre : par exemple, les éléments s32
deviennent des éléments f32
via la routine de bitcast, et un élément s32
devient quatre éléments s8
. Bitcast est implémenté en tant que conversion de bas niveau, de sorte que les machines avec différentes représentations à virgule flottante renvoient des résultats différents.
BitcastConvertType(operand, new_element_type)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau de type T avec dimensions 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 changera en fonction du ratio de la taille primitive avant et après la conversion.
Les éléments de type source et de destination ne doivent pas être des tuples.
Conversion de bitcast en type primitif de largeur différente
L'instruction HLO BitcastConvert
prend en charge le 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')
, il existe deux cas possibles.
Tout d'abord, lorsque B > B'
, la forme de sortie obtient une nouvelle dimension de taille mineure, 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 effectifs :
f16[2]{0} %output = f16[2]{0} bitcast-convert(f32[] %input)
Pour B' > B
, l'instruction exige que la dernière dimension logique de la forme d'entrée soit égale à B'/B
, et 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
Voir aussi XlaBuilder::Broadcast
.
Ajoute des dimensions à un tableau en dupliquant les données du tableau.
Broadcast(operand, broadcast_sizes)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
Le tableau à dupliquer |
broadcast_sizes |
ArraySlice<int64> |
Tailles des nouvelles dimensions |
Les nouvelles dimensions sont insérées à gauche. Autrement dit, 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}
.
Les nouvelles dimensions s'indexent dans des copies de l'opérande, c'est-à-dire
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 de forme f32[2, 3]
et toutes les valeurs dans le résultat seront 2.0f
.
BroadcastInDim
Voir aussi XlaBuilder::BroadcastInDim
.
Augmente 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> |
Tailles des dimensions de la forme cible |
broadcast_dimensions |
ArraySlice<int64> |
Dimension de la forme cible à laquelle chaque dimension de la forme opérande correspond |
Semblable à la diffusion, mais permet d'ajouter des dimensions n'importe où et d'étendre les dimensions existantes de taille 1.
Le 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 iième dimension de l'opérande est mappée à la dimension broadcast_dimension[i] de la forme de sortie. Les dimensions de operand
doivent avoir la taille 1 ou être de la même taille que la dimension de 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 diffuse ensuite 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.
Appeler
Voir aussi XlaBuilder::Call
.
Appelle un calcul avec les arguments donnés.
Call(computation, args...)
Arguments | Type | Sémantique |
---|---|---|
computation |
XlaComputation |
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
. Il est autorisé à ne pas avoir de args
.
Cholesky
Voir aussi XlaBuilder::Cholesky
.
Calcule la décomposition de Cholesky d'un lot de matrices symétriques (hermitiennes) définies positives.
Cholesky(a, lower)
Arguments | Type | Sémantique |
---|---|---|
a |
XlaOp |
Un tableau de rang > 2 d'un type complexe ou à virgule flottante. |
lower |
bool |
indiquer s'il faut utiliser le triangle supérieur ou inférieur de a . |
Si lower
est true
, calcule les matrices triangulaires inférieures l
telles que $a = l .
l^T$. Si lower
est false
, calcule les matrices triangulaires supérieures u
telles que\(a = u^T . u\).
Les données d'entrée ne sont lues que dans le 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 implémentation et peuvent prendre 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 dimensions mineures, sont des dimensions de lot.
Si a
n'est pas une valeur symétrique (hermitienne) définie positive, le résultat est défini par implémentation.
Pince
Voir aussi XlaBuilder::Clamp
.
Écrête 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 |
À partir d'un opérande et de valeurs minimale et maximale, renvoie l'opérande s'il se trouve dans la plage comprise entre la valeur minimale et la valeur maximale, sinon renvoie la valeur minimale si l'opérande est inférieur à cette plage ou la valeur maximale si l'opérande est supérieur à cette plage. Par exemple, clamp(a, x, b) = min(max(a, x), b)
.
Les trois tableaux doivent avoir la même forme. En tant que forme restreinte de diffusion, min
et/ou max
peuvent également être un scalaire de type T
.
Exemple avec les valeurs scalaires min
et max
:
let operand: s32[3] = {-1, 5, 9};
let min: s32 = 0;
let max: s32 = 6;
==>
Clamp(min, operand, max) = s32[3]{0, 5, 6};
Réduire
Consultez également XlaBuilder::Collapse
et l'opération tf.reshape
.
Réduit les dimensions d'un tableau en une seule dimension.
Collapse(operand, dimensions)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau de type T |
dimensions |
Vecteur int64 |
sous-ensemble consécutif et dans l'ordre des dimensions de T. |
La réduction remplace le sous-ensemble donné des 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 au moment de la compilation d'indices de dimension. Les index de dimensions doivent correspondre à un sous-ensemble consécutif de dimensions T, dans l'ordre (nombres de dimension faible à élevé). 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 seule nouvelle dimension, à la même position dans la séquence de dimensions que celles qu'elles remplacent, et dont la taille est égale au produit des tailles des dimensions d'origine. Le numéro de dimension le plus bas dans dimensions
est la dimension à variation la plus lente (la plus importante) dans le nid de boucles qui réduit ces dimensions, et le numéro de dimension le plus élevé est la variation la plus rapide (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, supposons que v soit 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
Voir aussi XlaBuilder::CollectivePermute
.
CollectivePermute est une opération collective qui envoie et reçoit des données répliquées sur plusieurs instances répliquées.
CollectivePermute(operand, source_target_pairs)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau d'entrée à N dimensions |
source_target_pairs |
Vecteur <int64, int64> |
Liste de paires (source_replica_id, target_replica_id). Pour chaque paire, l'opérande est envoyé du réplica source au réplica cible. |
Notez que les restrictions suivantes s'appliquent à source_target_pair
:
- Aucune paire ne doit avoir le même ID de réplication cible ni le même ID de réplication source.
- Si un ID de réplication n'est pas une cible dans aucune paire, la sortie de ce réplica est un tenseur composé de zéros de la même forme que l'entrée.
Concatenate
Voir aussi XlaBuilder::ConcatInDim
.
La concaténation compose un tableau à partir de plusieurs opérateurs de tableau. Le tableau est du même rang que chacun des opérandes de tableau d'entrée (qui doivent être du même rang les uns que les autres) 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 avec des dimensions [L0, L1, ...]. Nécessite N >= 1. |
dimension |
int64 |
Valeur comprise dans l'intervalle [0, N) qui nomme la dimension à concatenar entre les operands . |
À l'exception de dimension
, toutes les dimensions doivent être identiques. En effet, XLA n'est pas compatible avec les tableaux "découpés". Notez également que les valeurs de rang 0 ne peuvent pas être concaténées (car il est impossible de nommer la dimension le long de laquelle la concaténation se produit).
Exemple à une dimension :
Concat({ {2, 3}, {4, 5}, {6, 7} }, 0)
>>> {2, 3, 4, 5, 6, 7}
Exemple bidimensionnel:
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 :
Conditionnel
Voir aussi 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 |
Argument de type \(T_0\) |
true_computation |
XlaComputation |
XlaComputation de type \(T_0 \to S\) |
false_operand |
XlaOp |
Argument de type \(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.
Le 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 de true_computation
et false_computation
doit être le même.
Notez qu'un seul des champs true_computation
et false_computation
sera exécuté 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 un S32
inférieur à 0 ou supérieur ou égal à 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 est appelé avec branch_operands[b]
, qui doit être du même type. Le type de la valeur renvoyée par chaque branch_computations[b]
doit être le même.
Notez qu'un seul des branch_computations
sera exécuté en fonction de la valeur de branch_index
.
Conv (convolution)
Voir aussi XlaBuilder::Conv
.
Comme ConvWithGeneralPadding, mais le remplissage est spécifié de manière abrégée comme "SAME" ou "VALID". Une telle 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 sans tenir compte du pas de progression. La valeur "VALID" signifie simplement qu'il n'y a pas de remplissage.
ConvWithGeneralPadding (convolution)
Voir aussi 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-dimensionnelle se déplaçant sur une zone de base n-dimensionnelle, et un calcul est effectué pour chaque position possible de la fenêtre.
Arguments | Type | Sémantique |
---|---|---|
lhs |
XlaOp |
Tableau d'entrées de rang n+2 |
rhs |
XlaOp |
tableau des pondérations du noyau de rang n+2 |
window_strides |
ArraySlice<int64> |
Tableau n-d de pas de noyau |
padding |
ArraySlice< pair<int64,int64>> |
Tableau n-d de marges intérieures (faible, élevée) |
lhs_dilation |
ArraySlice<int64> |
Tableau de facteurs de dilatation de la gauche supérieure n-d |
rhs_dilation |
ArraySlice<int64> |
Tableau du facteur de dilatation de la partie droite de la matrice n-d |
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 la zone de base. C'est ce qu'on appelle l'entrée, même si, bien sûr, le membre de droite est également une entrée. Dans un réseau de neurones, il s'agit des activations d'entrée.
Les n+2 dimensions 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) de 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 la fenêtre se déplace.
L'argument rhs
est un tableau de rang n+2 décrivant le filtre/noyau/la fenêtre convolutif. Les dimensions sont les suivantes, dans l'ordre :
output-z
: dimensionz
de la sortie.input-z
: la taille de cette dimension multipliée parfeature_group_count
doit être égale à la taille de la dimensionz
en lh.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 longueur de la fenêtre de convolution dans les dimensions spatiales. Par exemple, si la foulée dans 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. Sa valeur absolue 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 de la dimension y
et padding[1]
spécifie la marge intérieure de la dimension x
. Chaque paire comporte la marge intérieure basse comme premier élément et la marge intérieure haute comme deuxième élément. La marge intérieure basse est appliquée dans la direction des indices inférieurs, tandis que la marge intérieure haute est appliquée dans la direction des indices supérieurs. Par exemple, si padding[1]
est (2,3)
, la deuxième dimension spatiale sera entourée de deux zéros à gauche et de trois zéros à droite. 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 dilatation à appliquer à la partie gauche et à la partie droite, respectivement, dans chaque dimension spatiale. Si le facteur de dilatation d'une dimension spatiale est d, alors 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 d'une valeur sans opération, ce qui signifie des zéros pour la convolution.
La dilatation du membre droit est également appelée convolution atrous. Pour en savoir plus, consultez tf.nn.atrous_conv2d
. La dilation des lhs est également appelée convolution
avec transposition. 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 des éléments géographiques d'entrée et de sortie. Si feature_group_count
est supérieur à 1, cela signifie que, conceptuellement, la dimension des caractéristiques d'entrée et de sortie et la dimension des caractéristiques de sortie rhs
sont réparties uniformément 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 groupes i sont utilisés ensemble pour 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 la convolution en profondeur, l'argument feature_group_count
est défini sur la dimension des caractéristiques d'entrée, et le filtre est restructuré 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 de lot de sortie doit être de taille input batch
/ batch_group_count
. batch_group_count
doit être un diviseur de la taille des éléments géographiques de sortie.
Les dimensions de la forme de sortie sont les suivantes, dans l'ordre suivant:
batch
: la taille de cette dimension multipliée parbatch_group_count
doit être égale à la taille de la dimensionbatch
exprimée en lh.z
: taille identique àoutput-z
sur le noyau (rhs
).spatial_dims
: une valeur pour chaque emplacement valide de la fenêtre de convolution.
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 nous concaténons la sortie selon 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 de convolution sont déterminés par les pas et la taille de la zone de base après le rembourrage.
Pour décrire ce qu'est une convolution, considérez une convolution 2D et choisissez des coordonnées batch
, z
, y
et x
fixes dans la sortie. (y,x)
correspond alors à la position d'un angle de la fenêtre dans la zone de base (par exemple, l'angle supérieur gauche, selon la façon dont vous interprétez les dimensions spatiales). Nous disposons maintenant d'une fenêtre bidimensionnelle, issue de la zone de la base, où chaque point 2D est associé à un vecteur unidimensionnel. Nous obtenons donc une boîte 3D. À partir du noyau de convolution, comme nous avons fixé la coordonnée de sortie z
, nous avons également une boîte 3D. Les deux boîtes ayant les mêmes dimensions, nous pouvons prendre la somme des produits par élément entre les deux boîtes (comme 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 cinq valeurs dans la sortie dans la dimension z
de la sortie. Ces valeurs diffèrent selon la partie du noyau de convolution utilisée. Une boîte 3D de valeurs distincte est utilisée pour chaque coordonnée output-z
. Vous pouvez donc considérer qu'il s'agit de cinq convolutions distinctes, chacune avec un filtre différent.
Voici un pseudo-code pour une convolution 2D avec remplissage et pas de défilement :
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
Voir aussi XlaBuilder::ConvertElementType
.
Semblable à une static_cast
par élément en C++, effectue une opération de conversion par élément d'une forme de données vers une forme cible. Les dimensions doivent correspondre, et la conversion est effectuée par élément. Par exemple, les éléments s32
deviennent des éléments f32
via une routine de conversion s32
vers f32
.
ConvertElementType(operand, new_element_type)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau de type T avec dimensions 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 tupels.
Une conversion telle que T=s32
en U=f32
exécutera une routine de conversion de type "int-to-float" de normalisation, telle que l'arrondi à l'intervalle le 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
Voir aussi 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, quelle que soit 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.
Point
Voir aussi 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 de vecteurs |
matrice [m x k] vecteur dot [k] |
vecteur [m] | multiplication à vecteur matriciel |
matrice [m x k] dot matrice [k x n] |
matrice [m x n] | multiplication matricielle |
L'opération effectue la somme des produits sur la deuxième dimension de lhs
(ou la première si elle a le rang 1) et la première dimension de rhs
. Il s'agit des 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 vecteur/matrice ou des multiplications matricielles/matrices.
DotGeneral
Voir aussi 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 |
les nombres de dimensions contractées |
Semblable à Dot, mais permet de spécifier des numéros de dimension de contrat et de lot pour lhs
et rhs
.
Champs DotDimensionNumbers | Type | Sémantique |
---|---|---|
lhs_contracting_dimensions
|
repeated int64 | lhs numéros de dimension contractuelle |
rhs_contracting_dimensions
|
Valeur répétée de type int64 | rhs numéros de dimension contractante |
lhs_batch_dimensions
|
repeated int64 | lhs numéros de dimension de lot |
rhs_batch_dimensions
|
repeated int64 | Numéros de dimension de lot rhs |
DotGeneral effectue la somme des produits sur les dimensions contractantes spécifiées dans dimension_numbers
.
Les numéros des dimensions de contraction associées à lhs
et rhs
n'ont pas besoin d'être identiques, mais doivent avoir les mêmes tailles de dimension.
Exemple avec des numéros de dimension en contraction :
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 la même taille de dimension.
Exemple avec des numéros de dimension de lot (taille de lot 2, matrices 2 x 2) :
lhs = { { {1.0, 2.0},
{3.0, 4.0} },
{ {5.0, 6.0},
{7.0, 8.0} } }
rhs = { { {1.0, 0.0},
{0.0, 1.0} },
{ {1.0, 0.0},
{0.0, 1.0} } }
DotDimensionNumbers dnums;
dnums.add_lhs_contracting_dimensions(2);
dnums.add_rhs_contracting_dimensions(1);
dnums.add_lhs_batch_dimensions(0);
dnums.add_rhs_batch_dimensions(0);
DotGeneral(lhs, rhs, dnums) -> { { {1.0, 2.0},
{3.0, 4.0} },
{ {5.0, 6.0},
{7.0, 8.0} } }
Entrée | Sortie | Sémantique |
---|---|---|
[b0, m, k] dot [b0, k, n] |
[b0, m, n] | matmul par lot |
[b0, b1, m, k] dot [b0, b1, k, n] |
[b0, b1, m, n] | matmul par lot |
Par conséquent, le numéro de dimension obtenu commence par la dimension de lot, puis par la dimension lhs
sans contrat/sans lot, et enfin par la dimension rhs
sans contrat/sans lot.
DynamicSlice
Voir aussi XlaBuilder::DynamicSlice
.
DynamicSlice extrait un sous-tableau du tableau d'entrée à l'état dynamique start_indices
. La taille de la tranche dans chaque dimension est transmise dans size_indices
, qui spécifie le point final des intervalles de tranche exclusifs dans chaque dimension : [start, start + size]. La forme de start_indices
doit être de rang == 1, avec une taille de dimension é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 indices de début 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 début + taille doit être inférieure ou égale à la taille de la dimension pour éviter d'encapsuler la taille modulo. |
Les indices de tranche efficaces 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érandes. Si la tranche est 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 bidimensionnel:
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
Voir aussi XlaBuilder::DynamicUpdateSlice
.
DynamicUpdateSlice génère un résultat correspondant à la valeur du tableau d'entrée operand
, avec une tranche update
écrasée à start_indices
.
La forme de update
détermine la forme du sous-tableau du résultat qui est mis à jour.
La forme de start_indices
doit être de rang = 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 de la tranche. Chaque dimension de la forme de mise à jour doit être strictement supérieure à zéro, et le début + la mise à jour doit être inférieur ou égal à la taille de l'opérande pour chaque dimension afin d'éviter de générer des indices de mise à jour hors limites. |
start_indices |
séquence de N XlaOp |
Liste de N entiers scalaires contenant les indices de début de la tranche pour chaque dimension. Cette valeur doit être supérieure ou égale à zéro. |
Les indices de tranche efficaces 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 le segment est 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 au niveau des éléments
Voir aussi XlaBuilder::Add
.
Un ensemble d'opérations arithmétiques binaires au niveau des éléments est accepté.
Op(lhs, rhs)
Où Op
est l'un des opérateurs suivants : Add
(addition), Sub
(soustraction), Mul
(multiplication), Div
(division), Pow
(puissance), Rem
(reste), Max
(maximum), Min
(minimum), And
(AND logique), Or
(OU logique), Xor
(XOR logique), ShiftLeft
(décalage à gauche), ShiftRightArithmetic
(décalage à droite arithmétique), ShiftRightLogical
(décalage à droite logique), Atan2
(arctangente à deux arguments) ou Complex
(combine les parties réelle et imaginaire en un nombre complexe)
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 savoir ce que signifie la compatibilité des formes. Le résultat d'une opération a une forme qui est le résultat 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 prises en charge, sauf si l'un des opérandes est un scalaire.
Lorsque Op
est Rem
, le signe du résultat est extrait du dividende, et la valeur absolue du résultat est toujours inférieure à la valeur absolue du diviseur.
Le dépassement de division d'entiers (division/reste par zéro signée/non signée ou une division/reste signée de INT_SMIN
avec -1
) génère une valeur définie par l'implémentation.
Une autre variante avec la prise en charge de la diffusion de différents rangs existe 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, pour ajouter une matrice à un vecteur).
L'opérande broadcast_dimensions
supplémentaire est une tranche d'entiers qui permet d'étendre le rang de l'opérande de rang inférieur au rang de l'opérande de rang supérieur. broadcast_dimensions
met en correspondance les dimensions de la forme de rang inférieur avec celles de la forme de rang supérieur. Les dimensions non mappées de la forme développée sont remplies de dimensions de taille 1. Le broadcasting de dimension dégénérée diffuse ensuite les formes le long de ces dimensions dégénérées pour égaliser les formes des deux opérandes. La sémantique est décrite en détail sur la page Diffusion.
Opérations de comparaison au niveau des éléments
Voir aussi XlaBuilder::Eq
.
Un ensemble d'opérations de comparaison binaires standards au niveau des éléments 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
(différent de), Ge
(supérieur ou égal à), Gt
(supérieur à), Le
(inférieur ou égal à) ou Lt
(inférieur à). Un autre ensemble d'opérateurs, EqTotalOrder, NeTotalOrder, GeTotalOrder, GtTotalOrder, LeTotalOrder et LtTotalOrder, fournit les mêmes fonctionnalités, à l'exception qu'ils acceptent également un ordre total sur les nombres à virgule flottante, en appliquant -NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.
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 de diffusion pour en savoir plus sur la compatibilité des formes. Le résultat d'une opération a une forme qui est le résultat 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 prises en charge, sauf si l'un des opérandes est un scalaire.
Une autre variante avec la prise en charge de la diffusion de différents rangs existe 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 comparer des opérations 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 la diffusion des opérandes. La sémantique est décrite en détail sur la page Diffusion.
Fonctions unaires au niveau des éléments
XlaBuilder accepte les fonctions unaires par élément suivantes :
Abs(operand)
x -> |x|
abs. par élément.
Cbrt(operand)
Opération racine cubique au niveau des éléments x -> cbrt(x)
.
Ceil(operand)
Plafond par élément x -> ⌈x⌉
.
Clz(operand)
Comptabilise les zéros au début par élément.
Cos(operand)
Cosinus élément par élément x -> cos(x)
.
Erf(operand)
Fonction d'erreur au niveau des éléments x -> erf(x)
, où
\(\text{erf}(x) = \frac{2}{\sqrt{\pi} }\int_0^x e^{-t^2} \, dt\).
Exp(operand)
x -> e^x
exponentiel naturel par élément.
Expm1(operand)
Exponentielle naturelle par élément moins un x -> e^x - 1
.
Floor(operand)
Prix plancher par é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)
Vérifie si chaque élément de operand
est fini, c'est-à-dire qu'il ne correspond pas à l'infini positif ou négatif, et qu'il n'est pas NaN
. Renvoie un tableau de valeurs PRED
de 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 x -> ln(x)
.
Log1p(operand)
Logarithme naturel décalé par élément x -> ln(1+x)
.
Logistic(operand)
Calcul de la fonction logistique par élément x ->
logistic(x)
.
Neg(operand)
Négation par élément x -> -x
.
Not(operand)
Négation logique par élément x -> !(x)
.
PopulationCount(operand)
Calcule le nombre de bits définis dans chaque élément de operand
.
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.
Round(operand)
Arrondi au niveau de l'élément en s'éloignant de zéro.
RoundNearestEven(operand)
Arrondi au niveau des éléments, qui est lié au pair le plus proche.
Rsqrt(operand)
Inverse élément par élément de l'opération racine carrée x -> 1.0 / sqrt(x)
.
Sign(operand)
: opération de signe au niveau des éléments 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
.
Sin(operand)
Sinus par élément x -> sin(x)
.
Sqrt(operand)
Racine carrée par élément x -> sqrt(x)
.
Tan(operand)
Tangente élément par élément x -> tan(x)
.
Tanh(operand)
Tangente hyperbolique par élément x -> tanh(x)
.
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 de la même forme. operand
peut être un scalaire (rang 0).
Fft
L'opération FFT XLA implémente les transformées de Fourier directe et inverse pour les entrées/sorties réelles et complexes. Les FFT multidimensionnelles sur jusqu'à trois axes sont acceptées.
Voir aussi 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 au domaine temporel des axes en cours de transformation. Cela est nécessaire en particulier pour que IRFFT redimensionne à 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 |
FFT complexe à complexe en temps réel. La forme reste inchangée. |
IFFT |
FFT inverse complexe à complexe. La forme reste inchangée. |
RFFT |
FFT directe réelle-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, ce qui omet la partie conjuguée inversée du signal transformé au-delà de la fréquence Nyquist. |
IRFFT |
FFT inverse réel-complexe (c'est-à-dire qu'elle prend un complexe et renvoie un réel). La forme de l'axe le plus interne est étendue à fft_length[-1] si fft_length[-1] est une valeur non nulle, ce qui permet d'inférer la partie du signal transformé au-delà de la fréquence Nyquist à partir de la conjugaison inverse des entrées 1 à fft_length[-1] // 2 + 1 . |
FFT multidimensionnelle
Lorsque plusieurs fft_length
sont fournis, cela revient à appliquer une cascade d'opérations FFT à chacun des axes les plus internes. Notez que pour les cas réels->complexes et complexes->réels, la transformation de l'axe le plus interne est (effectivement) effectuée en premier (RFFT, en dernier pour IRFFT). C'est pourquoi l'axe le plus interne est celui qui change de taille. Les autres transformations d'axe seront alors de type complexe->complexe.
Détails de mise en œuvre
Le processeur FFT du processeur repose sur TensorFFT d'Eigen. La FFT GPU utilise cuFFT.
Recueillir
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
Voir aussi XlaBuilder::Gather
.
Pour obtenir 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 recueillons des données. |
start_indices |
XlaOp |
Tableau contenant les index de départ des tranches collectées. |
index_vector_dim |
int64 |
Dimension dans start_indices qui "contient" les indices de début. Pour en savoir plus, consultez les sections ci-dessous. |
offset_dims |
ArraySlice<int64> |
Ensemble de dimensions dans la forme de sortie qui est décalé dans un tableau découpé à partir de l'opérande. |
slice_sizes |
ArraySlice<int64> |
slice_sizes[i] correspond aux limites du secteur de la dimension i . |
collapsed_slice_dims |
ArraySlice<int64> |
Ensemble des dimensions de chaque tranche qui sont réduites. Ces dimensions doivent avoir une taille de 1. |
start_index_map |
ArraySlice<int64> |
Carte décrivant comment mapper des index de start_indices avec des index légaux dans un opérande. |
indices_are_sorted |
bool |
Indique si les indices sont triés par l'appelant. |
Pour plus de commodité, nous étiquetons les dimensions du tableau de sortie qui ne sont pas dans offset_dims
comme batch_dims
.
Le résultat est un tableau de rang batch_dims.size
+ offset_dims.size
.
operand.rank
doit être égal à 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
possède une dimension 1
en fin de chaîne (c'est-à-dire que si start_indices
avait la forme [6,7]
et que index_vector_dim
est 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 certainesk
), nous choisissons les limites de dimension correspondantes en dehors destart_indices.shape
, en ignorantindex_vector_dim
(par exemple, choisissezstart_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 à partir deslice_sizes
après avoir pris en comptecollapsed_slice_dims
(c'est-à-dire que nous choisissonsadjusted_slice_sizes
[k
], oùadjusted_slice_sizes
correspond àslice_sizes
sans les limites de l'indicecollapsed_slice_dims
).
Formellement, l'indice d'opérande In
correspondant à un indice de sortie Out
donné est calculé comme suit :
Soit
G
= {Out
[k
] pourk
dansbatch_dims
}. UtilisezG
pour diviser un vecteurS
de sorte queS
[i
] =start_indices
[Combine(G
,i
)] où Combine(A, b) insère b à la positionindex_vector_dim
dans A. Notez que cette valeur est bien définie, même siG
est vide: siG
est vide,S
=start_indices
.Créez un indice 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
sinon.
Créez un index
O
in
dansoperand
en dispersant les indices aux 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
sinon.
In
estO
in
+S
in
, où + correspond à l'addition par élément.
remapped_offset_dims
est une fonction monotone avec un domaine [0
, offset_dims.size
) et une plage [0
, operand.rank
) \ collapsed_slice_dims
. Par exemple, Si offset_dims.size
est 4
, operand.rank
est 6
et collapsed_slice_dims
est {0
, 2
}, alors 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 start_indices
est trié (par ordre croissant, après avoir dispersé ses valeurs en fonction de 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
du tableau de sortie correspond à un élément E
du tableau d'opérandes, calculé comme suit :
Nous utilisons les dimensions de lot dans
Out
pour rechercher un indice de début à partir destart_indices
.Nous utilisons
start_index_map
pour mapper l'index de départ (dont la taille peut être inférieure à operand.rank) avec un index de départ "complet" dansoperand
.Nous extrayons de manière dynamique une tranche de taille
slice_sizes
à l'aide de l'index de début complet.Nous remodelons le segment en réduisant les dimensions
collapsed_slice_dims
. Étant donné que toutes les dimensions des secteurs réduits doivent avoir une limite de 1, ce remodelage est toujours légal.Nous utilisons les dimensions de décalage dans
Out
pour indexer cette tranche afin d'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. Les valeurs plus intéressantes pour index_vector_dim
ne modifient pas fondamentalement l'opération, mais rendent la représentation visuelle plus lourde.
Pour comprendre comment tous ces éléments s'assemblent, examinons un exemple qui rassemble cinq segments de la forme [8,6]
à partir d'un tableau [16,11]
. La position d'une tranche dans le tableau [16,11]
peut être représentée sous la forme d'un vecteur d'index de la forme S64[2]
, de sorte que l'ensemble de cinq positions peut être représenté sous la forme d'un tableau S64[5,2]
.
Le comportement de l'opération de rassemblement peut ensuite être représenté comme une transformation d'index qui prend [G
,O
0
,O
1
], un indice dans la forme de sortie, et le mappe sur un élément du tableau d'entrée comme suit :
Nous sélectionnons d'abord un vecteur (X
,Y
) à partir du tableau des indices de collecte à l'aide de G
.
L'élément du tableau de sortie à l'index [G
,O
0
,O
1
] est alors l'élément du tableau d'entrée à l'index [X
+O
0
,Y
+O
1
].
slice_sizes
est [8,6]
, qui détermine la plage d'O0
et d'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 indices de collecte peuvent être multidimensionnels. Par exemple, une version plus générale de l'exemple ci-dessus qui utilise un tableau "collecter des indices" de forme [4,5,2]
convertirait les index comme suit:
Là encore, il s'agit d'une tranche dynamique par lot G
0
et G
1
en tant que dimensions de lot. La taille de la tranche est toujours [8,6]
.
L'opération de collecte dans XLA généralise la sémantique informelle décrite ci-dessus comme suit :
Nous pouvons configurer les dimensions de la forme de sortie qui sont les dimensions de décalage (dimensions contenant
O
0
,O
1
dans le dernier exemple). Les dimensions de lot de sortie (dimensions contenantG
0
,G
1
dans le dernier exemple) sont définies comme étant les dimensions de sortie qui ne sont pas des dimensions de décalage.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 listées explicitement comme
collapsed_slice_dims
, doivent avoir une taille de tranche de1
. Étant donné qu'elles ont une taille de tranche de1
, le seul indice valide pour elles est0
, et leur suppression n'introduit aucune ambiguïté.La tranche extraite du tableau "Indices de rassemblement" ("
X
,Y
" dans le dernier exemple) peut avoir moins d'éléments que le rang du tableau d'entrée, et un mappage explicite dicte comment l'index doit être développé pour avoir le même rang que l'entrée.
Pour finir, nous utilisons (2) et (3) pour implémenter tf.gather_nd
:
G
0
et G
1
permettent de diviser un index de départ du tableau de collecte d'index comme d'habitude, sauf que l'index de départ ne contient qu'un seul élément, X
. De même, il n'y a qu'un seul indice de décalage de sortie avec la valeur O
0
. Toutefois, avant d'être utilisés comme index dans le tableau d'entrée, avant d'être utilisés comme index dans le tableau d'entrée, ils sont étendus selon " Collect Index Mapping Index Mapping Index Mapping (start_index_map
dans la description form)) et "remapped_offset_dims
X
X
0
0
0
0
0
0
0
0
0
O
O
O
O
G
G
G
G
1
1
GatherIndices
tf.gather_nd
slice_sizes
dans ce cas est [1,11]
. Intuitif, cela signifie que chaque indice X
du tableau des indices de collecte sélectionne une ligne entière, et le résultat est la concaténation de toutes ces lignes.
GetDimensionSize
Voir aussi XlaBuilder::GetDimensionSize
.
Renvoie la taille de la dimension donnée de l'opérande. L'opérande doit être au format tableau.
GetDimensionSize(operand, dimension)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
Tableau d'entrée à n dimensions |
dimension |
int64 |
Valeur de l'intervalle [0, n) qui spécifie la dimension |
SetDimensionSize
Voir aussi XlaBuilder::SetDimensionSize
.
Définit la taille dynamique de la dimension donnée de XlaOp. L'opérande doit être au format tableau.
SetDimensionSize(operand, size, dimension)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
un tableau d'entrée à N dimensions. |
size |
XlaOp |
int32 représentant la taille dynamique d'exécution. |
dimension |
int64 |
Valeur dans l'intervalle [0, n) spécifiant la dimension. |
Transmettez l'opérande en tant que résultat, avec une dimension dynamique suivie par le compilateur.
Les valeurs mises en forme sont 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
Voir aussi XlaBuilder::GetTupleElement
.
Indexe un tuple dans un tuple avec une valeur de constante au moment de la compilation.
La valeur doit être une constante au moment de la compilation afin que l'inférence de forme puisse déterminer le type de la valeur résultante.
Cette méthode est analogue à std::get<int N>(t)
en C++. Conceptuellement:
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
.
In-Feed
Voir aussi XlaBuilder::Infeed
.
Infeed(shape)
Argument | Type | Sémantique |
---|---|---|
shape |
Shape |
Format 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 streaming 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 entre les opérations Infeed. Par exemple, deux flux d'entrée dans le code ci-dessous ont une commande totale, car il existe une dépendance entre les boucles "while".
result1 = while (condition, init = init_value) {
Infeed(shape)
}
result2 = while (condition, init = result1) {
Infeed(shape)
}
Les formes de tuple imbriquées ne sont pas acceptées. Pour une forme de tuple vide, l'opération Infeed est en fait une opération sans effet et se poursuit sans lire aucune donnée de l'influx de l'appareil.
Iota
Voir aussi XlaBuilder::Iota
.
Iota(shape, iota_dimension)
Crée une valeur littérale constante sur l'appareil plutôt qu'un transfert hôte potentiellement important. Crée un tableau dont la forme est spécifiée et qui contient des valeurs commençant à zéro et augmentant de 1 le long de la dimension spécifiée. Pour les types à virgule flottante, le tableau produit est équivalent à ConvertElementType(Iota(...))
, où Iota
est de type entier et la conversion est au 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 ]]
Carte
Voir aussi 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 |
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 des dimensions de la carte |
Applique une fonction scalaire aux tableaux operands
donnés, produisant un tableau de mêmes dimensions où chaque élément est le résultat de la fonction mappée appliquée aux éléments correspondants des tableaux d'entrée.
La fonction mappée est un calcul arbitraire avec la restriction qu'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, à l'exception du type d'élément T, qui 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.
S'assure que toutes les entrées sont évaluées avant les opérateurs qui dépendent des sorties de la barrière.
Pad
Voir aussi 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 ajoutée |
padding_config |
PaddingConfig |
marge intérieure sur les deux bords (faible, haute) et entre les éléments de chaque dimension |
Élargit le tableau operand
donné en ajoutant une marge interne autour du tableau et entre les éléments du tableau avec le padding_value
donné. padding_config
spécifie la quantité de marge intérieure intérieure et de marge 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 respectivement la quantité de marge intérieure ajoutée à l'extrémité inférieure (à côté de l'indice 0) et à l'extrémité supérieure (à côté de l'indice le plus élevé) de chaque dimension. La valeur de la 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.
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. La marge intérieure se produit logiquement avant la marge extérieure. Par conséquent, en cas de marge extérieure négative, les éléments sont supprimés de l'opérande à marge intérieure.
Cette opération n'a aucun effet si les paires de marges extérieures sont toutes (0, 0) et que les valeurs de marge intérieure sont toutes égales à 0. La figure ci-dessous présente des exemples de différentes valeurs edge_padding
et interior_padding
pour un tableau à deux dimensions.
Recv
Voir aussi XlaBuilder::Recv
.
Recv(shape, channel_handle)
Arguments | Type | Sémantique |
---|---|---|
shape |
Shape |
forme des données à recevoir |
channel_handle |
ChannelHandle |
Identifiant unique pour chaque paire d'envoi/réception |
Réçoit les données de la forme donnée à partir d'une instruction Send
dans un autre calcul qui partage le même canal. Renvoie un XlaOp pour les données reçues.
L'API cliente de l'opération Recv
représente la communication synchrone.
Toutefois, l'instruction est décomposée en interne en deux instructions HLO (Recv
et RecvDone
) pour permettre les transferts de données asynchrones. Voir aussi 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
ayant 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 {receive buffer (shape), request identifier (U32)}, et il ne peut être utilisé que par une instruction RecvDone
.
RecvDone(HloInstruction context)
Compte tenu d'un contexte créé par une instruction Recv
, attend la fin du transfert de données et renvoie les données reçues.
Restreindre
Voir aussi 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 |
calcul 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é à réduire. |
Où :
- N doit être supérieur ou égal à 1.
- Le calcul doit être "à peu près" 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
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)
. La sortie 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 nombres à virgule flottante. Toutefois, si la plage des données est limitée, l'addition à virgule flottante est suffisamment proche de l'associativité pour la plupart des utilisations pratiques.
Exemples
Lorsque vous effectuez une réduction sur une dimension dans un seul tableau 1D avec des valeurs [10, 11,
12, 13]
, avec la fonction de réduction f
(il s'agit de computation
), vous pouvez effectuer le calcul suivant :
f(10, f(11, f(12, f(init_value, 13)))
mais il existe aussi de nombreuses autres possibilités,
f(init_value, f(f(10, f(init_value, 11)), f(f(init_value, 12), f(init_value, 13))))
Voici un exemple de pseudo-code approximatif de la façon dont la réduction peut être implémentée, 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'une matrice (tableau à deux dimensions). La forme a le rang 2, la dimension 0 de la taille 2 et la dimension 1 de la taille 3:
Résultats de la réduction des dimensions 0 ou 1 avec une fonction "add" :
Notez que les deux résultats de réduction sont des tableaux 1D. Le diagramme en affiche un en tant que colonne et un autre en tant que ligne, simplement pour faciliter la visualisation.
Voici un exemple plus complexe d'une matrice 3D. Son rang est de 3, la dimension 0 est de taille 4, la dimension 1 est de taille 2 et la dimension 2 est de taille 3. Pour plus de simplicité, les valeurs de 1 à 6 sont répliquées dans la dimension 0.
Comme dans l'exemple 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é regroupées sous la forme d'une valeur scalaire:
| 4 8 12 |
| 16 20 24 |
Si nous réduisons la dimension 2, nous obtenons également un tableau de rang-2 dans lequel toutes les valeurs de la dimension 2 ont été regroupées sous une forme scalaire:
| 6 15 |
| 6 15 |
| 6 15 |
| 6 15 |
Notez que l'ordre relatif entre les dimensions restantes dans l'entrée est conservé dans la sortie, mais certaines dimensions peuvent se voir attribuer de nouveaux nombres (puisque le classement change).
Nous pouvons également réduire plusieurs dimensions. Les dimensions de réduction d'ajout 0 et 1 génèrent le tableau 1D [20, 28, 36]
.
La réduction de la matrice 3D sur toutes ses dimensions produit le scalaire 84
.
Réduction variadique
Lorsque la valeur est N > 1
, l'application de la fonction "Réduire" est légèrement plus complexe, car elle est appliquée simultanément à toutes les entrées. Les opérandes sont fournies au calcul dans l'ordre suivant :
- Valeur réduite du premier opérande en cours d'exécution
- …
- Exécution de la valeur réduite pour l'opérande N'
- Valeur saisie pour le premier opérande
- …
- Valeur d'entrée pour l'opérande N
Prenons l'exemple de la fonction de réduction suivante, qui peut être utilisée pour calculer la valeur maximale et l'argmax d'un tableau unidimensionnel 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ée 1D 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 sur 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 (c'est-à-dire iota) permet de co-itérer sur les tableaux et de renvoyer un tuple contenant la valeur maximale et l'indice correspondant.
ReducePrecision
Voir aussi XlaBuilder::ReducePrecision
.
Modélise l'effet de la conversion des valeurs à virgule flottante en un format à précision inférieure (tel que IEEE-FP16) et de leur retour au format d'origine. Le nombre de bits d'exposant et de mantisse dans le format à précision inférieure peut être spécifié de manière arbitraire, bien que toutes les tailles de bits ne soient pas compatibles avec toutes les implémentations matérielles.
ReducePrecision(operand, mantissa_bits, exponent_bits)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
Tableau de type à virgule flottante T . |
exponent_bits |
int32 |
Nombre de bits d'exposant au format à précision inférieure |
mantissa_bits |
int32 |
nombre de bits de mantisse 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 avec le nombre donné de bits de mantisse (à l'aide de la sémantique "égalités à l'entier"), 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 avec une 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 les deux une mantisse nulle) et un nombre non négatif de bits de mantisse. Le nombre de bits d'exposant ou de mantisse peut dépasser la valeur correspondante pour le type T
. La partie correspondante de la conversion est alors simplement une opération sans opération.
ReduceScatter
Voir aussi XlaBuilder::ReduceScatter
.
ReduceScatter est une opération collective qui effectue efficacement une opération AllReduce, puis répartit le résultat en le divisant en blocs shard_count
le long de scatter_dimension
. L'instance répliquée i
du groupe d'instances répliquées reçoit la partition 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 entre les réplicas. |
computation |
XlaComputation |
Calcul de la réduction |
scatter_dimension |
int64 |
Dimension à répartir. |
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 |
int64 facultatif |
ID de canal facultatif pour la communication inter-modules |
- Lorsque
operand
est un tuple de tableaux, la réduction-dispersion est effectuée sur chaque élément du tuple. replica_groups
est une liste de groupes de réplicas entre lesquels la réduction est effectuée (l'ID de réplica pour le réplica actuel peut être récupéré à l'aide deReplicaId
). L'ordre des réplicas dans chaque groupe détermine l'ordre dans lequel le résultat de la réduction globale sera dispersé.replica_groups
doit être vide (dans ce cas, tous les réplicas appartiennent à un seul groupe) ou contenir le même nombre d'éléments que le nombre de réplicas. Lorsqu'il existe plusieurs groupes d'instances répliquées, ils doivent tous avoir la même taille. Par exemple,replica_groups = {0, 2}, {1, 3}
effectue une réduction entre les réplicas0
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 lorsquereplica_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
avec le mêmechannel_id
peuvent communiquer entre elles.
La forme de sortie est la forme d'entrée, avec scatter_dimension
réduit de shard_count
fois. Par exemple, s'il existe deux réplications et que l'opérande a les valeurs [1.0, 2.25]
et [3.0, 5.25]
respectivement sur les deux réplications, la valeur de sortie de cette opération où scatter_dim
est 0
sera [4.0]
pour le premier réplica et [7.5]
pour le second réplica.
ReduceWindow
Voir aussi XlaBuilder::ReduceWindow
.
Applique une fonction de réduction à tous les éléments de chaque fenêtre d'une séquence de N tableaux multidimensionnels, et produit un seul tableau multidimensionnel 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 sous la forme d'un ReduceWindow
. Comme pour Reduce
, le computation
appliqué est toujours transmis au init_values
sur la 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 |
Les N valeurs de départ de la réduction, une pour chacune des N opérandes. Pour en savoir plus, consultez la section 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 de dimension 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 la fenêtre |
padding |
Padding |
Type de marge intérieure pour la fenêtre (Padding::kSame, qui ajoute une marge intérieure pour avoir la même forme de sortie que l'entrée si la longueur de pas est de 1, ou Padding::kValid, qui n'utilise aucune marge intérieure et "arrête" la fenêtre une fois qu'elle ne convient 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
N > 1
,Collate(T_0, ..., T_{N-1})
est un tuple d'élémentsN
de type(T0,...T{N-1})
.
Le code et l'image ci-dessous illustrent un exemple d'utilisation de ReduceWindow
. L'entrée est une matrice de taille [4x6], et window_dimensions et window_stride_dimensions sont toutes deux [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 se situe à un élément de sa fenêtre adjacente. Pour spécifier qu'aucune fenêtre ne se chevauche, window_stride_dimensions doit être égal à window_dimensions. La figure ci-dessous illustre l'utilisation de deux valeurs de pas différentes. Le remplissage est appliqué à chaque dimension de l'entrée, et les calculs sont les mêmes que si l'entrée était fournie avec les dimensions qu'elle a après le remplissage.
Pour un exemple de remplissage non trivial, envisagez de calculer la valeur minimale de la fenêtre de réduction (valeur initiale : MAX_FLOAT
) avec la dimension 3
et la longueur de pas 2
sur le tableau d'entrée [10000, 1000, 100, 10, 1]
. Le remplissage kValid
calcule les valeurs minimales sur deux fenêtres valides : [10000, 1000, 100]
et [100, 10, 1]
, ce qui génère la sortie [100, 1]
. Le remplissage kSame
remplit d'abord le tableau afin que la forme après la fenêtre de réduction soit la même que l'entrée pour la première étape 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 reduce-window sur le tableau rembourré s'effectue sur trois fenêtres [MAX_VALUE, 10000, 1000]
, [1000, 100, 10]
et [10, 1, MAX_VALUE]
, et génère [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
Voir aussi XlaBuilder::ReplicaId
.
Renvoie l'ID unique (scalaire U32) du réplica.
ReplicaId()
L'ID unique de chaque instance dupliqué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 dupliquées exécutent le même programme, un appel ReplicaId()
dans le programme renvoie une valeur différente pour chaque instance dupliquée.
Remodeler
Consultez également XlaBuilder::Reshape
et l'opération Collapse
.
Remet en forme 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 |
Vecteur int64 |
l'ordre dans lequel les dimensions sont réduites ; |
new_sizes |
Vecteur int64 |
vecteur de tailles de nouvelles dimensions |
Conceptuellement, la fonction reshape aplatit d'abord un tableau en un vecteur unidimensionnel de valeurs de données, puis affine ce vecteur en une nouvelle forme. Les arguments d'entrée sont un tableau arbitraire de type T, un vecteur constant au moment de la compilation d'indices de dimension et un vecteur constant au moment de la compilation de tailles de dimension pour le résultat.
Si elles sont spécifiées, les valeurs du vecteur dimension
doivent être une permutation de toutes les dimensions de T. La valeur par défaut, si elle n'est pas spécifiée, est {0, ..., rank - 1}
. L'ordre des dimensions dans dimensions
va de la dimension à variation la plus lente (la plus importante) à la dimension à variation la plus rapide (la plus mineure) dans le nid de boucles, qui réduit le tableau d'entrée à 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 de dimension de l'opérande. Lorsque vous affinez le tableau récapitulatif dans le tableau multidimensionnel défini par new_sizes
, les dimensions de new_sizes
sont triées de la variation la plus lente (la plus importante) à la variation la plus rapide (la plus mineure).
Par exemple, supposons que v soit 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} } };
Dans un 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} };
Inversion
Voir aussi XlaBuilder::Rev
.
Rev(operand, dimensions)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
tableau de type T |
dimensions |
ArraySlice<int64> |
dimensions à inverser |
Inverse l'ordre des éléments du tableau operand
le long du dimensions
spécifié, générant un tableau de sortie de la même forme. Chaque élément du tableau d'opérandes à un index multidimensionnel est stocké dans le tableau de sortie à un index transformé. L'indice multidimensionnel est transformé en inversant l'indice dans chaque dimension à inverser (c'est-à-dire que si une dimension de taille N est l'une des dimensions à inverser, son indice i est transformé en N-1-i).
L'opération Rev
permet d'inverser le tableau de pondération de convolution le long des deux dimensions de fenêtre lors du calcul du gradient dans les réseaux de neurones.
RngNormal
Voir aussi XlaBuilder::RngNormal
.
Construit une sortie d'une forme donnée avec des nombres aléatoires générés selon la distribution normale \(N(\mu, \sigma)\) . Les paramètres \(\mu\) et \(\sigma\), ainsi que la forme de sortie, doivent avoir un 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 de la valeur générée |
shape |
Shape |
Forme de sortie de type T |
RngUniform
Voir aussi XlaBuilder::RngUniform
.
Construit une sortie d'une forme donnée avec des nombres aléatoires générés selon 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, entier ou à virgule flottante, et les types doivent être cohérents. Les backends de processeur et de GPU ne sont actuellement compatibles qu'avec les types de données F64, F32, F16, BF16, S64, U64, S32 et U32. De plus, les paramètres doivent avoir une valeur scalaire. Si \(b <= a\) le résultat est défini par 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 par défaut du backend) et 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 correspond à l'état initial de la génération de nombres aléatoires actuelle. La forme et les valeurs valides requises, ainsi que la valeur, dépendent de l'algorithme utilisé.
La sortie est garantie comme étant une fonction déterministe de l'état initial, mais elle n'est pas garantie comme étant déterministe entre les backends et les différentes versions du compilateur.
RngBitGenerator(algorithm, key, shape)
Arguments | Type | Sémantique |
---|---|---|
algorithm |
RandomAlgorithm |
Algorithme de 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 PRNG basé sur un compteur ThreeFry. La formeinitial_state
estu64[2]
avec des valeurs arbitraires. Salmon et al. SC 2011. Nombres aléatoires parallèles: c'est aussi simple que 1, 2, 3.rng_philox
: algorithme Philox pour générer des nombres aléatoires en parallèle. La formeinitial_state
estu64[3]
avec des valeurs arbitraires. Salmon et al. SC 2011. Nombres aléatoires parallèles : un jeu d'enfant.
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
.
Voir aussi 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 à répartir. |
scatter_indices |
XlaOp |
Tableau contenant les indices de début des tranches vers lesquelles les données doivent être dispersées. |
updates |
Séquence de N XlaOp |
N tableaux de types T_0, ..., T_N . updates[i] contient les valeurs à utiliser pour la diffusion operands[i] . |
update_computation |
XlaComputation |
Calcul à utiliser pour combiner les valeurs existantes du tableau d'entrée et les mises à jour lors de la dispersion. Ce calcul doit être de type T_0, ..., T_N, T_0, ..., T_N -> Collate(T_0, ..., T_N) . |
index_vector_dim |
int64 |
La dimension dans scatter_indices qui contient les index de départ. |
update_window_dims |
ArraySlice<int64> |
Ensemble des dimensions de la forme updates qui sont des dimensions de 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> |
Mappage des dimensions des indices de dispersion sur l'espace d'index des opérandes. Ce tableau est interprété comme mappant i sur scatter_dims_to_operand_dims[i] . Il doit s'agir d'un accompagnement individuel et total. |
indices_are_sorted |
bool |
Indique si le tri des index est garanti par l'appelant. |
unique_indices |
bool |
Indique si l'unicité des index est garantie par l'appelant. |
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
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
possède une dimension 1
à la fin.
Nous définissons les update_scatter_dims
de type ArraySlice<int64>
comme l'ensemble des dimensions de la forme updates
qui ne sont pas en update_window_dims
, dans l'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 les conditions suivantes :- 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 prise en compte deinserted_window_dims
(c'est-à-direadjusted_window_bounds
[k
], oùadjusted_window_bounds
contient les limites deoperand
avec les limites aux indicesinserted_window_dims
supprimées). - Si
i
est présent dansupdate_scatter_dims
(c'est-à-dire égal àupdate_scatter_dims
[k
] pour certainsk
), la limite de la dimensioni
dansupdates
doit être égale à celle correspondante descatter_indices
, en ignorantindex_vector_dim
(par exemple,scatter_indices.shape.dims
[k
], sik
est inférieur àindex_vector_dim
etscatter_indices.shape.dims
[k+1
] dans le cas contraire).
- Si
update_window_dims
doit être trié par ordre croissant, ne pas contenir de numéros de dimension en double 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 dimension identiques et se situer 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 index donné U
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'indexS
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 du mappagescatter_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
sinon.
- Créez un index
W
in
dans chaque tableauoperands
en dispersant les indices àupdate_window_dims
dansU
seloninserted_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
}, alorswindow_dims_to_operand_dims
est {0
→1
,1
→3
,2
→4
,3
→5
}.W
in
[_
] =0
sinon.
I
estW
in
+S
in
, où + correspond à l'addition par élément.
En résumé, l'opération de dispersion peut être définie comme suit.
- Initialisez
output
avecoperands
, c'est-à-dire pour tous les indicesJ
, pour tous les indicesO
dans le tableauoperands
[J
] :
output
[J
][O
] =operands
[J
][O
] - Pour chaque index
U
dans le tableauupdates
[J
] et l'indexO
correspondant dans le tableauoperand
[J
], siO
est un index 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 est non déterministe. Par conséquent, lorsque plusieurs indices dans updates
font référence au même indice dans operands
, la valeur correspondante dans output
n'est pas déterministe.
Notez que le premier paramètre transmis à update_computation
sera toujours la valeur actuelle du tableau output
et que le deuxième paramètre sera toujours la valeur du tableau updates
. Cela est particulièrement important lorsque update_computation
n'est pas commutatif.
Si indices_are_sorted
est défini sur "true", XLA peut supposer que les scatter_indices
sont triées (par ordre croissant, après avoir dispersé ses valeurs en fonction de scatter_dims_to_operand_dims
) par l'utilisateur. Si ce n'est pas le cas, la sémantique
est définie pour l'implémentation.
Si unique_indices
est défini sur "true", XLA peut supposer que tous les éléments dispersés sont uniques. XLA peut donc utiliser des opérations non atomiques. Si unique_indices
est défini sur "true" et que les indices vers lesquels la diffusion est effectuée ne sont pas uniques, la sémantique est définie par 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 qui sont extraits par l'opération de collecte correspondante.
Pour obtenir une description informelle détaillée et des exemples, consultez la section "Description informationnelle" sous Gather
.
Sélectionner
Voir aussi XlaBuilder::Select
.
Construit un tableau de sortie à partir des éléments de deux tableaux d'entrée, en fonction des valeurs d'un tableau de prédicats.
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 restreinte de diffusion, pred
peut être un scalaire de type PRED
. Dans ce cas, le tableau de sortie est entièrement extrait de on_true
si la valeur de pred
est true
, et de on_false
si la valeur de pred
est false
.
Exemple avec 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 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};
Les sélections entre des tupels sont acceptées. À cette fin, les tupels sont considérés comme des types scalaires. Si on_true
et on_false
sont des tupels (qui doivent avoir la même forme), pred
doit être un scalaire de type PRED
.
SelectAndScatter
Voir aussi 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 de chaque fenêtre, puis disperse le tableau source
sur les indices des éléments sélectionnés pour construire un tableau de sortie de la même forme que le tableau d'opérande. La fonction select
binaire 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é que le vecteur d'index du premier paramètre est lexicographiquement inférieur au vecteur d'index 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 respecter la transitivité (c'est-à-dire que 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épende pas de l'ordre des éléments parcourus pour une fenêtre donnée.
La fonction scatter
est appliquée à chaque indice sélectionné dans le tableau de sortie. Elle prend deux paramètres scalaires :
- Valeur actuelle à l'index sélectionné dans le tableau de sortie
- Valeur de dispersion de
source
qui s'applique à l'indice sélectionné
Il combine les deux paramètres et renvoie une valeur scalaire utilisée pour mettre à jour la valeur à l'index sélectionné dans le tableau de sortie. Initialement, 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
sur le tableau operand
. SelectAndScatter
peut être utilisé pour rétropropager les valeurs de gradient pour 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 de dimension 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 à son élément de destination |
La figure ci-dessous montre des exemples d'utilisation de SelectAndScatter
, avec la fonction select
qui calcule la valeur maximale parmi ses paramètres. Notez que lorsque les fenêtres se chevauchent, comme dans la figure 2 ci-dessous, un indice du tableau operand
peut être sélectionné plusieurs fois par différentes fenêtres. Dans la figure, l'élément de la valeur 9 est sélectionné par les deux fenêtres supérieures (bleu et rouge) et la fonction d'ajout binaire scatter
produit l'élément de sortie de la 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
Voir aussi XlaBuilder::Send
.
Send(operand, channel_handle)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
données à envoyer (tableau de type T) |
channel_handle |
ChannelHandle |
identifiant unique pour chaque paire envoi/réception |
Envoie les données d'opérande données à une instruction Recv
dans un autre calcul qui partage le même gestionnaire de canaux. 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. Voir é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 utilisé par une instruction SendDone
suivante pour attendre la fin du transfert de données. Le contexte est un tuple de {operand (shape), request identifier (U32)} et 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.
Planifier des instructions de diffusion sur une chaîne
L'ordre d'exécution des quatre instructions pour chaque canal (Recv
, RecvDone
, Send
et SendDone
) est le suivant.
Recv
se produit avantSend
Send
a lieu avantRecvDone
Recv
se produit avantRecvDone
Send
se produit avantSendDone
Lorsque les compilateurs backend génèrent une programmation 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 planifications ci-dessous entraînent des interblocages.
Tranche
Voir aussi XlaBuilder::Slice
.
Le découpage extrait un sous-tableau du tableau d'entrée. Le sous-tableau est du même rang que l'entrée et contient les valeurs dans un rectangle englobant dans le tableau d'entrée, où les dimensions et les indices du rectangle englobant sont donnés comme arguments à 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 indices de début 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 indices de fin (exclusifs) de la tranche pour chaque dimension. Chaque valeur doit être supérieure ou égale à la valeur start_indices respective de la dimension, et inférieure ou égale à la taille de la dimension. |
strides |
ArraySlice<int64> |
Liste de N entiers qui déterminent la longueur de la ligne d'entrée de la tranche. Le segment sélectionne tous les éléments strides[d] de la dimension d . |
Exemple à une dimension :
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} }
Trier
Voir aussi XlaBuilder::Sort
.
Sort(operands, comparator, dimension, is_stable)
Arguments | Type | Sémantique |
---|---|---|
operands |
ArraySlice<XlaOp> |
Les opérandes à trier. |
comparator |
XlaComputation |
Calcul du comparateur à utiliser. |
dimension |
int64 |
Dimension selon laquelle le tri doit être effectué. |
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érieure. Formellement, une fois le tableau trié, il contient pour toutes les positions d'index
i, j
aveci < j
qui estcomparator(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é selon la dimension fournie. Par exemple, pour un tenseur de rang 2 (une matrice), une valeur de dimension
0
triera indépendamment chaque colonne, et une valeur de dimension1
triera indépendamment chaque ligne. 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 du classement de niveau 1.
Si des opérandes n > 1
sont fournis:
Tous les opérandes
n
doivent être des tenseurs de même dimension. Les types d'éléments des tenseurs 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
correspond à la valeur à la positionj
de l'opérandek-th
. En règle générale, le comparateur compare les paramètres2 * k
et2 * k + 1
entre eux, et utilise éventuellement d'autres paires de paramètres pour départager les paramètres.Le résultat est un tuple constitué des opérandes dans l'ordre de tri (le long de la dimension fournie, comme ci-dessus). L'opérande
i-th
du tuple correspond à l'opérandei-th
de la fonction 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 la tuple ([1, 3], [50, 42], [1.1, -3.0])
.
Si is_stable
est défini sur "true", le tri est garanti 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 préservé. 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".
Transposer
Consultez également l'opération tf.reshape
.
Transpose(operand)
Arguments | Type | Sémantique |
---|---|---|
operand |
XlaOp |
Operande à transposer. |
permutation |
ArraySlice<int64> |
Permuter les dimensions |
Permute les dimensions de l'opérande avec la permutation donnée, par exemple ∀ i . 0 ≤ i < rank ⇒ input_dimensions[permutation[i]] = output_dimensions[i]
.
Cela revient à utiliser Reshape(operand, permutation, Permute(permutation, operand.shape.dimensions)).
TriangularSolve
Voir aussi XlaBuilder::TriangularSolve
.
Résoudre des systèmes d'équations linéaires avec des matrices de coefficients triangulaires inférieurs ou supérieurs par substitution avant ou arrière. En diffusant le long des dimensions de début, cette routine résout l'un des systèmes matriciels op(a) * x =
b
, ou x * op(a) = b
, pour la variable x
, selon a
et b
, où op(a)
est op(a) = a
, 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 de rang > 2 d'un type complexe ou à virgule flottante de forme [..., M, M] . |
b |
XlaOp |
Un tableau de rang > 2 du même type avec la forme [..., M, K] si left_side est défini sur "true", [..., K, M] sinon. |
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 |
d'utiliser le triangle supérieur ou inférieur de a . |
unit_diagonal |
bool |
Si true , les éléments de la diagonale de a sont supposés être 1 et ne sont pas accessibles. |
transpose_a |
Transpose |
d'utiliser a tel quel, de le transposer ou de prendre sa transposée conjuguée. |
Les données d'entrée ne sont lues que dans le 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 quoi.
Si le rang de a
et de b
est supérieur à 2, ils sont traités comme des lots de matrices, où toutes les dimensions à l'exception des 2 mineurs sont des dimensions de lot. a
et b
doivent avoir des dimensions de lot égales.
Tuple
Voir aussi XlaBuilder::Tuple
.
Un tuple contenant un nombre variable de poignées de données, chacune ayant sa propre forme.
Cela équivaut à std::tuple
en C++. Conceptuellement :
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écomposés (accessibles) via l'opération GetTupleElement
.
Bien que
Voir aussi 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 du paramètre de condition et body . |
Exécute de manière séquentielle le body
jusqu'à l'échec de l'condition
. Elle est semblable à une boucle while typique dans de nombreuses autres langues, à l'exception des différences et des restrictions listées ci-dessous.
- Un nœud
While
renvoie une valeur de typeT
, qui correspond au 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 la même pour 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 avec le nouveau résultat de body
à chaque itération ultérieure.
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 des réseaux de neurones. Le pseudo-code simplifié est illustré ci-dessous avec un graphique qui représente le calcul. Le code se trouve 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};
}