Sémantique de l'opération

La section suivante décrit la sémantique des opérations définies dans l'interface XlaBuilder. En règle générale, ces opérations correspondent à celles définies dans l'interface RPC de xla_data.proto.

Remarque sur la nomenclature : le type de données généralisé XLA traite d'un tableau N-dimensionnel contenant des éléments d'un type uniforme (tel que float 32 bits). Tout au long de la documentation, le terme array est utilisé pour désigner un tableau de dimensions arbitraires. 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.

Pour en savoir plus sur la structure d'une opération, consultez Formes et mise en page et Mise en page en mosaïque.

Abdos

Voir aussi XlaBuilder::Abs.

Valeur absolue par élément x -> |x|.

Abs(operand)

Pour en savoir plus sur StableHLO, consultez StableHLO – abs.

Ajouter

Voir aussi XlaBuilder::Add.

Effectue une addition élément par élément de lhs et rhs.

Add(lhs, rhs)

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 résulte 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.

Il existe une variante alternative avec une prise en charge de la diffusion de différentes dimensions pour l'opération Add :

Add(lhs,rhs, broadcast_dimensions)

Cette variante de l'opération doit être utilisée pour les opérations arithmétiques entre des tableaux de rangs différents (par exemple, l'addition d'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 de diffusion.

Pour en savoir plus sur StableHLO, consultez StableHLO – Ajouter.

AddDependency

Voir aussi HloInstruction::AddDependency.

AddDependency peut apparaître dans les dumps HLO, mais les utilisateurs finaux ne sont pas censés les construire manuellement.

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 à effets secondaires pour appliquer un ordre. AfterAll peut être utilisé comme jointure de jetons pour ordonner une opération après un ensemble d'opérations.

AfterAll(tokens)

Pour en savoir plus sur StableHLO, consultez StableHLO – after_all.

AllGather

Voir aussi XlaBuilder::AllGather.

Effectue la concaténation entre les répliques.

AllGather(operand, all_gather_dimension, shard_count, replica_groups, channel_id, layout, use_global_device_ids)

• replica_groups est une liste de groupes de répliques entre lesquels la concaténation est effectuée (l'ID de réplique de la réplique actuelle peut être récupéré à l'aide de ReplicaId). L'ordre des répliques dans chaque groupe détermine l'ordre dans lequel leurs entrées se trouvent dans le résultat. replica_groups doit être vide (dans ce cas, toutes les répliques appartiennent à un seul groupe, ordonné de 0 à N - 1) ou contenir le même nombre d'éléments que le nombre de répliques. Par exemple, replica_groups = {0, 2}, {1, 3} effectue une concaténation entre les répliques 0 et 2, et 1 et 3.
• shard_count correspond à la taille de chaque groupe de réplicas. Nous en avons besoin lorsque replica_groups est vide.
• channel_id est utilisé pour la communication intermodules : seules les opérations all-gather avec le même channel_id peuvent communiquer entre elles.
• use_global_device_ids Renvoie "true" si les ID de la configuration ReplicaGroup représentent un ID global (replica_id * partition_count + partition_id) au lieu d'un ID de réplica. Cela permet un regroupement plus flexible des appareils si cette réduction all-reduce est à la fois inter-partition et inter-réplique.

La forme de sortie est la forme d'entrée avec le all_gather_dimension rendu shard_count fois plus grand. Par exemple, s'il existe deux répliques et que l'opérande a respectivement la valeur [1.0, 2.5] et [3.0, 5.25] sur les deux répliques, 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 répliques.

L'API de AllGather est décomposée en interne en deux instructions HLO (AllGatherStart et AllGatherDone).

Voir aussi HloInstruction::CreateAllGatherStart.

AllGatherStart et AllGatherDone servent de primitives dans HLO. Ces opérations peuvent apparaître dans les dumps HLO, mais elles ne sont pas destinées à être construites manuellement par les utilisateurs finaux.

Pour en savoir plus sur StableHLO, consultez StableHLO – all_gather.

AllReduce

Voir aussi XlaBuilder::AllReduce.

Effectue un calcul personnalisé sur les répliques.

AllReduce(operand, computation, replica_groups, channel_id, shape_with_layout, use_global_device_ids)

• Lorsque operand est un tuple de tableaux, l'opération all-reduce est effectuée sur chaque élément du tuple.
• replica_groups est une liste de groupes de répliques entre lesquels la réduction est effectuée (l'ID de réplique de la réplique actuelle peut être récupéré à l'aide de ReplicaId). replica_groups doit être vide (dans ce cas, toutes les répliques appartiennent à un seul groupe) ou contenir le même nombre d'éléments que le nombre de répliques. Par exemple, replica_groups = {0, 2}, {1, 3} effectue une réduction entre les répliques 0 et 2, et 1 et 3.
• channel_id est utilisé pour la communication intermodules : seules les opérations all-reduce avec le même channel_id peuvent communiquer entre elles.
• shape_with_layout : force la mise en page d'AllReduce à la mise en page donnée. Cela permet de garantir la même mise en page pour un groupe d'opérations AllReduce compilées séparément.
• use_global_device_ids Renvoie "true" si les ID de la configuration ReplicaGroup représentent un ID global (replica_id * partition_count + partition_id) au lieu d'un ID de réplica. Cela permet un regroupement plus flexible des appareils si cette réduction all-reduce est à la fois inter-partition et inter-réplique.

La forme de sortie est identique à la forme d'entrée. Par exemple, s'il existe deux répliques et que l'opérande a respectivement la valeur [1.0, 2.5] et [3.0, 5.25] sur les deux répliques, la valeur de sortie de cette opération et du calcul de la somme sera [4.0, 7.75] sur les deux répliques. Si l'entrée est un tuple, la sortie l'est également.

Pour calculer le résultat de AllReduce, il faut une entrée de chaque réplica. Par conséquent, si un réplica exécute un nœud AllReduce plus de fois qu'un autre, le premier réplica attendra indéfiniment. Comme les réplicas exécutent tous le même programme, il n'y a pas beaucoup de façons pour que cela se produise, mais c'est possible lorsque la condition d'une boucle while dépend des données de infeed et que les données de infeed entraînent l'itération de la boucle while plus de fois sur un réplica que sur un autre.

L'API de AllReduce est décomposée en interne en deux instructions HLO (AllReduceStart et AllReduceDone).

Voir aussi HloInstruction::CreateAllReduceStart.

AllReduceStart et AllReduceDone servent de primitives dans HLO. Ces opérations peuvent apparaître dans les dumps HLO, mais elles ne sont pas destinées à être construites manuellement par les utilisateurs finaux.

CrossReplicaSum

Voir aussi XlaBuilder::CrossReplicaSum.

Effectue AllReduce avec un calcul de la somme.

CrossReplicaSum(operand, replica_groups)

Renvoie la somme de la valeur de l'opérande dans chaque sous-groupe de réplicas. Toutes les répliques fournissent une entrée à la somme et toutes les répliques reçoivent la somme résultante pour chaque sous-groupe.

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 :

1. La phase de dispersion. Sur chaque cœur, l'opérande est divisé en split_count blocs le long de split_dimensions, et les blocs sont répartis sur tous les cœurs (par exemple, le i-ième bloc est envoyé au i-ième cœur).
2. La phase de collecte. Chaque cœur concatène les blocs reçus le long de concat_dimension.

Les cœurs participants peuvent être configurés de la manière suivante :

• replica_groups : chaque ReplicaGroup contient une liste d'ID de répliques participant au calcul (l'ID de réplique de la réplique actuelle peut être récupéré à l'aide de ReplicaId). AllToAll sera appliqué aux sous-groupes dans l'ordre spécifié. Par exemple, replica_groups = { {1,2,3}, {4,5,0} } signifie qu'un AllToAll sera appliqué dans les répliques {1, 2, 3}, et dans la phase de collecte, et que les blocs reçus seront concaténés dans le même ordre : 1, 2, 3. Ensuite, un autre AllToAll sera appliqué aux répliques 4, 5 et 0, et l'ordre de concaténation sera également 4, 5 et 0. Si replica_groups est vide, toutes les répliques appartiennent à un seul groupe, dans l'ordre de concaténation de leur apparition.

Prérequis :

• La taille de la dimension de l'opérande sur split_dimension est divisible par split_count.
• La forme de l'opérande n'est pas un tuple.

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

Pour en savoir plus sur les formes et les mises en page, consultez xla::shapes.

Pour en savoir plus sur StableHLO, consultez StableHLO – all_to_all.

AllToAll – Exemple 1.

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 l'exemple ci-dessus, quatre cœurs participent à Alltoall. Sur chaque cœur, l'opérande est divisé en quatre parties le long de la dimension 1. Chaque partie a donc une forme f32[4,4]. Les quatre parties sont dispersées dans tous les cœurs. Chaque cœur concatène ensuite les parties reçues le long de la dimension 0, dans l'ordre des cœurs 0 à 4. La forme de la sortie sur chaque cœur est donc f32[16,4].

AllToAll – Exemple 2 – StableHLO

Dans l'exemple ci-dessus, deux répliques participent à AllToAll. Sur chaque réplique, l'opérande a la forme f32[2,4]. L'opérande est divisé en deux parties le long de la dimension 1. Chaque partie a donc une forme f32[2,2]. Les deux parties sont ensuite échangées entre les réplicas en fonction de leur position dans le groupe de réplicas. Chaque réplique collecte la partie correspondante des deux opérandes et les concatène le long de la dimension 0. Par conséquent, la sortie de chaque réplica a la forme f32[4,2].

RaggedAllToAll

Voir aussi XlaBuilder::RaggedAllToAll.

RaggedAllToAll effectue une opération collective all-to-all, où l'entrée et la sortie sont des Tensors irréguliers.

RaggedAllToAll(input, input_offsets, send_sizes, output, output_offsets, recv_sizes, replica_groups, channel_id)

Les Tensors irréguliers sont définis par un ensemble de trois Tensors :

• data : le tenseur data est "irrégulier" le long de sa dimension la plus externe, le long de laquelle chaque élément indexé a une taille variable.
• offsets : le tenseur offsets indexe la dimension la plus externe du tenseur data et représente le décalage de début de chaque élément irrégulier du tenseur data.
• sizes : le Tensor sizes représente la taille de chaque élément irrégulier du Tensor data, où la taille est spécifiée en unités de sous-éléments. Un sous-élément est défini comme le suffixe de la forme du Tensor "data" obtenu en supprimant la dimension "irrégulière" la plus externe.
• Les tenseurs offsets et sizes doivent avoir la même taille.

Exemple de Tensor irrégulier :

data: [8,3] =
{ {a,b,c},{d,e,f},{g,h,i},{j,k,l},{m,n,o},{p,q,r},{s,t,u},{v,w,x} }

offsets: [3] = {0, 1, 4}

sizes: [3] = {1, 3, 4}

// Index 'data' at 'offsets'[0], 'sizes'[0]' // {a,b,c}

// Index 'data' at 'offsets'[1], 'sizes'[1]' // {d,e,f},{g,h,i},{j,k,l}

// Index 'data' at 'offsets'[2], 'sizes'[2]' // {m,n,o},{p,q,r},{s,t,u},{v,w,x}

output_offsets doit être fragmenté de sorte que chaque réplica ait des décalages dans la perspective de sortie du réplica cible.

Pour le décalage de sortie i, le réplica actuel enverra une mise à jour input[input_offsets[i]:input_offsets[i]+input_sizes[i]] au réplica i qui sera écrit dans output_i[output_offsets[i]:output_offsets[i]+send_sizes[i]] dans le réplica i output.

Par exemple, si nous avons deux réplicas :

replica 0:
input: [1, 2, 2]
output:[0, 0, 0, 0]
input_offsets: [0, 1]
send_sizes: [1, 2]
output_offsets: [0, 0]
recv_sizes: [1, 1]

replica 1:
input: [3, 4, 0]
output: [0, 0, 0, 0]
input_offsets: [0, 1]
send_sizes: [1, 1]
output_offsets: [1, 2]
recv_sizes: [2, 1]

// replica 0's result will be: [1, 3, 0, 0]
// replica 1's result will be: [2, 2, 4, 0]

Le HLO ragged all-to-all comporte les arguments suivants :

• input : Tensor de données d'entrée irrégulières.
• output : Tensor de données de sortie irrégulier.
• input_offsets : Tensor des décalages d'entrée irréguliers.
• send_sizes : Tensor des tailles d'envoi irrégulières.
• output_offsets : tableau des décalages irréguliers dans la sortie de réplique cible.
• recv_sizes : Tensor des tailles de réception irrégulières.

Les Tensors *_offsets et *_sizes doivent tous avoir la même forme.

Deux formes sont acceptées pour les Tensors *_offsets et *_sizes :

• [num_devices], où ragged-all-to-all peut envoyer au maximum une mise à jour à chaque appareil distant du groupe de répliques. Exemple :

for (remote_device_id : replica_group) {
     SEND input[input_offsets[remote_device_id]],
     output[output_offsets[remote_device_id]],
     send_sizes[remote_device_id] }

• [num_devices, num_updates] où ragged-all-to-all peut envoyer jusqu'à num_updates mises à jour au même appareil à distance (chacune à des décalages différents), pour chaque appareil à distance du groupe de répliques.

Exemple :

for (remote_device_id : replica_group) {
    for (update_idx : num_updates) {
        SEND input[input_offsets[remote_device_id][update_idx]],
        output[output_offsets[remote_device_id][update_idx]]],
        send_sizes[remote_device_id][update_idx] } }

et

Voir aussi XlaBuilder::And.

Effectue un AND élément par élément de deux Tensors lhs et rhs.

And(lhs, rhs)

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 résulte 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.

Il existe une variante alternative avec une prise en charge de la diffusion multidimensionnelle pour And :

And(lhs,rhs, broadcast_dimensions)

Cette variante de l'opération doit être utilisée pour les opérations arithmétiques entre des tableaux de rangs différents (par exemple, l'addition d'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 de diffusion.

Pour en savoir plus sur StableHLO, consultez StableHLO.

Asynchrone

Voir aussi HloInstruction::CreateAsyncStart, HloInstruction::CreateAsyncUpdate, HloInstruction::CreateAsyncDone.

AsyncDone, AsyncStart et AsyncUpdate sont des instructions HLO internes utilisées pour les opérations asynchrones et servent de primitives dans HLO. Ces opérations peuvent apparaître dans les dumps HLO, mais elles ne sont pas destinées à être construites manuellement par les utilisateurs finaux.

Atan2

Voir aussi XlaBuilder::Atan2.

Effectue une opération atan2 au niveau des éléments sur lhs et rhs.

Atan2(lhs, rhs)

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 résulte 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.

Il existe une variante alternative avec une prise en charge de la diffusion multidimensionnelle pour Atan2 :

Atan2(lhs,rhs, broadcast_dimensions)

Cette variante de l'opération doit être utilisée pour les opérations arithmétiques entre des tableaux de rangs différents (par exemple, l'addition d'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 de diffusion.

Pour en savoir plus sur StableHLO, consultez StableHLO – atan2.

BatchNormGrad

Consultez également XlaBuilder::BatchNormGrad et l'article d'origine sur la normalisation par lot pour obtenir une description détaillée de l'algorithme.

Calcule les gradients de la normalisation par lot.

BatchNormGrad(operand, scale, batch_mean, batch_var, grad_output, epsilon, feature_index)

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 dans toutes les autres dimensions. feature_index doit être un index valide pour la dimension de caractéristique dans operand.

Les trois gradients sont définis par les formules suivantes (en supposant un tableau à quatre dimensions comme operand avec l'index de dimension de caractéristiques l, la taille du 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 batch_mean et batch_var représentent les valeurs des moments pour les dimensions spatiales et de lot.

Le type de sortie est un tuple de trois handles :

Pour en savoir plus sur StableHLO, consultez StableHLO – batch_norm_grad.

BatchNormInference

Consultez également XlaBuilder::BatchNormInference et l'article d'origine sur la normalisation par lot pour obtenir une description détaillée de l'algorithme.

Normalise un tableau sur les dimensions de lot et spatiales.

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

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 la moyenne et la variance pour toutes les autres dimensions, puis les utilise pour normaliser chaque élément de 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 entrées mean et variance comme valeurs estimées. L'objectif de cette opération est de réduire la latence dans l'inférence, d'où le nom BatchNormInference.

Le résultat est un tableau normalisé à n dimensions ayant la même forme que l'entrée operand.

Pour en savoir plus sur StableHLO, consultez StableHLO – batch_norm_inference.

BatchNormTraining

Pour obtenir une description détaillée de l'algorithme, consultez également XlaBuilder::BatchNormTraining et the original batch normalization paper.

Normalise un tableau sur les dimensions de lot et spatiales.

BatchNormTraining(operand, scale, offset, epsilon, feature_index)

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 la moyenne et la variance pour toutes les autres dimensions, puis les utilise pour normaliser chaque élément de operand. feature_index doit être un index valide pour la dimension de caractéristique dans operand.

L'algorithme se déroule comme suit pour chaque lot de operand \(x\) contenant m éléments avec w et h comme taille des dimensions spatiales (en supposant que operand est un tableau à quatre dimensions) :

• Calcule la moyenne du lot \(\mu_l\) pour chaque caractéristiqueldans la dimension de la caractéristique : \(\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}\)

• Calcule la variance du lot \(\sigma^2_l\) : $\sigma^2l=\frac{1}{mwh}\sum{i=1}^m\sum{j=1}^w\sum{k=1}^h (x_{ijkl} - \mu_l)^2$

• Normalisation, mise à l'échelle et décalage : \(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 :

batch_mean et batch_var sont des moments calculés sur les dimensions de lot et spatiales à l'aide des formules ci-dessus.

Pour en savoir plus sur StableHLO, consultez StableHLO – batch_norm_training.

Bitcast

Voir aussi HloInstruction::CreateBitcast.

Bitcast peut apparaître dans les dumps HLO, mais n'est pas destiné à être construit manuellement par les utilisateurs finaux.

BitcastConvertType

Voir aussi XlaBuilder::BitcastConvertType.

Semblable à un tf.bitcast dans TensorFlow, effectue une opération bitcast élément par élément d'une forme de données vers une forme cible. La taille des entrées et des sorties doit correspondre : par exemple, s32 éléments deviennent f32 éléments via la routine bitcast, et un élément s32 devient quatre éléments s8. Bitcast est implémenté en tant que cast de bas niveau. Par conséquent, les machines avec des représentations à virgule flottante différentes donneront des résultats différents.

BitcastConvertType(operand, new_element_type)

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 types d'éléments source et de destination ne doivent pas être des tuples.

Pour en savoir plus sur StableHLO, consultez StableHLO – bitcast_convert.

Conversion bitcast en type primitif de largeur différente

L'instruction HLO BitcastConvert est compatible avec 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. Comme l'opération entière est conceptuellement un bitcast et ne modifie pas les octets sous-jacents, la forme de l'élément de sortie doit changer. Pour B = sizeof(T), B' = sizeof(T'), deux cas sont possibles.

Tout d'abord, lorsque B > B', la forme de sortie obtient une nouvelle dimension 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. 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 élément par élément.

Annoncer

Voir aussi XlaBuilder::Broadcast.

Ajoute des dimensions à un tableau en dupliquant les données qu'il contient.

Broadcast(operand, broadcast_sizes)

Les nouvelles dimensions sont insérées à gauche. Par exemple, si broadcast_sizes a des valeurs {a0, ..., aN} et que la forme de l'opérande a des dimensions {b0, ..., bM}, la forme de la sortie a des dimensions {a0, ..., aN, b0, ..., bM}.

Les nouvelles dimensions indexent les copies de l'opérande, c'est-à-dire

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

Par exemple, si operand est un scalaire f32 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 du résultat seront 2.0f.

Pour en savoir plus sur StableHLO, consultez StableHLO – diffusion.

BroadcastInDim

Voir aussi XlaBuilder::BroadcastInDim.

Développe la taille et le nombre de dimensions d'un tableau en dupliquant les données du tableau.

BroadcastInDim(operand, out_dim_size, broadcast_dimensions)

Semblable à Broadcast, mais permet d'ajouter des dimensions n'importe où et d'étendre les dimensions existantes de taille 1.

operand est diffusé dans la forme décrite par out_dim_size. broadcast_dimensions mappe les dimensions de operand à celles de la forme cible, c'est-à-dire que la i-ième dimension de l'opérande est mappée à la i-ième dimension de la forme de sortie. Les dimensions de operand doivent être de taille 1 ou 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 est ensuite effectuée le long de ces dimensions dégénérées pour atteindre la forme de sortie. La sémantique est décrite en détail sur la page de diffusion.

Appeler

Voir aussi XlaBuilder::Call.

Appelle un calcul avec les arguments fournis.

Call(computation, operands...)

L'arité et les types de operands doivent correspondre aux paramètres de computation. Il est autorisé de ne pas avoir de operands.

CompositeCall

Voir aussi XlaBuilder::CompositeCall.

Encapsule une opération composée d'autres opérations StableHLO, en prenant des entrées et des composite_attributes et en produisant des résultats. La sémantique de l'opération est implémentée par l'attribut de décomposition. L'opération composite peut être remplacée par sa décomposition sans modifier la sémantique du programme. Dans les cas où l'intégration de la décomposition ne fournit pas la même sémantique d'opération, préférez utiliser custom_call.

Le champ de version (qui est défini sur 0 par défaut) est utilisé pour indiquer quand la sémantique d'un composite change.

Cette opération est implémentée en tant que kCall avec l'attribut is_composite=true. Le champ decomposition est spécifié par l'attribut computation. Les attributs frontend stockent les attributs restants précédés du préfixe composite..

Exemple d'opération CompositeCall :

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

CompositeCall(computation, operands..., name, attributes, version)

Le decomposition d'une opération n'est pas un champ appelé, mais apparaît plutôt comme un attribut to_apply qui pointe vers la fonction contenant l'implémentation de niveau inférieur, c'est-à-dire to_apply=%funcname.

Pour en savoir plus sur la composition et la décomposition, consultez la spécification StableHLO.

Cbrt

Voir aussi XlaBuilder::Cbrt.

Opération de racine cubique au niveau des éléments x -> cbrt(x).

Cbrt(operand)

Cbrt accepte également l'argument facultatif result_accuracy :

Cbrt(operand, result_accuracy)

Pour en savoir plus sur result_accuracy, consultez Précision des résultats.

Pour en savoir plus sur StableHLO, consultez StableHLO – cbrt.

Ceil

Voir aussi XlaBuilder::Ceil.

Plafond par élément x -> ⌈x⌉.

Ceil(operand)

Pour en savoir plus sur StableHLO, consultez StableHLO – ceil.

Cholesky

Voir aussi XlaBuilder::Cholesky.

Calcule la décomposition de Cholesky d'un lot de matrices définies positives symétriques (hermitiennes).

Cholesky(a, lower)

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 qu'à partir du triangle inférieur/supérieur de a, en fonction de la valeur de lower. Les valeurs de l'autre triangle sont ignorées. Les données de sortie sont renvoyées dans le même triangle. Les valeurs de l'autre triangle sont définies par l'implémentation et peuvent être n'importe quoi.

Si a comporte plus de deux dimensions, 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 défini comme positif, symétrique (hermitien), le résultat est défini par l'implémentation.

Pour en savoir plus sur StableHLO, consultez StableHLO – cholesky.

Limiter

Voir aussi XlaBuilder::Clamp.

Limite un opérande à une plage comprise entre une valeur minimale et une valeur maximale.

Clamp(min, operand, max)

Étant donné un opérande et des valeurs minimale et maximale, renvoie l'opérande s'il se trouve dans la plage comprise entre le minimum et le maximum. 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. Vous pouvez également utiliser min et/ou max comme scalaire de type T, en tant que forme restreinte de diffusion.

Exemple avec min et max scalaires :

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

Pour en savoir plus sur StableHLO, consultez StableHLO – clamp.

Réduire

Voir aussi XlaBuilder::Collapse. et l'opération tf.reshape.

Réduit les dimensions d'un tableau à une seule dimension.

Collapse(operand, dimensions)

Collapse 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 de constantes connues à la compilation d'indices de dimension. Les indices de dimension doivent être un sous-ensemble consécutif (du plus petit au plus grand) des dimensions de T. Ainsi, {0, 1, 2}, {0, 1} ou {1, 2} sont des ensembles de dimensions valides, mais {1, 0} ou {0, 2} ne le sont pas. Elles sont remplacées par une seule et même dimension, à la même position dans la séquence de dimensions que celles qu'elles remplacent, avec une taille de dimension égale au produit des tailles de dimensions d'origine. Le plus petit nombre de dimension dans dimensions correspond à la dimension à la variation la plus lente (la plus importante) dans la structure de boucle qui réduit ces dimensions, et le plus grand nombre de dimension correspond à la dimension à la variation la plus rapide (la moins importante). Consultez l'opérateur tf.reshape si vous avez besoin d'un ordre de réduction plus général.

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

Clz

Voir aussi XlaBuilder::Clz.

Compte les zéros en début de nombre pour chaque élément.

Clz(operand)

CollectiveBroadcast

Voir aussi XlaBuilder::CollectiveBroadcast.

Diffuse les données sur les répliques. Les données sont envoyées du premier ID de réplica de chaque groupe aux autres ID du même groupe. Si un ID de réplique ne figure dans aucun groupe de répliques, la sortie de cette réplique est un Tensor composé de 0 dans shape.

CollectiveBroadcast(operand, replica_groups, channel_id)

Pour en savoir plus sur StableHLO, consultez StableHLO – collective_broadcast.

CollectivePermute

Voir aussi XlaBuilder::CollectivePermute.

CollectivePermute est une opération collective qui envoie et reçoit des données entre les répliques.

CollectivePermute(operand, source_target_pairs, channel_id, inplace)

Notez que les restrictions suivantes s'appliquent à source_target_pairs :

• Deux paires ne doivent pas avoir le même ID de réplica cible ni le même ID de réplica source.
• Si un ID de réplique n'est la cible d'aucune paire, la sortie sur cette réplique est un Tensor composé de 0(s) avec la même forme que l'entrée.

L'API de l'opération CollectivePermute est décomposée en interne en deux instructions HLO (CollectivePermuteStart et CollectivePermuteDone).

Voir aussi HloInstruction::CreateCollectivePermuteStart.

CollectivePermuteStart et CollectivePermuteDone servent de primitives dans HLO. Ces opérations peuvent apparaître dans les dumps HLO, mais elles ne sont pas destinées à être construites manuellement par les utilisateurs finaux.

Pour en savoir plus sur StableHLO, consultez StableHLO – collective_permute.

Comparer

Voir aussi XlaBuilder::Compare.

Effectue une comparaison élément par élément de lhs et rhs des éléments suivants :

Eq

Voir aussi XlaBuilder::Eq.

Effectue une comparaison égal à élément par élément de lhs et rhs.

\(lhs = rhs\)

Eq(lhs, rhs)

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 résulte 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.

Il existe une variante alternative avec une prise en charge de la diffusion multidimensionnelle pour l'équation :

Eq(lhs,rhs, broadcast_dimensions)

Cette variante de l'opération doit être utilisée pour les opérations arithmétiques entre des tableaux de rangs différents (par exemple, l'addition d'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 de diffusion.

Une commande totale sur les nombres à virgule flottante existe pour Eq, en appliquant :

\[-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.\]

EqTotalOrder(lhs,rhs, broadcast_dimensions)

Pour en savoir plus sur StableHLO, consultez StableHLO – Comparer.

Ne

Voir aussi XlaBuilder::Ne.

Effectue une comparaison not equal-to (différent de) au niveau des éléments de lhs et rhs.

\(lhs != rhs\)

Ne(lhs, rhs)

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 résulte 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.

Il existe une variante alternative avec une prise en charge de la diffusion multidimensionnelle pour Ne :

Ne(lhs,rhs, broadcast_dimensions)

Cette variante de l'opération doit être utilisée pour les opérations arithmétiques entre des tableaux de rangs différents (par exemple, l'addition d'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 de diffusion.

Une commande totale sur les nombres à virgule flottante existe pour Ne, en appliquant :

\[-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.\]

NeTotalOrder(lhs,rhs, broadcast_dimensions)

Pour en savoir plus sur StableHLO, consultez StableHLO – Comparer.

Ge

Voir aussi XlaBuilder::Ge.

Effectue une comparaison greater-or-equal-than élément par élément de lhs et rhs.

\(lhs >= rhs\)

Ge(lhs, rhs)

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 résulte 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.

Il existe une variante alternative avec une prise en charge de la diffusion de différentes dimensions pour Ge :

Ge(lhs,rhs, broadcast_dimensions)

Cette variante de l'opération doit être utilisée pour les opérations arithmétiques entre des tableaux de rangs différents (par exemple, l'addition d'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 de diffusion.

Une commande totale sur les nombres à virgule flottante existe pour Gt, en appliquant :

\[-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.\]

GtTotalOrder(lhs,rhs, broadcast_dimensions)

Pour en savoir plus sur StableHLO, consultez StableHLO – Comparer.

Gt

Voir aussi XlaBuilder::Gt.

Effectue une comparaison supérieure à élément par élément de lhs et rhs.

\(lhs > rhs\)

Gt(lhs, rhs)

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 résulte 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.

Il existe une variante alternative avec une prise en charge de la diffusion multidimensionnelle pour Gt :

Gt(lhs,rhs, broadcast_dimensions)

Cette variante de l'opération doit être utilisée pour les opérations arithmétiques entre des tableaux de rangs différents (par exemple, l'addition d'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 de diffusion.

Pour en savoir plus sur StableHLO, consultez StableHLO – Comparer.

Le

Voir aussi XlaBuilder::Le.

Effectue une comparaison less-or-equal-than élément par élément de lhs et rhs.

\(lhs <= rhs\)

Le(lhs, rhs)

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 résulte 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.

Il existe une variante alternative avec une prise en charge de la diffusion de différentes dimensions pour Le :

Le(lhs,rhs, broadcast_dimensions)

Cette variante de l'opération doit être utilisée pour les opérations arithmétiques entre des tableaux de rangs différents (par exemple, l'addition d'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 de diffusion.

La prise en charge d'un ordre total sur les nombres à virgule flottante existe pour Le, en appliquant :

\[-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.\]

LeTotalOrder(lhs,rhs, broadcast_dimensions)

Pour en savoir plus sur StableHLO, consultez StableHLO – Comparer.

Lt

Voir aussi XlaBuilder::Lt.

Effectue une comparaison inférieure à élément par élément de lhs et rhs.

\(lhs < rhs\)

Lt(lhs, rhs)

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 résulte 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.

Il existe une variante alternative avec une prise en charge de la diffusion multidimensionnelle pour Lt :

Lt(lhs,rhs, broadcast_dimensions)

Cette variante de l'opération doit être utilisée pour les opérations arithmétiques entre des tableaux de rangs différents (par exemple, l'addition d'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 de diffusion.

Une commande totale sur les nombres à virgule flottante existe pour Lt, en appliquant :

\[-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.\]

LtTotalOrder(lhs,rhs, broadcast_dimensions)

Pour en savoir plus sur StableHLO, consultez StableHLO – Comparer.

Complexe

Voir aussi XlaBuilder::Complex.

Effectue une conversion élément par élément en valeur complexe à partir d'une paire de valeurs réelles et imaginaires, lhs et rhs.

Complex(lhs, rhs)

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 résulte 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.

Il existe une variante alternative avec une prise en charge de la diffusion multidimensionnelle pour Complex :

Complex(lhs,rhs, broadcast_dimensions)

Cette variante de l'opération doit être utilisée pour les opérations arithmétiques entre des tableaux de rangs différents (par exemple, l'addition d'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 de diffusion.

Pour en savoir plus sur StableHLO, consultez StableHLO – complexe.

ConcatInDim (Concaténation)

Voir aussi XlaBuilder::ConcatInDim.

La concaténation compose un tableau à partir de plusieurs opérandes de tableau. Le tableau comporte le même nombre de dimensions que chacun des opérandes du tableau d'entrée (qui doivent avoir le même nombre de dimensions les uns que les autres) et contient les arguments dans l'ordre dans lequel ils ont été spécifiés.

Concatenate(operands..., dimension)

À l'exception de dimension, toutes les dimensions doivent être identiques. En effet, XLA n'est pas compatible avec les tableaux "irréguliers". Notez également que les valeurs de dimension 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)
//Output:  {2, 3, 4, 5, 6, 7}

Exemple bidimensionnel :

let a = { {1, 2},
         {3, 4},
         {5, 6} };

let b = { {7, 8} };

Concat({a, b}, 0)

//Output:  { {1, 2},
//          {3, 4},
//          {5, 6},
//          {7, 8} }

Diagramme :

Pour en savoir plus sur StableHLO, consultez StableHLO – concatenate.

Conditionnel

Voir aussi XlaBuilder::Conditional.

Conditional(predicate, true_operand, true_computation, false_operand, false_computation)

Exécute true_computation si predicate est true, false_computation si predicate 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 identique.

Notez que seul l'un des éléments true_computation et false_computation sera exécuté en fonction de la valeur de predicate.

Conditional(branch_index, branch_computations, branch_operands)

Exécute branch_computations[branch_index] et renvoie le résultat. Si branch_index est un S32 qui est < 0 ou >= N, alors branch_computations[N-1] est exécuté en tant que branche par défaut.

Chaque branch_computations[b] doit accepter un seul argument de type \(T_b\) et sera appelé avec branch_operands[b], qui doit être du même type. Le type de la valeur renvoyée de chaque branch_computations[b] doit être identique.

Notez qu'une seule des branch_computations sera exécutée en fonction de la valeur de branch_index.

Pour en savoir plus sur StableHLO, consultez StableHLO – if.

Constante

Voir aussi XlaBuilder::ConstantLiteral.

Génère un output à partir d'un literal constant.

Constant(literal)

Pour en savoir plus sur StableHLO, consultez StableHLO – constant.

ConvertElementType

Voir aussi XlaBuilder::ConvertElementType.

Semblable à un static_cast élément par élément en C++, ConvertElementType effectue une opération de conversion élément par élément d'une forme de données vers une forme cible. Les dimensions doivent correspondre, et la conversion s'effectue élément 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)

Les dimensions de l'opérande et de la forme cible doivent correspondre. Les types d'éléments source et de destination ne doivent pas être des tuples.

Une conversion telle que T=s32 vers U=f32 effectuera une routine de conversion int-to-float normalisatrice telle que round-to-nearest-even.

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

Pour en savoir plus sur StableHLO, consultez StableHLO – Convertir.

Conv (convolution)

Voir aussi XlaBuilder::Conv.

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.

Conv Met en file d'attente une instruction de convolution dans le calcul, qui utilise les nombres de dimensions de convolution par défaut sans dilatation.

Le remplissage est spécifié de manière abrégée sous la forme SAME ou VALID. Le remplissage SAME ajoute des zéros à l'entrée (lhs) afin que la sortie ait la même forme que l'entrée sans tenir compte du pas. Le remplissage VALID signifie simplement qu'il n'y a pas de remplissage.

Conv(lhs, rhs, window_strides, padding, feature_group_count, batch_group_count, precision_config, preferred_element_type)

Différents niveaux de contrôle sont disponibles pour Conv :

• ConvWithGeneralPadding
• ConvWithGeneralDimensions
• ConvGeneral
• ConvGeneralDilated

Soit n le nombre de dimensions spatiales. L'argument lhs est un tableau à (n+2) dimensions 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 dimensions n+2 sont, dans cet ordre :

• batch : chaque coordonnée de cette dimension représente une entrée indépendante pour laquelle la convolution est effectuée.
• z/depth/features : chaque position (y,x) de la zone de base est associée à un vecteur, qui est inclus dans cette dimension.
• spatial_dims : décrit les dimensions spatiales n qui définissent la zone de base sur laquelle la fenêtre se déplace.

L'argument rhs est un tableau à (n+2) dimensions décrivant le filtre/noyau/fenêtre de convolution. Les dimensions sont, dans cet ordre :

• output-z : dimension z de la sortie.
• input-z : la taille de cette dimension multipliée par feature_group_count doit être égale à la taille de la dimension z dans lhs.
• spatial_dims : décrit les dimensions spatiales n qui définissent la fenêtre n-d qui se déplace sur la zone de base.

L'argument window_strides spécifie le pas de la fenêtre de convolution dans les dimensions spatiales. Par exemple, si le pas de la première dimension spatiale est de 3, la fenêtre ne peut être placée qu'à des coordonnées où le premier index spatial est divisible par 3.

L'argument padding spécifie la quantité de marge intérieure à appliquer à la zone de base. La quantité de marge intérieure peut être négative. La valeur absolue de la marge intérieure négative indique le nombre d'éléments à supprimer de la dimension spécifiée avant d'effectuer la convolution. padding[0] spécifie la marge intérieure pour la dimension y et padding[1] spécifie la marge intérieure pour la dimension x. Chaque paire 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 le sens des indices les plus faibles, tandis que la marge intérieure haute est appliquée dans le sens des indices les plus élevés. Par exemple, si padding[1] est défini sur (2,3), une marge intérieure de deux zéros sera ajoutée à gauche et de trois zéros à droite dans la deuxième dimension spatiale. L'utilisation d'un 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 respectivement aux côtés gauche et droit dans chaque dimension spatiale. Si le facteur de dilatation dans une dimension spatiale est d, alors d-1 trous sont implicitement placés entre chacune des entrées de cette dimension, ce qui augmente la taille du tableau. Les trous sont remplis avec une valeur no-op, qui pour la convolution signifie des zéros.

La dilatation du côté droit est également appelée convolution à trous. Pour en savoir plus, consultez tf.nn.atrous_conv2d. La dilatation du côté gauche est également appelée convolution transposée. Pour en savoir plus, consultez tf.nn.conv2d_transpose.

L'argument feature_group_count (valeur par défaut : 1) peut être utilisé pour les convolutions groupées. feature_group_count doit être un diviseur de la dimension de la caractéristique d'entrée et de sortie. Si feature_group_count est supérieur à 1, cela signifie que, conceptuellement, la dimension de la caractéristique d'entrée et de sortie, ainsi que la dimension de la caractéristique de sortie rhs, sont divisées de manière égale en plusieurs groupes feature_group_count, chaque groupe étant constitué d'une sous-séquence consécutive de caractéristiques. La dimension de la caractéristique d'entrée rhs doit être égale à la dimension de la 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-th 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éristiques de sortie.

Pour la convolution depthwise, l'argument feature_group_count serait défini sur la dimension de la caractéristique d'entrée, et le filtre serait remodelé de [filter_height, filter_width, in_channels, channel_multiplier] à [filter_height, filter_width, 1, in_channels * channel_multiplier]. Pour en savoir plus, consultez tf.nn.depthwise_conv2d.

L'argument batch_group_count (valeur par défaut 1) peut être utilisé pour les filtres groupés lors de la rétropropagation. batch_group_count doit être un diviseur de la taille de la dimension de lot lhs (entrée). Si batch_group_count est supérieur à 1, cela signifie que la dimension du lot de sortie doit être de taille input batch / batch_group_count. La valeur batch_group_count doit être un diviseur de la taille de la caractéristique de sortie.

La forme de sortie présente les dimensions suivantes, dans cet ordre :

• batch : la taille de cette dimension multipliée par batch_group_count doit être égale à la taille de la dimension batch dans lhs.
• z : même taille que 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 comment fonctionne le champ batch_group_count. En effet, nous découpons chaque batch lhs en batch_group_count groupes et faisons de même pour les caractéristiques de sortie. Ensuite, pour chacun de ces groupes, nous effectuons des convolutions par paires et concaténons la sortie le long de la dimension de la 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 foulées et la taille de la zone de base après le remplissage.

Pour décrire ce qu'est une convolution, considérons une convolution 2D et choisissons des coordonnées batch, z, y, x fixes dans la sortie. (y,x) correspond à la position d'un angle de la fenêtre dans la zone de base (par exemple, l'angle supérieur gauche, selon votre interprétation des dimensions spatiales). Nous avons maintenant une fenêtre 2D, extraite de la zone de base, où chaque point 2D est associé à un vecteur 1D, ce qui nous donne une boîte 3D. À partir du noyau de convolution, puisque nous avons fixé la coordonnée de sortie z, nous avons également une boîte 3D. Les deux boîtes ont les mêmes dimensions. Nous pouvons donc calculer la somme des produits élément 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 de valeurs 3D distincte est utilisée pour chaque coordonnée output-z. Vous pouvez donc le considérer comme cinq convolutions distinctes, chacune avec un filtre différent.

Voici un pseudo-code pour une convolution 2D avec remplissage et foulée :

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

precision_config est utilisé pour indiquer la configuration de la précision. Le niveau indique si le matériel doit tenter de générer davantage d'instructions de code machine pour fournir une émulation de type de données plus précise si nécessaire (par exemple, émuler f32 sur un TPU qui ne prend en charge que les matmuls bf16). Les valeurs peuvent être DEFAULT, HIGH ou HIGHEST. Pour en savoir plus, consultez les sections sur les MXU.

preferred_element_type est un élément scalaire de types de sortie de précision supérieure/inférieure utilisé pour l'accumulation. preferred_element_type recommande le type d'accumulation pour l'opération donnée, mais cela n'est pas garanti. Cela permet à certains backends matériels de s'accumuler dans un autre type et de se convertir au type de sortie préféré.

Pour en savoir plus sur StableHLO, consultez StableHLO – convolution.

ConvWithGeneralPadding

Voir aussi XlaBuilder::ConvWithGeneralPadding.

ConvWithGeneralPadding(lhs, rhs, window_strides, padding, feature_group_count, batch_group_count, precision_config, preferred_element_type)

Identique à Conv où la configuration de la marge intérieure est explicite.

ConvWithGeneralDimensions

Voir aussi XlaBuilder::ConvWithGeneralDimensions.

ConvWithGeneralDimensions(lhs, rhs, window_strides, padding, dimension_numbers, feature_group_count, batch_group_count, precision_config, preferred_element_type)

Identique à Conv où les numéros de dimension sont explicites.

ConvGeneral

Voir aussi XlaBuilder::ConvGeneral.

ConvGeneral(lhs, rhs, window_strides, padding, dimension_numbers, feature_group_count, batch_group_count, precision_config, preferred_element_type)

Identique à Conv où les numéros de dimension et la configuration du remplissage sont explicites

ConvGeneralDilated

Voir aussi XlaBuilder::ConvGeneralDilated.

ConvGeneralDilated(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count, batch_group_count, precision_config, preferred_element_type, window_reversal)

Identique à Conv où la configuration du remplissage, les facteurs de dilatation et les numéros de dimension sont explicites.

Copier

Voir aussi HloInstruction::CreateCopyStart.

Copy est décomposé en interne en deux instructions HLO : CopyStart et CopyDone. Copy, CopyStart et CopyDone servent de primitives dans HLO. Ces opérations peuvent apparaître dans les dumps HLO, mais elles ne sont pas destinées à être construites manuellement par les utilisateurs finaux.

COS

Voir aussiXlaBuilder::Cos.

Cosinus par élément x -> cos(x).

Cos(operand)

Cos accepte également l'argument facultatif result_accuracy :

Cos(operand, result_accuracy)

Pour en savoir plus sur result_accuracy, consultez Précision des résultats.

Pour en savoir plus sur StableHLO, consultez StableHLO – cosinus.

Cosh

Voir aussi XlaBuilder::Cosh.

Cosinus hyperbolique élément par élément x -> cosh(x).

Cosh(operand)

Cosh accepte également l'argument facultatif result_accuracy :

Cosh(operand, result_accuracy)

Pour en savoir plus sur result_accuracy, consultez Précision des résultats.

CustomCall

Voir aussi XlaBuilder::CustomCall.

Appelez une fonction fournie par l'utilisateur dans un calcul.

La documentation CustomCall est disponible dans Informations pour les développeurs : appels personnalisés XLA.

Pour en savoir plus sur StableHLO, consultez StableHLO – custom_call.

Div

Voir aussi XlaBuilder::Div.

Effectue une division élément par élément du dividende lhs et du diviseur rhs.

Div(lhs, rhs)

Le dépassement de capacité de la division entière (division/reste signé/non signé par zéro ou division/reste signé de INT_SMIN avec -1) produit une valeur définie par l'implémentation.

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 résulte 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.

Il existe une variante alternative avec une prise en charge de la diffusion multidimensionnelle pour Div :

Div(lhs,rhs, broadcast_dimensions)

Cette variante de l'opération doit être utilisée pour les opérations arithmétiques entre des tableaux de rangs différents (par exemple, l'addition d'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 de diffusion.

Pour en savoir plus sur StableHLO, consultez StableHLO – division.

Domaine

Voir aussi HloInstruction::CreateDomain.

Domain peut apparaître dans les dumps HLO, mais il n'est pas destiné à être construit manuellement par les utilisateurs finaux.

Point

Voir aussi XlaBuilder::Dot.

Dot(lhs, rhs, precision_config, preferred_element_type)

La sémantique exacte de cette opération dépend des rangs des opérandes :

L'opération effectue la somme des produits sur la deuxième dimension de lhs (ou la première si elle comporte une dimension) 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 de vecteurs/matrices ou des multiplications de matrices/matrices.

precision_config est utilisé pour indiquer la configuration de la précision. Le niveau indique si le matériel doit tenter de générer davantage d'instructions de code machine pour fournir une émulation de type de données plus précise si nécessaire (par exemple, émuler f32 sur un TPU qui ne prend en charge que les matmuls bf16). Les valeurs peuvent être DEFAULT, HIGH ou HIGHEST. Pour en savoir plus, consultez les sections sur les MXU.

preferred_element_type est un élément scalaire de types de sortie de précision supérieure/inférieure utilisé pour l'accumulation. preferred_element_type recommande le type d'accumulation pour l'opération donnée, mais cela n'est pas garanti. Cela permet à certains backends matériels de s'accumuler dans un type différent et de se convertir au type de sortie préféré.

Pour en savoir plus sur StableHLO, consultez StableHLO – dot.

DotGeneral

Voir aussi XlaBuilder::DotGeneral.

DotGeneral(lhs, rhs, dimension_numbers, precision_config, preferred_element_type)

Semblable à Dot, mais permet de spécifier les nombres de dimensions de contraction et de lot pour lhs et rhs.

DotGeneral effectue la somme des produits sur les dimensions de contraction spécifiées dans dimension_numbers.

Les numéros de dimension de contrat associés de lhs et rhs n'ont pas besoin d'être identiques, mais doivent avoir les mêmes tailles de dimension.

Exemple avec des nombres de dimensions contractés :

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 nombres de dimensions de lot associés de lhs et rhs doivent avoir les mêmes tailles de dimensions.

Exemple avec des nombres de dimensions de lot (taille de lot 2, matrices 2x2) :

lhs = { { {1.0, 2.0},
          {3.0, 4.0} },
        { {5.0, 6.0},
          {7.0, 8.0} } }

rhs = { { {1.0, 0.0},
          {0.0, 1.0} },
        { {1.0, 0.0},
          {0.0, 1.0} } }

DotDimensionNumbers dnums;
dnums.add_lhs_contracting_dimensions(2);
dnums.add_rhs_contracting_dimensions(1);
dnums.add_lhs_batch_dimensions(0);
dnums.add_rhs_batch_dimensions(0);

DotGeneral(lhs, rhs, dnums) -> {
    { {1.0, 2.0},
      {3.0, 4.0} },
    { {5.0, 6.0},
      {7.0, 8.0} } }

Il s'ensuit que le nombre de dimensions résultant commence par la dimension de lot, puis par la dimension lhs non contractante/non de lot, et enfin par la dimension rhs non contractante/non de lot.

precision_config est utilisé pour indiquer la configuration de la précision. Le niveau indique si le matériel doit tenter de générer davantage d'instructions de code machine pour fournir une émulation de type de données plus précise si nécessaire (par exemple, émuler f32 sur un TPU qui ne prend en charge que les matmuls bf16). Les valeurs peuvent être DEFAULT, HIGH ou HIGHEST. Pour en savoir plus, consultez les sections sur les MXU.

preferred_element_type est un élément scalaire de types de sortie de précision supérieure/inférieure utilisé pour l'accumulation. preferred_element_type recommande le type d'accumulation pour l'opération donnée, mais cela n'est pas garanti. Cela permet à certains backends matériels de s'accumuler dans un autre type et de se convertir au type de sortie préféré.

Pour en savoir plus sur StableHLO, consultez StableHLO – dot_general.

ScaledDot

Voir aussi XlaBuilder::ScaledDot.

ScaledDot(lhs, lhs_scale, rhs, rhs_scale, dimension_number, precision_config,preferred_element_type)

Semblable à DotGeneral.

Crée une opération de produit scalaire avec les opérandes "lhs", "lhs_scale", "rhs" et "rhs_scale", avec les dimensions de contraction et de lot spécifiées dans "dimension_numbers".

RaggedDot

Voir aussi XlaBuilder::RaggedDot.

Pour obtenir une répartition du calcul RaggedDot, consultez StableHLO – chlo.ragged_dot.

DynamicReshape

Voir aussi XlaBuilder::DynamicReshape.

Cette opération est fonctionnellement identique à reshape, mais la forme du résultat est spécifiée de manière dynamique via output_shape.

DynamicReshape(operand, dim_sizes, new_size_bounds, dims_are_dynamic)

Pour en savoir plus sur StableHLO, consultez StableHLO – dynamic_reshape.

DynamicSlice

Voir aussi XlaBuilder::DynamicSlice.

DynamicSlice extrait un sous-tableau du tableau d'entrée à start_indices dynamique. La taille de la tranche dans chaque dimension est transmise dans size_indices, qui spécifie le point de terminaison des intervalles de tranche exclusifs dans chaque dimension : [start, start + size). La forme de start_indices doit être unidimensionnelle, avec une taille de dimension égale au nombre de dimensions de operand.

DynamicSlice(operand, start_indices, slice_sizes)

Les indices de tranche effectifs sont calculés en appliquant la transformation suivante pour chaque indice i dans [1, N) avant d'effectuer le découpage :

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

Cela garantit que la tranche extraite est toujours dans les limites du tableau d'opérandes. Si la tranche est dans les limites avant l'application de la transformation, celle-ci n'a aucun effet.

Exemple à une dimension :

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

DynamicSlice(a, s, {2});
// Result: {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});
//Result:
// { { 7.0,  8.0},
//   {10.0, 11.0} }

Pour en savoir plus sur StableHLO, consultez StableHLO – dynamic_slice.

DynamicUpdateSlice

Voir aussi XlaBuilder::DynamicUpdateSlice.

DynamicUpdateSlice génère un résultat qui correspond à 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 unidimensionnelle, avec une taille de dimension égale au nombre de dimensions de operand.

DynamicUpdateSlice(operand, update, start_indices)

Les indices de tranche effectifs sont calculés en appliquant la transformation suivante pour chaque indice i dans [1, N) avant d'effectuer le découpage :

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 du tableau d'opérandes. Si la tranche est dans les limites avant l'application de la transformation, celle-ci n'a aucun effet.

Exemple à une dimension :

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)
// Result: {0.0, 1.0, 5.0, 6.0, 4.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 u =
{ {12.0, 13.0},
  {14.0, 15.0},
  {16.0, 17.0} }

let s = {1, 1}

DynamicUpdateSlice(b, u, s)
// Result:
// { {0.0,  1.0,  2.0},
//   {3.0, 12.0, 13.0},
//   {6.0, 14.0, 15.0},
//   {9.0, 16.0, 17.0} }

Pour en savoir plus sur StableHLO, consultez StableHLO – dynamic_update_slice.

Erf

Voir aussi XlaBuilder::Erf.

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\).

Erf(operand)

Erf accepte également l'argument facultatif result_accuracy :

Erf(operand, result_accuracy)

Pour en savoir plus sur result_accuracy, consultez Précision des résultats.

Exp.

Voir aussi XlaBuilder::Exp.

Exponentielle naturelle par élément x -> e^x.

Exp(operand)

Exp accepte également l'argument facultatif result_accuracy :

Exp(operand, result_accuracy)

Pour en savoir plus sur result_accuracy, consultez Précision des résultats.

Pour en savoir plus sur StableHLO, consultez StableHLO – exponentiel.

Expm1

Voir aussi XlaBuilder::Expm1.

Exponentielle naturelle par élément moins un x -> e^x - 1.

Expm1(operand)

Expm1 accepte également l'argument facultatif result_accuracy :

Expm1(operand, result_accuracy)

Pour en savoir plus sur result_accuracy, consultez Précision des résultats.

Pour en savoir plus sur StableHLO, consultez StableHLO – exponential_minus_one.

Fft

Voir aussi XlaBuilder::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 sont acceptées sur un maximum de trois axes.

Fft(operand, ftt_type, fft_length)

Pour en savoir plus sur StableHLO, consultez StableHLO – fft.

FFT multidimensionnelle

Lorsque plusieurs fft_length sont fournis, cela équivaut à appliquer une cascade d'opérations FFT à chacun des axes les plus intérieurs. Notez que pour les cas réel->complexe et complexe->réel, la transformation de l'axe le plus intérieur est (effectivement) effectuée en premier (RFFT ; en dernier pour IRFFT), c'est pourquoi l'axe le plus intérieur est celui dont la taille change. Les autres transformations d'axe seront alors complexes->complexes.

Détails de mise en œuvre

La FFT du processeur est basée sur TensorFFT d'Eigen. GPU FFT utilise cuFFT.

Étage

Voir aussi XlaBuilder::Floor.

Plancher par élément x -> ⌊x⌋.

Floor(operand)

Pour en savoir plus sur StableHLO, consultez StableHLO – floor.

Fusion

Voir aussi HloInstruction::CreateFusion.

L'opération Fusion représente les instructions HLO et sert de primitive dans HLO. Cette opération peut apparaître dans les dumps HLO, mais n'est pas destinée à être construite manuellement par les utilisateurs finaux.

Collecter

L'opération XLA gather assemble plusieurs tranches (chacune à un décalage d'exécution potentiellement différent) d'un tableau d'entrée.

Pour en savoir plus sur StableHLO, consultez StableHLO – gather.

Sémantique générale

Voir aussi XlaBuilder::Gather. Pour une description plus intuitive, consultez la section "Description informelle" ci-dessous.

gather(operand, start_indices, dimension_numbers, slice_sizes, indices_are_sorted)

Pour plus de commodité, nous désignons les dimensions du tableau de sortie qui ne figurent pas dans offset_dims par batch_dims.

La sortie est un tableau avec batch_dims.size + offset_dims.size dimensions.

La valeur operand.rank doit être égale à la somme de offset_dims.size et collapsed_slice_dims.size. De plus, slice_sizes.size doit être égal à operand.rank.

Si index_vector_dim est égal à start_indices.rank, nous considérons implicitement que start_indices a une dimension 1 de fin (c'est-à-dire que si start_indices était de 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 le long de la dimension i sont calculées comme suit :

1. Si i est présent dans batch_dims (c'est-à-dire égal à batch_dims[k] pour certains k), nous sélectionnons les limites de dimension correspondantes dans start_indices.shape, en ignorant index_vector_dim (c'est-à-dire en sélectionnant start_indices.shape.dims[k] si k < index_vector_dim et start_indices.shape.dims[k+1] sinon).

2. Si i est présent dans offset_dims (c'est-à-dire égal à offset_dims[k] pour un certain k), nous choisissons la limite correspondante dans slice_sizes après avoir tenu compte de collapsed_slice_dims (c'est-à-dire que nous choisissons adjusted_slice_sizes[k] où adjusted_slice_sizes est slice_sizes avec les limites aux index collapsed_slice_dims supprimées).

Formellement, l'index d'opérande In correspondant à un index de sortie donné Out est calculé comme suit :

1. Soit G = { Out[k] pour k dans batch_dims }. Utilisez G pour extraire un vecteur S tel que S[i] = start_indices[Combine(G, i)] où Combine(A, b) insère b à la position index_vector_dim dans A. Notez que cela est bien défini même si G est vide : si G est vide, alors S = start_indices.

2. Créez un index de départ, Sin, dans operand en utilisant S en dispersant S à l'aide de start_index_map. Plus précisément :

   1. Sin[start_index_map[k]] = S[k] si k < start_index_map.size.

   2. Sin[_] = 0 sinon.

3. Crée un index Oin dans operand en dispersant les index au niveau des dimensions de décalage dans Out selon l'ensemble collapsed_slice_dims. Plus précisément :

   1. Oin[remapped_offset_dims(k)] = Out[offset_dims[k]] si k < offset_dims.size (remapped_offset_dims est défini ci-dessous).

   2. Oin[_] = 0 sinon.

4. In est Oin + Sin, où + correspond à l'addition élément par élément.

remapped_offset_dims est une fonction monotone avec le domaine [0, offset_dims.size) et la plage [0, operand.rank) \ collapsed_slice_dims. Par exemple, offset_dims.size est 4, operand.rank est 6 et collapsed_slice_dims est {0, 2}, 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 sont triés (par ordre croissant, après la dispersion de ses valeurs selon start_index_map) par l'utilisateur. Si ce n'est pas le cas, la sémantique est définie par l'implémentation.

Description informelle et exemples

De manière informelle, chaque index 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 index de départ à partir de start_indices.

• Nous utilisons start_index_map pour mapper l'index de départ (dont la taille peut être inférieure à operand.rank) à un index de départ "complet" dans operand.

• Nous découpons dynamiquement une tranche de taille slice_sizes en utilisant l'index de début complet.

• Nous remodelons le slice en réduisant les dimensions collapsed_slice_dims. Étant donné que toutes les dimensions de tranche réduites doivent avoir une limite de 1, cette redimension est toujours légale.

• 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 sortie Out.

index_vector_dim est défini sur start_indices.rank – 1 dans tous les exemples suivants. Des valeurs plus intéressantes pour index_vector_dim ne modifient pas fondamentalement l'opération, mais rendent la représentation visuelle plus complexe.

Pour mieux comprendre comment tout cela s'articule, examinons un exemple qui collecte cinq tranches de forme [8,6] à partir d'un tableau [16,11]. La position d'un segment dans le tableau [16,11] peut être représentée sous la forme d'un vecteur d'index de forme S64[2]. L'ensemble de cinq positions peut donc être représenté sous la forme d'un tableau S64[5,2].

Le comportement de l'opération de collecte peut ensuite être représenté comme une transformation d'index qui prend [G,O0,O1], un index dans la forme de sortie, et le mappe à un élément du tableau d'entrée de la manière suivante :

Nous sélectionnons d'abord un vecteur (X,Y) à partir du tableau d'indices de collecte à l'aide de G. L'élément du tableau de sortie à l'index [G,O0,O1] est alors l'élément du tableau d'entrée à l'index [X+O0,Y+O1].

slice_sizes est [8,6], ce qui détermine la plage de O₀ et O₁, qui à son tour détermine les limites de la tranche.

Cette opération de collecte agit comme une tranche dynamique par lot avec G comme dimension de lot.

Les index de collecte peuvent être multidimensionnels. Par exemple, une version plus générale de l'exemple ci-dessus utilisant un tableau "gather indices" de forme [4,5,2] traduirait les indices comme suit :

Ici encore, cela sert de tranche dynamique de lot G0 et G1 en tant que dimensions de lot. La taille du segment est toujours [8,6].

L'opération de collecte dans XLA généralise la sémantique informelle décrite ci-dessus de la manière suivante :

1. Nous pouvons configurer les dimensions de la forme de sortie qui sont les dimensions de décalage (dimensions contenant O0, O1 dans le dernier exemple). Les dimensions de lot de sortie (dimensions contenant G0, G1 dans le dernier exemple) sont définies comme les dimensions de sortie qui ne sont pas des dimensions de décalage.

2. Le nombre de dimensions de décalage de sortie explicitement présentes dans la forme de sortie peut être inférieur au nombre de dimensions d'entrée. Ces dimensions "manquantes", qui sont listées explicitement comme collapsed_slice_dims, doivent avoir une taille de tranche de 1. Comme leur taille de tranche est de 1, le seul index valide pour elles est 0. Leur élision n'introduit aucune ambiguïté.

3. La tranche extraite du tableau "Gather Indices" ((X, Y) dans le dernier exemple) peut comporter moins d'éléments que le nombre de dimensions du tableau d'entrée. Un mappage explicite indique comment l'index doit être étendu pour avoir le même nombre de dimensions que l'entrée.

Pour le dernier exemple, nous utilisons (2) et (3) pour implémenter tf.gather_nd :

G0 et G1 sont utilisés pour extraire un index de départ du tableau d'index de collecte comme d'habitude, sauf que l'index de départ ne comporte qu'un seul élément, X. De même, il n'existe qu'un seul index de décalage de sortie avec la valeur O0. Toutefois, avant d'être utilisées comme index dans le tableau d'entrée, elles sont développées conformément à "Gather Index Mapping" (start_index_map dans la description formelle) et "Offset Mapping" (remapped_offset_dims dans la description formelle) en [X,0] et [0,O0] respectivement, ce qui donne [X,O0]. En d'autres termes, l'index de sortie [G0,G1,O0] correspond à l'index d'entrée [GatherIndices[G0,G1,0],O0], ce qui nous donne la sémantique pour tf.gather_nd.

Le slice_sizes pour cette demande est [1,11]. Intuitivement, cela signifie que chaque index X du tableau d'index de collecte sélectionne une ligne entière, et que 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 de forme matricielle.

GetDimensionSize(operand, dimension)

Pour en savoir plus sur StableHLO, consultez StableHLO – get_dimension_size.

GetTupleElement

Voir aussi XlaBuilder::GetTupleElement.

Index dans un tuple avec une valeur constante au moment de la compilation.

La valeur doit être une constante de compilation afin que l'inférence de forme puisse déterminer le type de la valeur résultante.

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

GetTupleElement(tuple_data, index)

Pour en savoir plus sur StableHLO, consultez StableHLO – get_tuple_element.

Imag

Voir aussi XlaBuilder::Imag.

Partie imaginaire élément par élément d'une forme complexe (ou réelle). x -> imag(x). Si l'opérande est de type à virgule flottante, renvoie 0.

Imag(operand)

Pour en savoir plus sur StableHLO, consultez StableHLO – imag.

In-Feed

Voir aussi XlaBuilder::Infeed.

Infeed(shape, config)

Lit un seul élément de données à partir de l'interface de streaming Infeed implicite de l'appareil, interprète les données comme la forme et la mise en page données, et 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 Infeed dans le code ci-dessous ont un ordre total, 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 tuples imbriquées ne sont pas acceptées. Pour une forme de tuple vide, l'opération Infeed est effectivement une opération sans effet et se poursuit sans lire aucune donnée à partir de l'Infeed de l'appareil.

Pour en savoir plus sur StableHLO, consultez StableHLO – Infeed.

Iota

Voir aussi XlaBuilder::Iota.

Iota(shape, iota_dimension)

Crée un littéral constant sur l'appareil plutôt qu'un transfert d'hôte potentiellement volumineux. Crée un tableau avec la forme spécifiée et contenant des valeurs commençant à zéro et augmentant d'une unité 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 de type à virgule flottante.

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

Pour en savoir plus sur StableHLO, consultez StableHLO – iota.

IsFinite

Voir aussi XlaBuilder::IsFinite.

Teste si chaque élément de operand est fini, c'est-à-dire s'il n'est pas une infinité positive ou négative, et s'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.

IsFinite(operand)

Pour en savoir plus sur StableHLO, consultez StableHLO – is_finite.

Journal

Voir aussi XlaBuilder::Log.

Logarithme naturel par élément x -> ln(x).

Log(operand)

Log accepte également l'argument facultatif result_accuracy :

Log(operand, result_accuracy)

Pour en savoir plus sur result_accuracy, consultez Précision des résultats.

Pour en savoir plus sur StableHLO, consultez StableHLO – journal.

Log1p

Voir aussi XlaBuilder::Log1p.

Logarithme naturel décalé par élément x -> ln(1+x).

Log1p(operand)

Log1p accepte également l'argument facultatif result_accuracy :

Log1p(operand, result_accuracy)

Pour en savoir plus sur result_accuracy, consultez Précision des résultats.

Pour en savoir plus sur StableHLO, consultez StableHLO – log_plus_one.

Logistique

Voir aussi XlaBuilder::Logistic.

Calcul de la fonction logistique par élément x -> logistic(x).

Logistic(operand)

Logistic accepte également l'argument facultatif result_accuracy :

Logistic(operand, result_accuracy)

Pour en savoir plus sur result_accuracy, consultez Précision des résultats.

Pour en savoir plus sur StableHLO, consultez StableHLO – logistique.

Carte

Voir aussi XlaBuilder::Map.

Map(operands..., computation, dimensions)

Applique une fonction scalaire aux tableaux operands donnés, ce qui produit 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, sauf que le type d'élément T est remplacé par S.

Par exemple, Map(op1, op2, op3, computation, par1) mappe elem_out <- computation(elem1, elem2, elem3, par1) à chaque index (multidimensionnel) des tableaux d'entrée pour générer le tableau de sortie.

Pour en savoir plus sur StableHLO, consultez StableHLO – map.

Max

Voir aussi XlaBuilder::Max.

Effectue une opération max au niveau des éléments sur les Tensors lhs et rhs.

Max(lhs, rhs)

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 résulte 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.

Une variante alternative avec une prise en charge de la diffusion de différentes dimensions existe pour Max :

Max(lhs,rhs, broadcast_dimensions)

Cette variante de l'opération doit être utilisée pour les opérations arithmétiques entre des tableaux de rangs différents (par exemple, l'addition d'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 de diffusion.

Pour en savoir plus sur StableHLO, consultez StableHLO – maximum.

Min

Voir aussi XlaBuilder::Min.

Effectue une opération min au niveau des éléments sur lhs et rhs.

Min(lhs, rhs)

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 résulte 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.

Il existe une variante alternative avec une prise en charge de la diffusion multidimensionnelle pour Min :

Min(lhs,rhs, broadcast_dimensions)

Cette variante de l'opération doit être utilisée pour les opérations arithmétiques entre des tableaux de rangs différents (par exemple, l'addition d'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 de diffusion.

Pour en savoir plus sur StableHLO, consultez StableHLO – minimum.

Mul

Voir aussi XlaBuilder::Mul.

Effectue le produit élément par élément de lhs et rhs.

Mul(lhs, rhs)

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 résulte 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.

Il existe une variante alternative avec une prise en charge de la diffusion multidimensionnelle pour Mul :

Mul(lhs,rhs, broadcast_dimensions)

Cette variante de l'opération doit être utilisée pour les opérations arithmétiques entre des tableaux de rangs différents (par exemple, l'addition d'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 de diffusion.

Pour en savoir plus sur StableHLO, consultez StableHLO – multiply.

Neg

Voir aussi XlaBuilder::Neg.

Négation par élément x -> -x.

Neg(operand)

Pour en savoir plus sur StableHLO, consultez StableHLO – négation.

Non

Voir aussi XlaBuilder::Not.

Logique par élément NOT x -> !(x).

Not(operand)

Pour en savoir plus sur StableHLO, consultez StableHLO – not.

OptimizationBarrier

Voir aussi XlaBuilder::OptimizationBarrier.

Empêche tout pass d'optimisation de déplacer des calculs au-delà de la barrière.

OptimizationBarrier(operand)

Garantit que toutes les entrées sont évaluées avant tout opérateur qui dépend des sorties de la barrière.

Pour en savoir plus sur StableHLO, consultez StableHLO – optimization_barrier.

Ou

Voir aussi XlaBuilder::Or.

Effectue un OR élément par élément de lhs et rhs .

Or(lhs, rhs)

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 résulte 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.

Il existe une variante alternative avec une prise en charge de la diffusion multidimensionnelle pour Or :

Or(lhs,rhs, broadcast_dimensions)

Cette variante de l'opération doit être utilisée pour les opérations arithmétiques entre des tableaux de rangs différents (par exemple, l'addition d'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 de diffusion.

Pour en savoir plus sur StableHLO, consultez StableHLO.

Outfeed

Voir aussi XlaBuilder::Outfeed.

Écrit les entrées dans le flux de sortie.

Outfeed(operand, shape_with_layout, outfeed_config)

shape_with_layout communique la forme mise en page que nous voulons sortir.

Pour en savoir plus sur StableHLO, consultez StableHLO – Outfeed.

Pad

Voir aussi XlaBuilder::Pad.

Pad(operand, padding_value, padding_config)

Développe le tableau operand donné en ajoutant une marge autour du tableau ainsi qu'entre les éléments du tableau avec le padding_value donné. padding_config spécifie la quantité de marge intérieure et de marge extérieure pour chaque dimension.

PaddingConfig est un champ répété de PaddingConfigDimension, qui contient trois champs pour chaque dimension : edge_padding_low, edge_padding_high et interior_padding.

edge_padding_low et edge_padding_high spécifient la quantité de marge intérieure ajoutée respectivement à l'extrémité inférieure (à côté de l'index 0) et à l'extrémité supérieure (à côté de l'index le plus élevé) de chaque dimension. La valeur de la marge intérieure peut être négative. La valeur absolue d'une marge intérieure négative indique le nombre d'éléments à supprimer de la dimension spécifiée.

interior_padding spécifie la marge intérieure ajoutée entre deux éléments dans chaque dimension. Elle ne peut pas être négative. La marge intérieure se produit logiquement avant la marge extérieure. Ainsi, en cas de marge extérieure négative, les éléments sont supprimés de l'opérande avec marge intérieure.

Cette opération est sans effet si les paires de marge intérieure sont toutes (0, 0) et que les valeurs de marge intérieure sont toutes égales à 0. La figure ci-dessous montre des exemples de différentes valeurs edge_padding et interior_padding pour un tableau bidimensionnel.

Pour en savoir plus sur StableHLO, consultez StableHLO – pad.

Paramètre

Voir aussi XlaBuilder::Parameter.

Parameter représente un argument d'entrée pour un calcul.

PartitionID

Voir aussi XlaBuilder::BuildPartitionId.

Produit partition_id du processus actuel.

PartitionID(shape)

PartitionID peut apparaître dans les dumps HLO, mais il n'est pas destiné à être construit manuellement par les utilisateurs finaux.

Pour en savoir plus sur StableHLO, consultez StableHLO – partition_id.

PopulationCount

Voir aussi XlaBuilder::PopulationCount.

Calcule le nombre de bits définis dans chaque élément de operand.

PopulationCount(operand)

Pour en savoir plus sur StableHLO, consultez StableHLO – popcnt.

Pow

Voir aussi XlaBuilder::Pow.

Effectue une exponentiation élément par élément de lhs par rhs.

Pow(lhs, rhs)

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 résulte 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.

Il existe une variante alternative avec une prise en charge de la diffusion multidimensionnelle pour Pow :

Pow(lhs,rhs, broadcast_dimensions)

Cette variante de l'opération doit être utilisée pour les opérations arithmétiques entre des tableaux de rangs différents (par exemple, l'addition d'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 de diffusion.

Pour en savoir plus sur StableHLO, consultez StableHLO – power.

Réel

Voir aussi XlaBuilder::Real.

Partie réelle élément par élément d'une forme complexe (ou réelle). x -> real(x). Si l'opérande est de type à virgule flottante, Real renvoie la même valeur.

Real(operand)

Pour en savoir plus sur StableHLO, consultez StableHLO – réel.

Recv

Voir aussi XlaBuilder::Recv.

Recv, RecvWithTokens et RecvToHost sont des opérations qui servent de primitives de communication dans HLO. Ces opérations apparaissent généralement dans les dumps HLO en tant que partie des transferts d'entrée/sortie de bas niveau ou inter-appareils, mais elles ne sont pas destinées à être construites manuellement par les utilisateurs finaux.

Recv(shape, handle)

Reçoit les données de la forme donnée à partir d'une instruction Send dans un autre calcul qui partage le même handle de canal. Renvoie un XlaOp pour les données reçues.

Pour en savoir plus sur StableHLO, consultez StableHLO – recv.

RecvDone

Voir aussi HloInstruction::CreateRecv et HloInstruction::CreateRecvDone.

Comme pour Send, l'API cliente de l'opération Recv représente une 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.

Recv(const Shape& shape, int64 channel_id)

Alloue les ressources nécessaires pour recevoir des données à partir d'une instruction Send avec le même channel_id. Renvoie un contexte pour les ressources allouées, qui est utilisé par une instruction RecvDone suivante pour attendre la fin du transfert de données. Le contexte est un tuple de {forme du tampon de réception, identifiant de la requête (U32)}. Il ne peut être utilisé que par une instruction RecvDone.

Étant donné un contexte créé par une instruction Recv, attend la fin du transfert de données et renvoie les données reçues.

Réduire

Voir aussi XlaBuilder::Reduce.

Applique une fonction de réduction à un ou plusieurs tableaux en parallèle.

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

Où :

• N doit être supérieur ou égal à 1.
• Le calcul doit être "à peu près" associatif (voir ci-dessous).
• Toutes les matrices d'entrée doivent avoir les mêmes dimensions.
• Toutes les valeurs initiales doivent former une identité sous computation.
• Si N = 1, Collate(T) est T.
• Si N > 1, Collate(T_0, ..., T_{N-1}) est un tuple de N éléments de type T.

Cette opération réduit une ou plusieurs dimensions de chaque tableau d'entrée en scalaires. Le nombre de dimensions de chaque tableau renvoyé est de number_of_dimensions(operand) - len(dimensions). Le résultat de l'opération est Collate(Q_0, ..., Q_N), où Q_i est un tableau de type T_i dont les dimensions sont décrites ci-dessous.

Différents backends sont autorisés à réassocier le calcul de la 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 de données est limitée, l'addition à virgule flottante est suffisamment associative pour la plupart des utilisations pratiques.

Pour en savoir plus sur StableHLO, consultez StableHLO – reduce.

Exemples

Lorsque vous réduisez une dimension dans un tableau 1D unique avec des valeurs [10, 11, 12, 13], avec la fonction de réduction f (qui est computation), cela peut être calculé comme suit :

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

mais il existe également de nombreuses autres possibilités, par exemple :

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

Voici un exemple de pseudo-code approximatif montrant comment 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 number of dimensions of the result.
for r0 in range(result_shape[0]), r1 in range(result_shape[1]), ...:
  # Initialize this result element
  result[r0, r1...] <- 0

  # Iterate over all the reduction dimensions
  for d0 in range(dimensions[0]), d1 in range(dimensions[1]), ...:
    # Increment the result element with the value of the operand's element.
    # The index of the operand's element is constructed from all ri's and di's
    # in the right order (by construction ri's and di's together index over the
    # whole operand shape).
    result[r0, r1...] += operand[ri... di]

Voici un exemple de réduction d'un tableau 2D (matrice). La forme comporte deux dimensions : la dimension 0 de taille 2 et la dimension 1 de 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 à une dimension. Le diagramme montre une colonne et une ligne uniquement pour des raisons de commodité visuelle.

Voici un exemple plus complexe de tableau 3D. Son nombre de dimensions est de 3, avec une dimension 0 de taille 4, une dimension 1 de taille 2 et une dimension 2 de taille 3. Pour plus de simplicité, les valeurs de 1 à 6 sont répliquées dans la dimension 0.

Comme dans l'exemple en 2D, nous pouvons réduire une seule dimension. Si nous réduisons la dimension 0, par exemple, nous obtenons un tableau à deux dimensions dans lequel toutes les valeurs de la dimension 0 ont été regroupées dans un scalaire :

|  4   8  12 |
| 16  20  24 |

Si nous réduisons la dimension 2, nous obtenons également un tableau à deux dimensions dans lequel toutes les valeurs de la dimension 2 ont été réduites à un 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 que de nouveaux numéros peuvent être attribués à certaines dimensions (puisque le nombre de dimensions change).

Nous pouvons également réduire plusieurs dimensions. L'ajout de dimensions de réduction 0 et 1 produit le tableau 1D [20, 28, 36].

La réduction du tableau 3D sur toutes ses dimensions produit le scalaire 84.

Réduction variadique

Lorsque N > 1, l'application de la fonction de réduction est légèrement plus complexe, car elle est appliquée simultanément à toutes les entrées. Les opérandes sont fournis au calcul dans l'ordre suivant :

• Exécution de la valeur réduite pour le premier opérande
• …
• Exécution de la valeur réduite pour le N-ième opérande
• Valeur d'entrée du premier opérande
• …
• Valeur d'entrée pour le Nième opérande

Par exemple, considérez la fonction de réduction suivante, qui peut être utilisée pour calculer le max et l'argmax d'un tableau à une dimension 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 équivaut à 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) co-itérera sur les tableaux et renverra un tuple contenant la valeur maximale et l'index correspondant.

ReducePrecision

Voir aussi XlaBuilder::ReducePrecision.

Modélise l'effet de la conversion de valeurs à virgule flottante dans un format de précision inférieure (tel que IEEE-FP16) et de leur reconversion dans le format d'origine. Le nombre de bits de l'exposant et de la mantisse dans le format de 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, exponent_bits, mantissa_bits)

Le résultat est un tableau de type T. Les valeurs d'entrée sont arrondies à la valeur représentable la plus proche avec le nombre de bits de mantisse donné (en utilisant la sémantique "ties to even"), et toutes les valeurs qui dépassent la plage spécifiée par le nombre de bits d'exposant sont ramenées à l'infini positif ou négatif. Les valeurs NaN sont conservées, mais elles peuvent être converties en valeurs NaN canoniques.

Le format de précision inférieure doit comporter au moins un bit d'exposant (afin de distinguer une valeur nulle d'une valeur infinie, car les deux ont une mantisse nulle) et doit comporter 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 effet.

Pour en savoir plus sur StableHLO, consultez StableHLO – reduce_precision.

ReduceScatter

Voir aussi XlaBuilder::ReduceScatter.

ReduceScatter est une opération collective qui effectue efficacement un AllReduce, puis disperse le résultat en le divisant en blocs shard_count le long de scatter_dimension. La réplique i du groupe de répliques reçoit le shard ith.

ReduceScatter(operand, computation, scatter_dimension, shard_count, replica_groups, channel_id, layout, use_global_device_ids)

• Lorsque operand est un tuple de tableaux, le reduce-scatter est effectué sur chaque élément du tuple.
• replica_groups est une liste de groupes de répliques entre lesquels la réduction est effectuée (l'ID de réplique de la réplique actuelle peut être récupéré à l'aide de ReplicaId). L'ordre des répliques dans chaque groupe détermine l'ordre dans lequel le résultat all-reduce sera dispersé. replica_groups doit être vide (dans ce cas, toutes les répliques appartiennent à un même groupe) ou contenir le même nombre d'éléments que le nombre de répliques. Lorsqu'il existe plusieurs groupes de répliques, ils doivent tous être de la même taille. Par exemple, replica_groups = {0, 2}, {1, 3} effectue une réduction entre les répliques 0 et 2, et 1 et 3, puis distribue le résultat.
• shard_count correspond à la taille de chaque groupe de réplicas. Nous en avons besoin lorsque replica_groups est vide. Si replica_groups n'est pas vide, shard_count doit être égal à la taille de chaque groupe de réplicas.
• channel_id est utilisé pour la communication intermodules : seules les opérations reduce-scatter avec le même channel_id peuvent communiquer entre elles.
• layout Pour en savoir plus sur les mises en page, consultez xla::shapes.
• use_global_device_ids est un indicateur spécifié par l'utilisateur. Lorsque false(valeur par défaut), les nombres dans replica_groups sont ReplicaId lorsque true, les replica_groups représentent un ID global de (ReplicaID*partition_count +partition_id). Par exemple :
  • Avec deux répliques et quatre partitions,
  • replica_groups={ {0,1,4,5},{2,3,6,7} } and use_global_device_ids=true
  • group[0] = (0,0), (0,1), (1,0), (1,1)
  • group[1] = (0,2), (0,3), (1,2), (1,3)
  • où chaque paire est (replica_id, partition_id).

La forme de sortie est la forme d'entrée avec le scatter_dimension shard_count fois plus petit. Par exemple, s'il existe deux répliques et que l'opérande a respectivement la valeur [1.0, 2.25] et [3.0, 5.25] sur les deux répliques, la valeur de sortie de cette opération où scatter_dim est 0 sera [4.0] pour la première réplique et [7.5] pour la seconde.

Pour en savoir plus sur StableHLO, consultez StableHLO – reduce_scatter.

ReduceScatter – Exemple 1 – StableHLO

Dans l'exemple ci-dessus, deux répliques participent à ReduceScatter. Sur chaque réplica, l'opérande a la forme f32[2,4]. Une opération all-reduce (somme) est effectuée sur les répliques, ce qui produit une valeur réduite de forme f32[2,4] sur chaque réplique. Cette valeur réduite est ensuite divisée en deux parties le long de la dimension 1, de sorte que chaque partie a une forme f32[2,2]. Chaque réplica du groupe de processus reçoit la partie correspondant à sa position dans le groupe. Par conséquent, la sortie de chaque réplique a la forme f32[2,2].

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, en produisant un seul tableau multidimensionnel ou un tuple de N tableaux multidimensionnels en sortie. Chaque tableau de sortie comporte 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)

Où :

• N doit être supérieur ou égal à 1.
• Toutes les matrices d'entrée doivent avoir les mêmes dimensions.
• Si N = 1, Collate(T) est T.
• Si N > 1, Collate(T_0, ..., T_{N-1}) est un tuple de N éléments de type (T0,...T{N-1}).

Pour en savoir plus sur StableHLO, consultez StableHLO – reduce_window.

ReduceWindow - Exemple 1

L'entrée est une matrice de taille [4x6], et window_dimensions et window_stride_dimensions sont [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);