Spécification StableHLO

StableHLO est un ensemble d'opérations pour les opérations de haut niveau (HLO) dans les modèles de machine learning (ML). StableHLO sert de couche de portabilité entre différents frameworks et compilateurs de ML : les frameworks de ML qui produisent des programmes StableHLO sont compatibles avec les compilateurs de ML qui consomment des programmes StableHLO.

Notre objectif est de simplifier et d'accélérer le développement du ML en créant une meilleure interopérabilité entre les différents frameworks de ML (tels que TensorFlow, JAX et PyTorch) et les compilateurs de ML (tels que XLA et IREE). À cette fin, ce document fournit une spécification pour le langage de programmation StableHLO.

Cette spécification comporte trois sections principales. Tout d'abord, la section Programmes décrit la structure des programmes StableHLO, qui se composent de fonctions StableHLO, elles-mêmes composées d'opérations StableHLO. Dans cette structure, la section Ops spécifie la sémantique des opérations individuelles. La section Exécution fournit la sémantique pour toutes ces opérations exécutées ensemble dans un programme. Enfin, la section Notation aborde la notation utilisée tout au long de la spécification.

Pour afficher les spécifications d'une version précédente de StableHLO, ouvrez le dépôt au niveau de la version taguée qui vous intéresse. Par exemple, la spécification StableHLO v0.19.0. Pour afficher les modifications apportées à chaque mise à jour mineure de StableHLO, consultez le journal des versions dans VhloDialect.td.

Programmes

Program ::= {Func}

Les programmes StableHLO se composent d'un nombre arbitraire de fonctions StableHLO. Vous trouverez ci-dessous un exemple de programme avec une fonction @main qui comporte trois entrées (%image, %weights et %bias) et une sortie. Le corps de la fonction comporte six opérations.

func.func @main(
  %image: tensor<28x28xf32>,
  %weights: tensor<784x10xf32>,
  %bias: tensor<1x10xf32>
) -> tensor<1x10xf32> {
  %0 = "stablehlo.reshape"(%image) : (tensor<28x28xf32>) -> tensor<1x784xf32>
  %1 = "stablehlo.dot"(%0, %weights) : (tensor<1x784xf32>, tensor<784x10xf32>) -> tensor<1x10xf32>
  %2 = "stablehlo.add"(%1, %bias) : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32>
  %3 = "stablehlo.constant"() {value = dense<0.0> : tensor<1x10xf32>} : () -> tensor<1x10xf32>
  %4 = "stablehlo.maximum"(%2, %3) : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32>
  "func.return"(%4): (tensor<1x10xf32>) -> ()
}

Fonctions

Func        ::= 'func' '.' 'func' FuncId FuncInputs FuncOutputs '{' FuncBody '}'
FuncInputs  ::= '(' [FuncInput {',' FuncInput}] `)`
FuncInput   ::= ValueId ':' ValueType
FuncOutputs ::= ['->' FuncOutput, {',' FuncOutput}]
FuncOutput  ::= ValueType
FuncBody    ::= {Op}

Les fonctions StableHLO (également appelées fonctions nommées) comportent un identifiant, des entrées/sorties et un corps. À l'avenir, nous prévoyons d'introduire des métadonnées supplémentaires pour les fonctions afin d'améliorer la compatibilité avec HLO (#425, #626, #740, #744).

Identifiants

FuncId  ::= '@' letter {letter | digit}
ValueId ::= '%' digit {digit}
          | '%' letter {letter | digit}
letter  ::= 'a' | ... | 'z' | 'A' | ... | 'Z' | '_'
digit   ::= '0' | ... | '9'

Les identifiants StableHLO sont semblables aux identifiants de nombreux langages de programmation, avec deux particularités : 1) tous les identifiants ont des sigles qui distinguent les différents types d'identifiants, 2) les identifiants de valeur peuvent être entièrement numériques pour simplifier la génération de programmes StableHLO.

Types

Type         ::= ValueType | NonValueType
ValueType    ::= TensorType | QuantizedTensorType | TokenType | TupleType | BufferType
NonValueType ::= TensorElementType | QuantizedTensorElementType | FunctionType | StringType

Les types StableHLO sont classés dans les types de valeurs (également appelés types de première classe) qui représentent les valeurs StableHLO et les types non-valeurs qui décrivent d'autres éléments du programme. Les types StableHLO sont semblables à ceux de nombreux langages de programmation, la principale particularité étant la nature spécifique au domaine de StableHLO, qui entraîne des résultats inhabituels (par exemple, les types scalaires ne sont pas des types de valeurs).

TensorType ::= 'tensor' '<' Shape TensorElementType '>'
Shape ::= {DimensionSize 'x'}
DimensionSize ::= digit {digit} | '?'

Les types Tensor représentent des tenseurs, c'est-à-dire des tableaux multidimensionnels. Ils ont une forme et un type d'élément, où une forme représente des tailles de dimensions non négatives ou inconnues dans l'ordre croissant des dimensions correspondantes (également appelées axes) numérotées de 0 à R-1. Le nombre de dimensions R est appelé rang. Par exemple, tensor<2x3xf32> est un type de Tensor avec une forme 2x3 et un type d'élément f32. Il comporte deux dimensions (ou, en d'autres termes, deux axes) : la dimension 0 et la dimension 1, dont les tailles sont respectivement de 2 et 3. Son rang est 2.

Les formes peuvent être partiellement ou complètement inconnues (dynamiques), par exemple tensor<?x2xf64> est partiellement inconnue et tensor<?x?xf64> est complètement inconnue. Les tailles de dimensions dynamiques sont représentées à l'aide d'un ?. Il n'est pas possible de supprimer le classement des formes.

À l'avenir, nous prévoyons d'étendre les types de Tensor au-delà des tailles de dimension et des types d'éléments, par exemple pour inclure les mises en page (#629) et la parcimonie (#1078).

QuantizedTensorType ::= 'tensor' '<' Shape QuantizedTensorElementType '>'
QuantizedTensorElementType ::= '!quant.uniform' '<'
                  QuantizationStorageType
                  ['<' QuantizationStorageMin ':' QuantizationStorageMax '>']
                  ':' QuantizationExpressedType
                  [':' QuantizationDimension]
                  ',' QuantizationParameters '>'
QuantizationStorageType ::= IntegerType
QuantizationStorageMin ::= IntegerLiteral
QuantizationStorageMax ::= IntegerLiteral
QuantizationExpressedType ::= FloatType
QuantizationDimension ::= IntegerLiteral
QuantizationParameters ::= QuantizationParameter
                         | '{' QuantizationParameter {',' QuantizationParameter} '}'
QuantizationParameter ::= QuantizationScale [':' QuantizationZeroPoint]
QuantizationScale ::= FloatLiteral
QuantizationZeroPoint ::= IntegerLiteral

Les types d'éléments quantifiés représentent des valeurs entières d'un type de stockage dans la plage allant de storage_min à storage_max (inclus) qui correspondent à des valeurs à virgule flottante d'un type exprimé. Pour une valeur entière donnée i, la valeur à virgule flottante correspondante f peut être calculée sous la forme f = (i - zero_point) * scale, où scale et zero_point sont appelés paramètres de quantification. Les éléments storage_min et storage_max sont facultatifs dans la grammaire, mais ont respectivement les valeurs par défaut min_value(storage_type) et max_value(storage_type). Les types d'éléments quantifiés sont soumis aux contraintes suivantes :

• (C1) type(storage_min) = storage_type.
• (C2) type(storage_max) = storage_type.
• (C3) min_value(storage_type) <= storage_min < storage_max <= max_value(storage_type).
• (C4) type(scales...) = expressed_type.
• (C5) 0 < scales.
• (C6) is_finite(scales...).
• (C7) storage_min <= zero_points <= storage_max.
• (C8) type(zero_points...) = storage_type.
• (C9) size(scales) = size(zero_points).
• (C10) Si is_empty(quantization_dimension), alors size(scales) = 1.
• (C11) 0 <= quantization_dimension.

Pour le moment, QuantizationScale est une constante à virgule flottante, mais il existe un fort intérêt pour les échelles basées sur des nombres entiers, représentées par des multiplicateurs et des décalages. Nous prévoyons d'étudier cette possibilité dans un avenir proche (#1404).

Une discussion est en cours sur la sémantique de QuantizationZeroPoint, y compris le type, les valeurs et la possibilité d'avoir un ou plusieurs points zéro dans un type de Tensor quantifié. En fonction des résultats de cette discussion, la spécification concernant les zéro points pourra être modifiée à l'avenir (#1405).

Une autre discussion en cours porte sur la sémantique de QuantizationStorageMin et QuantizationStorageMax afin de déterminer si des contraintes doivent être imposées à ces valeurs et à celles des Tensors quantifiés (#1406).

Enfin, nous prévoyons d'explorer la représentation des échelles inconnues et des points zéro, de la même manière que nous prévoyons d'explorer la représentation des tailles de dimensions inconnues (#1407).

Les types de tenseurs quantifiés représentent des tenseurs avec des éléments quantifiés. Ces Tensors sont exactement les mêmes que les Tensors standards, sauf que leurs éléments ont des types d'éléments quantifiés au lieu de types d'éléments standards.

Dans les Tensors quantifiés, la quantification peut être par Tensor, c'est-à-dire avec un scale et un zero_point pour l'ensemble du Tensor, ou par axe, c'est-à-dire avec plusieurs scales et zero_points, une paire par tranche d'une dimension particulière quantization_dimension. Plus formellement, dans un Tensor t avec quantification par axe, il existe des tranches dim(t, quantization_dimension) de quantization_dimension : t[:, ..., 0, ..., :], t[:, ..., 1, ..., :], etc. Tous les éléments de la tranche i utilisent scales[i] et zero_points[i] comme paramètres de quantification. Les types de Tensor quantifiés sont soumis aux contraintes suivantes :

• Pour la quantification par Tensor :
  • Aucune autre contrainte.
• Pour la quantification par axe :
  • (C12) quantization_dimension < rank(self).
  • (C13) dim(self, quantization_dimension) = size(scales).

TokenType ::= 'token'

Les types de jetons représentent des jetons, c'est-à-dire des valeurs opaques produites et consommées par certaines opérations. Les jetons sont utilisés pour imposer un ordre d'exécution aux opérations, comme décrit dans la section Exécution.

TupleType ::= 'tuple' '<' TupleElementTypes '>'
TupleElementTypes ::= [ValueType {',' ValueType}]

Les types de tampon représentent des tampons. Par exemple, dans XLA, les tampons sont des tableaux multidimensionnels avec un stockage cohérent. Comme les types Tensor, les types Buffer ont une forme et un type d'élément, où une forme représente des tailles de dimension non négatives ou inconnues dans l'ordre croissant des dimensions correspondantes (également appelées axes) numérotées de 0 à R-1. Le nombre de dimensions R est appelé rang. Par exemple, memref<2x3xf32> est un type de tampon avec une forme 2x3 et un type d'élément f32. Il comporte deux dimensions (ou, en d'autres termes, deux axes) : la dimension 0 et la dimension 1, dont les tailles sont respectivement de 2 et 3. Son rang est 2.

Les tampons peuvent être alloués à l'aide d'un custom_call à CreateBuffer ou Pin, et désalloués à l'aide d'un custom_call à Unpin. Seules les opérations custom_call peuvent lire et écrire le contenu des tampons. Pour en savoir plus, consultez custom_call.

Les types de tuples représentent des tuples, c'est-à-dire des listes hétérogènes. Les tuples sont une fonctionnalité ancienne qui n'existe que pour la compatibilité avec HLO. Dans HLO, les tuples sont utilisés pour représenter les entrées et sorties variadiques. Dans StableHLO, les entrées et sorties variadiques sont prises en charge de manière native. La seule utilisation des tuples dans StableHLO est de représenter de manière exhaustive l'ABI HLO où, par exemple, T, tuple<T> et tuple<tuple<T>> peuvent être très différents selon une implémentation particulière. À l'avenir, nous prévoyons de modifier l'ABI HLO, ce qui pourrait nous permettre de supprimer les types de tuples de StableHLO (#598).

TensorElementType ::= BooleanType | IntegerType | FloatType | ComplexType
BooleanType ::= 'i1'
IntegerType ::= SignedIntegerType | UnsignedIntegerType
SignedIntegerType ::= 'si2' | 'si4' | 'si8' | 'si16' | 'si32' | 'si64'
UnsignedIntegerType ::= 'ui2' | 'ui4' | 'ui8' | 'ui16' | 'ui32' | 'ui64'
FloatType ::= 'f4E2M1FN' | 'f6E2M3FN' | 'f6E3M2FN' | 'f8E3M4' | 'f8E4M3'
            | 'f8E4M3FN' | 'f8E4M3FNUZ' | 'f8E4M3B11FNUZ' | 'f8E5M2'
            | 'f8E5M2FNUZ' | 'f8E8M0FNU' | 'bf16' | 'f16' | 'f32' | 'f64'
TensorFloat32 ::= 'tf32'
ComplexType ::= 'complex' '<' ComplexElementType '>'
ComplexElementType ::= 'f32' | 'f64'

Les types d'éléments représentent les éléments des types Tensor. Contrairement à de nombreux langages de programmation, ces types ne sont pas de première classe dans StableHLO. Cela signifie que les programmes StableHLO ne peuvent pas représenter directement les valeurs de ces types (par conséquent, il est idiomatique de représenter les valeurs scalaires de type T avec des valeurs de tenseur de dimension 0 de type tensor<T>).

• Le type booléen représente les valeurs booléennes true et false.
• Les types entiers peuvent être signés (si) ou non signés (ui) et avoir l'une des largeurs de bits acceptées (2, 4, 8, 16, 32 ou 64). Les types siN signés représentent des valeurs entières allant de -2^(N-1) à 2^(N-1)-1 inclus, et les types uiN non signés représentent des valeurs entières allant de 0 à 2^N-1 inclus.
• Les types à virgule flottante peuvent être l'un des suivants :
  • f8E3M4, f8E4M3 et f8E5M2 sont des nombres à virgule flottante de 8 bits suivant les conventions IEEE-754.
  • Types f8E4M3FN et f8E5M2 correspondant respectivement aux encodages E4M3 et E5M2 du format FP8 décrits dans Formats FP8 pour le deep learning.
  • Types f8E4M3FNUZ et f8E5M2FNUZ correspondant aux encodages E4M3 et E5M2 des formats FP8 décrits dans Formats numériques 8 bits pour les réseaux de neurones profonds.
  • Type f8E4M3B11FNUZ correspondant à l'encodage E4M3 des formats FP8 décrits dans Hybrid 8-bit Floating Point (HFP8) Training and Inference for Deep Neural Networks.
  • Type bf16 correspondant au format bfloat16 décrit dans BFloat16 : le secret des hautes performances sur les Cloud TPU.
  • Les types f16, f32 et f64 correspondent respectivement aux formats binary16 ("demi-précision"), binary32 ("simple précision") et binary64 ("double précision") décrits dans la norme IEEE 754.
  • Le type tf32 correspond au format TensorFloat32 et est limité dans StableHLO.
  • Types MX (microsizing) f4E2M1FN, f6E2M3FN, f6E3M2FN et f8E8M0FNU décrits dans la spécification des formats de microsizing OCP.
• Les types complexes représentent des valeurs complexes qui ont une partie réelle et une partie imaginaire du même type d'élément. Les types complexes acceptés sont complex<f32> (les deux parties sont de type f32) et complex<f64> (les deux parties sont de type f64).

FunctionType ::= '(' InputTypes ')' '->' '(' OutputTypes ')'
InputTypes ::= [ValueType {',' ValueType}]
OutputTypes ::= [ValueType {',' ValueType}]

Les types de fonction représentent les fonctions nommées et anonymes. Ils ont des types d'entrée (la liste des types à gauche de ->) et des types de sortie (la liste des types à droite de ->). Dans de nombreux langages de programmation, les types de fonctions sont de première classe, mais pas dans StableHLO.

StringType ::= 'string'

Le type String représente des séquences d'octets. Contrairement à de nombreux langages de programmation, le type de chaîne n'est pas de première classe dans StableHLO et n'est utilisé que pour spécifier des métadonnées statiques pour les éléments de programme.

Opérations

Les opérations StableHLO (également appelées ops) représentent un ensemble fermé d'opérations de haut niveau dans les modèles de machine learning. Comme indiqué ci-dessus, la syntaxe StableHLO s'inspire fortement de MLIR, qui n'est pas forcément l'alternative la plus ergonomique, mais qui est sans doute la mieux adaptée à l'objectif de StableHLO, qui est de créer une meilleure interopérabilité entre les frameworks de ML et les compilateurs de ML.

Op            ::= [OpOutputs] OpName OpInputs ':' OpSignature
OpName        ::= '"' 'stablehlo' '.' OpMnemonic '"'
OpMnemonic    ::= 'abs' | 'add' | ...

Les opérations StableHLO (également appelées ops) ont un nom, des entrées/sorties et une signature. Le nom se compose du préfixe stablehlo. et d'un mnémonique qui identifie de manière unique l'une des opérations compatibles. Vous trouverez ci-dessous la liste complète des opérations compatibles.

OpInputs        ::= OpInputValues OpInputFuncs OpInputAttrs
OpInputValues   ::= '(' [OpInputValue {',' OpInputValue}] ')'
OpInputValue    ::= ValueId
OpInputFuncs    ::= ['(' OpInputFunc {',' OpInputFunc} ')']
OpInputAttrs    ::= ['{' OpInputAttr {',' OpInputAttr} '}']
OpOutputs       ::= [OpOutput {',' OpOutput} '=']
OpOutput        ::= ValueId

Les opérations consomment des entrées et produisent des sorties. Les entrées sont classées dans les catégories suivantes : valeurs d'entrée (calculées lors de l'exécution), fonctions d'entrée (fournies de manière statique, car dans StableHLO, les fonctions ne sont pas des valeurs de première classe) et attributs d'entrée (également fournis de manière statique). Le type d'entrées et de sorties consommées et produites par une opération dépend de son mnémonique. Par exemple, l'opération add consomme deux valeurs d'entrée et produit une valeur de sortie. En comparaison, l'opération select_and_scatter consomme trois valeurs d'entrée, deux fonctions d'entrée et trois attributs d'entrée.

OpInputFunc ::= '{' Unused FuncInputs ':' FuncBody '}'
Unused      ::= '^' digit {digit}
              | '^' letter {letter | digit}

Les fonctions d'entrée (également appelées fonctions anonymes) sont très semblables aux fonctions nommées, sauf qu'elles n'ont pas d'identifiant (d'où le nom "anonyme") et qu'elles ne déclarent pas de types de sortie (les types de sortie sont déduits de l'opération return dans la fonction).

La syntaxe des fonctions d'entrée inclut une partie actuellement inutilisée (voir la production Unused ci-dessus), qui est là pour la compatibilité avec MLIR. Dans MLIR, il existe un concept plus général de "régions" qui peuvent comporter plusieurs "blocs" d'opérations connectés entre eux par des opérations de saut. Ces blocs ont des ID qui correspondent à la production Unused, ce qui permet de les distinguer les uns des autres. StableHLO ne comporte pas d'opérations de saut. La partie correspondante de la syntaxe MLIR n'est donc pas utilisée (mais elle est toujours présente).

OpInputAttr      ::= OpInputAttrName '=' OpInputAttrValue
OpInputAttrName  ::= letter {letter | digit}
OpInputAttrValue ::= Constant

Les attributs d'entrée ont un nom et une valeur qui correspond à l'une des constantes acceptées. Il s'agit du principal moyen de spécifier des métadonnées statiques pour les éléments de programme. Par exemple, l'opération concatenate utilise l'attribut dimension pour spécifier la dimension le long de laquelle ses valeurs d'entrée sont concaténées. De même, l'opération slice utilise plusieurs attributs tels que start_indices et limit_indices pour spécifier les limites utilisées pour segmenter la valeur d'entrée.

Pour le moment, les programmes StableHLO en circulation contiennent parfois des attributs qui ne sont pas décrits dans ce document. À l'avenir, nous prévoyons d'intégrer ces attributs à l'opset StableHLO ou de leur interdire d'apparaître dans les programmes StableHLO. En attendant, voici la liste de ces attributs :

• layout (#629).
• mhlo.frontend_attributes (#628).
• mhlo.sharding (#619).
• output_operand_aliases (#740).
• Métadonnées de localisation (#594).

OpSignature ::= '(' [ValueType {',' ValueType}] ')' '->' '(' [ValueType {',' ValueType}] ')'

La signature d'opération se compose des types de toutes les valeurs d'entrée (la liste des types à gauche de ->) et des types de toutes les valeurs de sortie (la liste des types à droite de ->). À proprement parler, les types d'entrée sont redondants, et les types de sortie le sont presque toujours également (car pour la plupart des opérations StableHLO, les types de sortie peuvent être déduits des entrées). Néanmoins, la signature op fait délibérément partie de la syntaxe StableHLO pour assurer la compatibilité avec MLIR.

Vous trouverez ci-dessous un exemple d'op dont le mnémonique est select_and_scatter. Il consomme trois valeurs d'entrée (%operand, %source et %init_value), deux fonctions d'entrée et trois attributs d'entrée (window_dimensions, window_strides et padding). Notez que la signature de l'opération n'inclut que les types de ses valeurs d'entrée (mais pas les types de fonctions et d'attributs d'entrée qui sont fournis en ligne).

%result = "stablehlo.select_and_scatter"(%operand, %source, %init_value) ({
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = "stablehlo.compare"(%arg0, %arg1) {
      comparison_direction = #stablehlo<comparison_direction GE>
    } : (tensor<i32>, tensor<i32>) -> tensor<i1>
    "stablehlo.return"(%0) : (tensor<i1>) -> ()
}, {
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
    "stablehlo.return"(%0) : (tensor<i32>) -> ()
}) {
  window_dimensions = dense<[3, 1]> : tensor<2xi64>,
  window_strides = dense<[2, 1]> : tensor<2xi64>,
  padding = dense<[[0, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<4x2xi32>, tensor<2x2xi32>, tensor<i32>) -> tensor<4x2xi32>

Constantes

Constant ::= BooleanConstant
           | IntegerConstant
           | FloatConstant
           | ComplexConstant
           | TensorConstant
           | QuantizedTensorConstant
           | StringConstant
           | EnumConstant

Les constantes StableHLO ont un littéral et un type qui représentent ensemble une valeur StableHLO. En général, le type fait partie de la syntaxe constante, sauf lorsqu'il n'y a pas d'ambiguïté (par exemple, une constante booléenne a sans ambiguïté le type i1, tandis qu'une constante entière peut avoir plusieurs types possibles).

BooleanConstant ::= BooleanLiteral
BooleanLiteral  ::= 'true' | 'false'

Les constantes booléennes représentent les valeurs booléennes true et false. Les constantes booléennes sont de type i1.

IntegerConstant   ::= IntegerLiteral ':' IntegerType
IntegerLiteral    ::= ['-' | '+'] DecimalDigits
                    | ['-' | '+'] '0x' HexadecimalDigits
DecimalDigits     ::= decimalDigit {decimalDigit}
HexadecimalDigits ::= hexadecimalDigit {hexadecimalDigit}
decimalDigit      ::= '0' | ... | '9'
hexadecimalDigit  ::= decimalDigit | 'a' | ... | 'f' | 'A' | ... | 'F'

Les constantes entières représentent des valeurs entières sous forme de chaînes utilisant la notation décimale ou hexadécimale. Les autres bases, par exemple binaire ou octale, ne sont pas acceptées. Les constantes entières sont soumises aux contraintes suivantes :

• (C1) is_wellformed(integer_literal, integer_type).

FloatConstant  ::= FloatLiteral ':' FloatType
FloatLiteral   ::= SignPart IntegerPart FractionalPart ScientificPart
                 | '0x' [HexadecimalDigits]
SignPart       ::= ['-' | '+']
IntegerPart    ::= DecimalDigits
FractionalPart ::= ['.' [DecimalDigits]]
ScientificPart ::= [('e' | 'E') ['-' | '+'] DecimalDigits]

Les constantes à virgule flottante représentent des valeurs à virgule flottante via des chaînes qui utilisent la notation décimale ou scientifique. De plus, la notation hexadécimale peut être utilisée pour spécifier directement les bits sous-jacents dans le format à virgule flottante du type correspondant. Les constantes à virgule flottante sont soumises aux contraintes suivantes :

• (C1) Si une notation non hexadécimale est utilisée, is_wellformed(float_literal, float_type).
• (C2) Si la notation hexadécimale est utilisée, size(hexadecimal_digits) = num_bits(float_type) / 4.

ComplexConstant ::= ComplexLiteral ':' ComplexType
ComplexLiteral  ::= '(' RealPart ',' ImaginaryPart ')'
RealPart        ::= FloatLiteral
ImaginaryPart   ::= FloatLiteral

Les constantes complexes représentent des valeurs complexes à l'aide de listes composées d'une partie réelle (en premier) et d'une partie imaginaire (en second). Par exemple, (1.0, 0.0) : complex<f32> représente 1.0 + 0.0i et (0.0, 1.0) : complex<f32> représente 0.0 + 1.0i. L'ordre dans lequel ces parties sont ensuite stockées en mémoire est défini par l'implémentation. Les constantes complexes sont soumises aux contraintes suivantes :

• (C1) is_wellformed(real_part, complex_element_type(complex_type)).
• (C2) is_wellformed(imaginary_part, complex_element_type(complex_type)).

TensorConstant ::= TensorLiteral ':' TensorType
TensorLiteral  ::= 'dense' '<' (DenseLiteral | ElementLiteral) '>'
DenseLiteral   ::= DenseDimension | DenseElements
DenseDimension ::= '[' [DenseLiteral {',' DenseLiteral}] ']'
DenseElements  ::= [ElementLiteral {',' ElementLiteral}]
ElementLiteral ::= BooleanLiteral | IntegerLiteral | FloatLiteral | ComplexLiteral

Les constantes Tensor représentent les valeurs Tensor à l'aide de listes imbriquées spécifiées à l'aide de la notation NumPy. Par exemple, dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> représente une valeur de Tensor avec le mappage suivant des index aux éléments : {0, 0} => 1, {0, 1} => 2, {0, 2} => 3, {1, 0} => 4, {1, 1} => 5, {1, 2} => 6. L'ordre dans lequel ces éléments sont ensuite stockés en mémoire est défini par l'implémentation. Les constantes Tensor sont soumises aux contraintes suivantes :

• (C1) has_syntax(tensor_literal, element_type(tensor_type)), où :
  • has_syntax(element_literal: Syntax, element_type: Type) = is_wellformed(element_literal, type).
  • has_syntax(tensor_literal: List, element_type: Type) = has_syntax(tensor_literal..., element_type).
• (C2) has_shape(tensor_literal, shape(tensor_type)), où :
  • has_shape(element_literal: Syntax, []) = true.
  • has_shape(tensor_literal: List, shape: List) = size(tensor_literal) = shape[0] and has_shape(tensor_literal..., shape[1:]).
  • Sinon, false.

QuantizedTensorConstant ::= QuantizedTensorLiteral ':' QuantizedTensorType
QuantizedTensorLiteral  ::= 'dense' '<' (DenseLiteral | ElementLiteral) '>'

Les constantes de tenseur quantifié représentent des valeurs de tenseur quantifié en utilisant la même notation que les constantes de tenseur, avec des éléments spécifiés comme constantes de leur type de stockage. Les constantes de Tensor quantifié sont soumises aux contraintes suivantes :

• (C1) has_syntax(quantized_tensor_literal, storage_type(quantized_tensor_type)).
• (C2) has_shape(quantized_tensor_literal, shape(quantized_tensor_type)).

StringConstant  ::= StringLiteral
StringLiteral   ::= '"' {stringCharacter | escapeSequence} '"'
stringCharacter ::= all ASCII characters except '\00', '\01', ... '\1f' and '"'
escapeSequence  ::= '\' ('"' | '\' | 'n' | 't' | (hexadecimalDigit hexadecimalDigit))

Les littéraux de chaîne sont constitués d'octets spécifiés à l'aide de caractères ASCII et de séquences d'échappement. Elles sont indépendantes de l'encodage. L'interprétation de ces octets est donc définie par l'implémentation. Les littéraux de chaîne sont de type string.

Opérations

abs

Sémantique

Effectue une opération abs au niveau des éléments sur le Tensor operand et génère un Tensor result. En fonction du type d'élément, effectue les actions suivantes :

• Pour les entiers signés : module entier.
• Pour les valeurs float : abs à partir de la norme IEEE-754.
• Pour les nombres complexes : module complexe.
• Pour les types quantifiés : dequantize_op_quantize(abs, operand, type(result)).

Entrées

Sorties

Contraintes

• (C1) shape(result) = shape(operand).
• (C2) baseline_element_type(result) est défini comme suit :
  • complex_element_type(element_type(operand)) si is_complex(operand).
  • Sinon, baseline_element_type(operand).

Exemples

// %operand: [-2, 0, 2]
%result = "stablehlo.abs"(%operand)< : (t>ens>or3xi32<) - t>ensor3xi32
// %result: [2, 0, 2]

 Autres exemples

add

Sémantique

Effectue l'addition élément par élément de deux Tensors lhs et rhs, et produit un Tensor result. En fonction du type d'élément, effectue les actions suivantes :

• Pour les valeurs booléennes : OR logique.
• Pour les nombres entiers : addition d'entiers.
• Pour les valeurs float : addition à partir de la norme IEEE-754.
• Pour les nombres complexes : addition complexe.
• Pour les types quantifiés : dequantize_op_quantize(add, lhs, rhs, type(result)).

Entrées

Sorties

Contraintes

• Si l'opération utilise des Tensors non quantifiés :
  • (C1) type(lhs) = type(rhs) = type(result).
• Si l'opération utilise des Tensors quantifiés :
  • (C2) is_quantized(lhs) and is_quantized(rhs) and is_quantized(result).
  • (C3) storage_type(lhs) = storage_type(rhs) = storage_type(result).
  • (C4) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C5) (is_per_axis_quantized(lhs) or is_per_axis_quantized(rhs)) = is_per_axis_quantized(result).
  • (C6) Si is_per_axis_quantized(lhs), alors quantization_dimension(lhs) = quantization_dimension(result).
  • (C7) Si is_per_axis_quantized(rhs), alors quantization_dimension(rhs) = quantization_dimension(result).

Exemples

// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.add"(%lhs, %rhs)< : (ten>sor2x2xi<32, ten>sor>2x2xi32<) - ten>sor2x2xi32
// %result: [[6, 8], [10, 12]]

 Autres exemples

after_all

Sémantique

Garantit que les opérations produisant le inputs sont exécutées avant toute opération dépendant de result. L'exécution de cette opération n'a aucun effet. Elle n'existe que pour établir des dépendances de données entre result et inputs.

Entrées

Sorties

Exemples

// %input0: !stablehlo.token
// %input1: !stablehlo.token
%result = "stablehlo.after_all"(%input0, %input1) : (!stablehlo.token, !stablehl>o.token) - !stablehlo.token

 Autres exemples

all_gather

Sémantique

Dans chaque groupe de processus de la grille de processus StableHLO, concatène les valeurs des Tensors operands de chaque processus le long de all_gather_dim et produit des Tensors results.

L'opération divise la grille de processus StableHLO en process_groups, qui est défini comme suit :

• cross_replica(replica_groups) si channel_id <= 0 and use_global_device_ids = false.
• cross_replica_and_partition(replica_groups) si channel_id > 0 and use_global_device_ids = false.
• flattened_ids(replica_groups) si channel_id > 0 and use_global_device_ids = true.

Ensuite, dans chaque process_group :

• operands...@receiver = [operand@sender for sender in process_group] pour tous les receiver dans process_group.
• results...@process = concatenate(operands...@process, all_gather_dim) pour tous les process dans process_group.

Entrées

Sorties

Contraintes

• (C1) 0 <= all_gather_dim < rank(operands...).
• (C2) is_unique(replica_groups).
• (C3) size(replica_groups) est défini comme suit :
  • num_replicas si cross_replica est utilisé.
  • num_replicas si cross_replica_and_partition est utilisé.
  • num_processes si flattened_ids est utilisé.
• (C4) 0 <= replica_groups < size(replica_groups).
• (C5) Si use_global_device_ids = true, alors channel_id > 0.
• (C6) type(results...) = type(operands...) sauf :
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1).

Exemples

// num_replicas: 2
// num_partitions: 1
// %operand0@(0, 0): [[1, 2], [3, 4]]
// %operand0@(1, 0): [[5, 6], [7, 8]]
// %operand1@(0, 0): [[11, 12], [13, 14]]
// %operand1@(1, 0): [[15, 16], [17, 18]]
%result:2 = "stablehlo.all_gather"(%operand0, %operand1) {
  all_gather_dim = 1 : i64,
  replica_grou<ps = den>se[[0, 1]<] : ten>sor1x2xi64,
  // channel_id = 0
  channel_handle = #stablehlo.chan<nel_handlehandle = 0>, type = 0
  // use_global_device_ids = false
}< : (ten>sor2x2xi<64, ten>sor>2x2xi64)< - (ten>sor2x4xi<64, ten>sor2x4xi64)
// %result0@(0, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result0@(1, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result1@(0, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]
// %result1@(1, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]

 Autres exemples

all_reduce

Sémantique

Dans chaque groupe de processus de la grille de processus StableHLO, applique une fonction de réduction computation aux valeurs des Tensors operands de chaque processus et produit des Tensors results.

L'opération divise la grille de processus StableHLO en process_groups, qui est défini comme suit :

• cross_replica(replica_groups) si channel_id <= 0 and use_global_device_ids = false.
• cross_replica_and_partition(replica_groups) si channel_id > 0 and use_global_device_ids = false.
• flattened_ids(replica_groups) si channel_id > 0 and use_global_device_ids = true.

Ensuite, dans chaque process_group :

• results...@process[result_index] = exec(schedule) pour un arbre binaire donné schedule où :
  • exec(node) = computation(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value.
• schedule est un arbre binaire défini par l'implémentation dont la traversée dans l'ordre est to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0])).

Entrées

Sorties

Contraintes

• (C1) is_unique(replica_groups).
• (C2) size(replica_groups) est défini comme suit :
  • num_replicas si cross_replica est utilisé.
  • num_replicas si cross_replica_and_partition est utilisé.
  • num_processes si flattened_ids est utilisé.
• (C3) 0 <= replica_groups < size(replica_groups).
• (C4) Si use_global_device_ids = true, alors channel_id > 0.
• (C5) computation est de type (tensor<E>, tensor<E>) -> (tensor<E>) où is_promotable(element_type(operand), E).
• (C6) shape(results...) = shape(operands...).
• (C7) element_type(results...) = E.

Exemples

// num_replicas: 2
// num_partitions: 1
// %operand0@(0, 0): [1, 2, 3, 4]
// %operand0@(1, 0): [5, 6, 7, 8]
// %operand1@(0, 0): [9, 10, 11, 12]
// %operand1@(1, 0): [13, 14, 15, 16]
%result:2 = "stablehlo.all_reduce"(%operand0, %operand0) ({
  ^bb0(%ar<g0:> tensori64, %ar<g1:> tensori64):
    %0 = "stablehlo.add"(%ar<g0,> %arg1) <: (>ten>sori64,< te>nsori64) - tensori64
    "stable<hlo>.re>turn"(%0) : (tensori64) - ()<
}) {
  >replica_g<roups => dense[[0, 1]] : tensor1x2xi64,
  // channel_id = 0
  channel_hand<le = #stablehlo.chan>nel_handlehandle = 0, type = 0
  // use_global_<devic>e_ids = <false>
} >: (tenso<r4xi6>4, tenso<r4xi6>4) - (tensor4xi64, tensor4xi64)
// %result0@(0, 0): [6, 8, 10, 12]
// %result0@(1, 0): [6, 8, 10, 12]
// %result1@(0, 0): [22, 24, 26, 28]
// %result1@(1, 0): [22, 24, 26, 28]

 Autres exemples

all_to_all

Sémantique

Dans chaque groupe de processus de la grille de processus StableHLO, divise les valeurs des tenseurs operands le long de split_dimension en parties, répartit les parties divisées entre les processus, concatène les parties réparties le long de concat_dimension et produit des tenseurs results. L'opération divise la grille de processus StableHLO en process_groups, qui est défini comme suit :

• cross_replica(replica_groups) si channel_id <= 0.
• cross_partition(replica_groups) si channel_id > 0.

Ensuite, dans chaque process_group :

• split_parts...@sender = split(operands...@sender, split_count, split_dimension) pour tous les sender dans process_group.
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group] oùreceiver_index = process_group.index(receiver).
• results...@process = concatenate(scattered_parts...@process, concat_dimension).

Entrées

Sorties

Contraintes

• (C1) 0 <= split_dimension < rank(operands...).
• (C2) dim(operands..., split_dimension) % split_count = 0.
• (C3) 0 <= concat_dimension < rank(operands...).
• (C4) 0 < split_count.
• (C5) is_unique(replica_groups).
• (C6) size(replica_groups) est défini comme suit :
  • num_replicas si cross_replica est utilisé.
  • num_partitions si cross_partition est utilisé.
• (C7) 0 <= replica_groups < size(replica_groups).
• (C8) dim(replica_groups, 1) = split_count.
• (C9) type(results...) = type(operands...) sauf si split_dimension != concat_dimension :
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count.
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count.

Exemples

// num_replicas: 2
// num_partitions: 1
// %operand1@(0, 0): [[1, 2, 3, 4],
//                    [5, 6, 7, 8]]
// %operand1@(1, 0): [[9, 10, 11, 12],
//                    [13, 14, 15, 16]]
// %operand2@(0, 0): [[17, 18, 19, 20],
//                    [21, 22, 23, 24]]
// %operand2@(1, 0): [[25, 26, 27, 28],
//                    [29, 30, 31, 32]]
%result:2 = "stablehlo.all_to_all"(%operand1, %operand2) {
  split_dimension = 1 : i64,
  concat_dimension = 0 : i64,
  split_count = 2 : i64,
  replica_grou<ps = den>se[[0, 1]<] : ten>sor1x2xi64
  // channel_id = 0
}< : (ten>sor2x4xi<64, ten>sor>2x4xi64)< - (ten>sor4x2xi<64, ten>sor4x2xi64)
// %result#0@(0, 0): [[1, 2], [5, 6], [9, 10], [13, 14]]
// %result#0@(1, 0): [[3, 4], [7, 8], [11, 12], [15, 16]]
// %result#1@(0, 0): [[17, 18], [21, 22], [25, 26], [29, 30]]
// %result#1@(1, 0): [[19, 20], [23, 24], [27, 28], [31, 32]]

 Autres exemples

et

Sémantique

Effectue un AND élément par élément de deux Tensors lhs et rhs, et produit un Tensor result. En fonction du type d'élément, effectue les actions suivantes :

• Pour les valeurs booléennes : AND logique.
• Pour les nombres entiers : AND au niveau du bit.

Entrées

Sorties

Contraintes

• (C1) type(lhs) = type(rhs) = type(result).

Exemples

// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.and"(%lhs, %rhs)< : (ten>sor2x2xi<32, ten>sor>2x2xi32<) - ten>sor2x2xi32
// %result: [[1, 2], [3, 0]]

 Autres exemples

atan2

Sémantique

Effectue une opération atan2 au niveau des éléments sur les Tensors lhs et rhs, et génère un Tensor result. En fonction du type d'élément, effectue les actions suivantes :

• Pour les valeurs float : atan2 à partir de la norme IEEE-754.
• Pour les nombres complexes : atan2 complexe.
• Pour les types quantifiés : dequantize_op_quantize(atan2, lhs, rhs, type(result)).

Entrées

Sorties

Contraintes

• (C1) baseline_type(lhs) = baseline_type(rhs) = baseline_type(result).

Exemples

// %lhs: [0.0, 1.0, -1.0]
// %rhs: [0.0, 0.0, 0.0]
%result = "stablehlo.atan2"(%lhs, %rhs)< : (t>ensor3xf<64, t>ens>or3xf64<) - t>ensor3xf64
// %result: [0.0, 1.57079637, -1.57079637] // [0.0, pi/2, -pi/2]

 Autres exemples

batch_norm_grad

Sémantique

Calcule les gradients de plusieurs entrées de batch_norm_training en rétropropagation à partir de grad_output et produit les Tensors grad_operand, grad_scale et grad_offset. Plus formellement, cette opération peut être exprimée comme une décomposition en opérations StableHLO existantes à l'aide de la syntaxe Python comme suit :

def compute_sum(operand, feature_index):
  (sum,) = reduce(
      inputs=[operand],
      init_values=[constant(0, element_type(operand))],
      dimensions=[i for i in range(rank(operand)) if i != feature_index],
      body=lambda x, y: add(x, y))
  return sum

def compute_mean(operand, feature_index):
  sum = compute_sum(operand, feature_index)
  divisor = constant(size(operand) / dim(operand, feature_index),
                     element_type(operand))
  divisor_bcast = broadcast_in_dim(divisor, [], type(sum))
  return divide(sum, divisor_bcast)

def batch_norm_grad(operand, scale, mean, variance, grad_output, epsilon, feature_index):
  # Broadcast inputs to type(operand)
  scale_bcast = broadcast_in_dim(scale, [feature_index], type(operand))
  mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
  variance_bcast = broadcast_in_dim(variance, [feature_index], type(operand))
  epsilon_bcast = broadcast_in_dim(constant(epsilon, element_type(operand)), [],
                                   type(operand))

  # Perform normalization using the provided `mean` and `variance`
  # Intermediate values will be useful for computing gradients
  centered_operand = subtract(operand, mean_bcast)
  stddev = sqrt(add(variance_bcast, epsilon_bcast))
  normalized_operand = divide(centered_operand, stddev)

  # Use the implementation from batchnorm_expander.cc in XLA
  # Temporary variables have exactly the same names as in the C++ code
  elements_per_feature = broadcast_in_dim(
      constant(divide(size(operand), dim(operand, feature_index)),
               element_type(grad_output)),
      [], type(operand))
  i1 = multiply(grad_output, elements_per_feature)
  i2 = broadcast_in_dim(
      compute_sum(grad_output, feature_index), [feature_index], type(operand))
  i3 = broadcast_in_dim(
      compute_sum(multiply(grad_output, centered_operand), feature_index),
      [feature_index], type(operand))
  i4 = multiply(i3, centered_operand)
  i5 = divide(i4, add(variance_bcast, epsilon_bcast))
  i6 = subtract(subtract(i1, i2), i5)

  grad_operand =
      multiply(divide(divide(scale_bcast, stddev), elements_per_feature), i6)
  grad_scale =
      compute_sum(multiply(grad_output, normalized_operand), feature_index)
  grad_offset = compute_sum(grad_output, feature_index)

  return grad_operand, grad_scale, grad_offset

Pour les types quantifiés, effectue dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, mean, variance, grad_output: batch_norm_grad(operand, scale, mean, variance, grad_output, epsilon, feature_index), operand, scale, mean, variance, grad_output, type(grad_operand), type(grad_scale), type(feature_index)).

Entrées

Sorties

Contraintes

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, mean, variance, grad_output, grad_operand, grad_scale et grad_offset ont le même baseline_element_type.
• (C3) operand, grad_output et grad_operand ont la même forme.
• (C4) scale, mean, variance, grad_scale et grad_offset ont la même forme.
• (C5) size(scale) = dim(operand, feature_index).

Exemples

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %mean: [2.0, 3.0]
// %variance: [1.0, 1.0]
// %grad_output: [
//                [[0.1, 0.1], [0.1, 0.1]],
//                [[0.1, 0.1], [0.1, 0.1]]
//               ]
%grad_operand, %grad_scale, %grad_offset =
"stablehlo.batch_norm_grad"(%operand, %scale, %mean, %variance, %grad_output) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ensor2xf64,
 <    tenso>r2x>2x2xf64)< - (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf64)
// %grad_operand: [
//                 [[0.0, 0.0], [0.0, 0.0]],
//                 [[0.0, 0.0], [0.0, 0.0]]
//                ]
// %grad_scale:  [0.0, 0.0]
// %grad_offset: [0.4, 0.4]

batch_norm_inference

Sémantique

Normalise le tenseur operand dans toutes les dimensions, à l'exception de la dimension feature_index, et produit un tenseur result. Plus formellement, cette opération peut être exprimée comme une décomposition en opérations StableHLO existantes à l'aide de la syntaxe Python comme suit :

def batch_norm_inference(operand, scale, offset, mean, variance, epsilon, feature_index):
  # Broadcast inputs to shape(operand)
  scale_bcast = broadcast_in_dim(scale, [feature_index], type(operand))
  offset_bcast = broadcast_in_dim(offset, [feature_index], type(operand))
  mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
  variance_bcast = broadcast_in_dim(variance, [feature_index], type(operand))
  epsilon_bcast = broadcast_in_dim(constant(epsilon, element_type(operand)), [],
                                   type(operand))

  # Perform normalization using the provided `mean` and `variance` instead of
  # computing them like `batch_norm_training` does.
  centered_operand = subtract(operand, mean_bcast)
  stddev = sqrt(add(variance_bcast, epsilon_bcast))
  normalized_operand = divide(centered_operand, stddev)
  return add(multiply(scale_bcast, normalized_operand), offset_bcast)

Pour les types quantifiés, effectue dequantize_op_quantize(lambda operand, scale, offset, mean, variance: batch_norm_inference(operand, scale, offset, mean, variance, epsilon, feature_index), operand, scale, offset, mean, variance, type(result)).

Entrées

Sorties

Contraintes

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, mean, variance et result ont le même baseline_element_type.
• (C3) size(scale) = dim(operand, feature_index).
• (C4) size(offset) = dim(operand, feature_index).
• (C5) size(mean) = dim(operand, feature_index).
• (C6) size(variance) = dim(operand, feature_index).
• (C7) baseline_type(operand) = baseline_type(result).

Exemples

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %offset: [1.0, 1.0]
// %mean: [2.0, 3.0]
// %variance: [1.0, 1.0]
%result = "stablehlo.batch_norm_inference"(%operand, %scale, %offset, %mean, %variance) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ens>or2xf64<) - tenso>r2x2x2xf64
// %result: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]

batch_norm_training

Sémantique

Calcule la moyenne et la variance pour toutes les dimensions, à l'exception de la dimension feature_index, et normalise le Tensor operand pour produire les Tensors output, batch_mean et batch_var. Plus formellement, cette opération peut être exprimée comme une décomposition des opérations StableHLO existantes à l'aide de la syntaxe Python comme suit :

def compute_mean(operand, feature_index):
  (sum,) = reduce(
      inputs=[operand],
      init_values=[constant(0, element_type(operand))],
      dimensions=[i for i in range(rank(operand)) if i != feature_index],
      body=lambda x, y: add(x, y))
  divisor = constant(size(operand) / dim(operand, feature_index),
                     element_type(operand))
  divisor_bcast = broadcast_in_dim(divisor, [], type(sum))
  return divide(sum, divisor_bcast)

def compute_variance(operand, feature_index):
  mean = compute_mean(operand, feature_index)
  mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
  centered_operand = subtract(operand, mean_bcast)
  return compute_mean(mul(centered_operand, centered_operand), feature_index)

def batch_norm_training(operand, scale, offset, epsilon, feature_index):
  mean = compute_mean(operand, feature_index)
  variance = compute_variance(operand, feature_index)
  return batch_norm_inference(operand, scale, offset, mean, variance, epsilon,
                              feature_index),
         mean, variance

Pour les types quantifiés, effectue dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, offset: batch_norm_training(operand, scale, offset, epsilon, feature_index), operand, scale, offset, type(output), type(batch_mean), type(batch_var)).

Entrées

Sorties

Contraintes

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, batch_mean, batch_var et output ont le même baseline_element_type.
• (C3) size(scale) = dim(operand, feature_index).
• (C4) size(offset) = dim(operand, feature_index).
• (C5) size(batch_mean) = dim(operand, feature_index).
• (C6) size(batch_var) = dim(operand, feature_index).
• (C7) baseline_type(output) = baseline_type(operand).

Exemples

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %offset: [1.0, 1.0]
%output, %batch_mean, %batch_var = "stablehlo.batch_norm_training"(%operand, %scale, %offset) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ens>or2xf64) -
 <   (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf64)
// %output: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]
// %batch_mean: [2.0, 3.0]
// %batch_var: [1.0, 1.0]

bitcast_convert

Sémantique

Effectue une opération bitcast sur le Tensor operand et produit un Tensor result où les bits de l'ensemble du Tensor operand sont réinterprétés à l'aide du type du Tensor result.

Plus formellement, étant donné E = element_type(operand), E' = element_type(result) et R = rank(operand) :

• Si num_bits(E') < num_bits(E), bits(result[i0, ..., iR-1, :]) = bits(operand[i0, ..., iR-1]).
• Si num_bits(E') > num_bits(E), bits(result[i0, ..., iR-2]) = bits(operand[i0, ..., iR-2, :]).
• Si num_bits(E') = num_bits(E), bits(result[i0, ..., iR-1]) = bits(operand[i0, ..., iR-1]).

bits renvoie la représentation en mémoire d'une valeur donnée. Son comportement est défini par l'implémentation, car la représentation exacte des Tensors et des types d'éléments est également définie par l'implémentation.

Entrées

Sorties

Contraintes

• (C1) Étant donné E = is_quantized(operand) ? storage_type(operand) : element_type(operand), E' = is_quantized(result) ? storage_type(result) : element_type(result) et R = rank(operand) :
  • Si num_bits(E') = num_bits(E), shape(result) = shape(operand).
  • Si num_bits(E') < num_bits(E) :
  • rank(result) = R + 1.
  • dim(result, i) = dim(operand, i) pour tous les 0 <= i < R.
  • dim(result, R) * num_bits(E') = num_bits(E).
  • Si num_bits(E') > num_bits(E) :
  • rank(result) = R - 1.
  • dim(result, i) = dim(operand, i) pour tous les 0 <= i < R.
  • dim(operand, R - 1) * num_bits(E) = num_bits(E').
• (C2) Si is_complex(operand) or is_complex(result), alorsis_complex(operand) and is_complex(result).

Exemples

// %operand: 0x0123456789ABCDEF
%result = "stablehlo.bitcast_convert"(%operand)< : >(te>nsorf64<) - t>ensor4xf16
// %result: [0xCDEF, 0x89AB, 0x4567, 0x0123] // little-endian representation

 Autres exemples

broadcast_in_dim

Sémantique

Développe les dimensions et/ou le rang d'un Tensor d'entrée en dupliquant les données dans le Tensor operand et produit un Tensor result. Plus formellement, result[result_index] = operand[operand_index] où pour tous les d dans axes(operand) :

• operand_index[d] = 0 si dim(operand, d) = 1.
• Sinon, operand_index[d] = result_index[broadcast_dimensions[d]].

Entrées

Sorties

Contraintes

• (C1) element_type(result) est donné par :
  • element_type(operand), si !is_per_axis_quantized(operand).
  • element_type(operand), sauf si quantization_dimension(operand), scales(operand) et zero_points(operand) diffèrent de quantization_dimension(result), scales(result) et zero_points(result), respectivement.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Pour tous les d dans axes(operand) :
  • dim(operand, d) = 1 ou
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Si is_per_axis_quantized(result) :
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Si dim(operand, quantization_dimension(operand)) = 1, alorsscales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))).

Exemples

// %operand: [
//            [1, 2, 3]
//           ]
%result = "stablehlo.broadcast_in_dim"(%operand) {
  broadcast_dimensio<ns = arra>yi64: 2, 1
}< : (ten>sor>1x3xi32<) - tenso>r2x3x2xi32
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

 Autres exemples

coque

Sémantique

Génère le résultat de l'exécution d'une seule fonction à partir de branches en fonction de la valeur de index. Plus précisément, result = selected_branch(), où :

• selected_branch = branches[index] si 0 <= index < size(branches).
• Sinon, selected_branch = branches[-1].

Entrées

Sorties

Contraintes

• (C1) 0 < size(branches).
• (C2) input_types(branches...) = [].
• (C3) same(output_types(branches...)).
• (C4) type(results...) = output_types(branches[0]).

Exemples

// %index: -1
// %result_branch0: [0, 0]
// %result_branch1: [1, 1]
%result0, %result1 = "stablehlo.case"(%index) ({
  "stablehlo.return"(%result_branch0, %resul<t_bra>nch0) : <(tens>or2>xi64, tensor2xi64) - ()
}, {
  "stablehlo.return"(%result_branc<h1, %>result_b<ranch>1) >: (tensor2xi64, <ten>sor>2xi64) -< ()
}>) : (ten<sori3>2) - (tensor2xi64, tensor2xi64)
// %result0: [1, 1]
// %result1: [1, 1]

 Autres exemples

cbrt

Sémantique

Effectue une opération de racine cubique élément par élément sur le Tensor operand et produit un Tensor result. En fonction du type d'élément, effectue les actions suivantes :

• Pour les valeurs float : rootn(x, 3) à partir de la norme IEEE-754.
• Pour les nombres complexes : racine cubique complexe.
• Pour les types quantifiés : dequantize_op_quantize(cbrt, operand, type(result))

Entrées

Sorties

Contraintes

• (C1) baseline_type(operand) = baseline_type(result).

Exemples

// %operand: [0.0, 1.0, 8.0, 27.0]
%result = "stablehlo.cbrt"(%operand)< : (t>ens>or4xf64<) - t>ensor4xf64
// %result: [0.0, 1.0, 2.0, 3.0]

 Autres exemples

ceil

Sémantique

Effectue un plafond élément par élément du Tensor operand et produit un Tensor result. Implémente l'opération roundToIntegralTowardPositive à partir de la spécification IEEE-754. Pour les types quantifiés, effectue dequantize_op_quantize(ceil, operand, type(result)).

Entrées

Sorties

Contraintes

• (C1) baseline_type(operand) = baseline_type(result).

Exemples

// %operand: [-0.8166, -0.2530, 0.2530, 0.8166, 2.0]
%result = "stablehlo.ceil"(%operand)< : (t>ens>or5xf32<) - t>ensor5xf32
// %result: [-0.0, -0.0, 1.0, 1.0, 2.0]

 Autres exemples

cholesky

Sémantique

Calcule la décomposition de Cholesky d'un lot de matrices.

Plus formellement, pour tous les i dans index_space(result), result[i0, ..., iR-3, :, :] est une décomposition de Cholesky de a[i0, ..., iR-3, :, :], sous la forme d'une matrice triangulaire inférieure (si lower est true) ou supérieure (si lower est false). Les valeurs de sortie dans le triangle opposé, c'est-à-dire le triangle supérieur strict ou le triangle inférieur strict, sont définies par l'implémentation.

Si i existe et que la matrice d'entrée n'est pas une matrice hermitienne définie positive, le comportement n'est pas défini.

Pour les types quantifiés, effectue dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)).

Entrées

Sorties

Contraintes

• (C1) baseline_type(a) = baseline_type(result).
• (C2) 2 <= rank(a).
• (C3) dim(a, -2) = dim(a, -1).

Exemples

// %a: [
//      [1.0, 2.0, 3.0],
//      [2.0, 20.0, 26.0],
//      [3.0, 26.0, 70.0]
//     ]
%result = "stablehlo.cholesky"(%a) {
  lower = true
}< : (ten>sor>3x3xf32<) - ten>sor3x3xf64
// %result: [
//           [1.0, 0.0, 0.0],
//           [2.0, 4.0, 0.0],
//           [3.0, 5.0, 6.0]
//          ]

limiter

Sémantique

Limite chaque élément du Tensor operand entre une valeur minimale et maximale, et produit un Tensor result. Plus formellement, result[result_index] = minimum(maximum(operand[result_index], min_element), max_element), où min_element = rank(min) = 0 ? min[] : min[result_index], max_element = rank(max) = 0 ? max[] : max[result_index]. Pour les types quantifiés, effectue dequantize_op_quantize(clamp, min, operand, max, type(result)).

L'imposition d'un ordre sur les nombres complexes implique une sémantique surprenante. Nous prévoyons donc de supprimer la prise en charge des nombres complexes pour cette opération à l'avenir (#560).

Entrées

Sorties

Contraintes

• (C1) rank(min) = 0 or shape(min) = shape(operand).
• (C2) rank(max) = 0 or shape(max) = shape(operand).
• (C3) baseline_element_type(min) = baseline_element_type(operand) = baseline_element_type(max).
• (C4) baseline_type(operand) = baseline_type(result).

Exemples

// %min: [5, 10, 15]
// %operand: [3, 13, 23]
// %max: [10, 15, 20]
%result = "stablehlo.clamp"(%min, %operand, %max)< : (t>ensor3xi<32, t>ensor3xi<32, t>ens>or3xi32<) - t>ensor3xi32
// %result: [5, 13, 20]

 Autres exemples

collective_broadcast

Sémantique

Dans chaque groupe de processus de la grille de processus StableHLO, envoyez la valeur du tenseur operand du processus source aux processus cibles et générez un tenseur result.

L'opération divise la grille de processus StableHLO en process_groups, qui est défini comme suit :

• cross_replica(replica_groups) si channel_id <= 0.
• cross_partition(replica_groups) si channel_id > 0.

result@process est ensuite donné par :

• operand@process_groups[i, 0] s'il existe un i tel que le processus soit dans process_groups[i].
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) dans le cas contraire.

Entrées

Sorties

Contraintes

• (C1) is_unique(replica_groups).
• (C2) 0 <= replica_groups < N où N est défini comme suit :
  • num_replicas si cross_replica est utilisé.
  • num_partitions si cross_partition est utilisé.
• (C3) type(result) = type(operand).

Exemples

// num_replicas: 4
// num_partitions: 1
// %operand@(0, 0): [[1, 2]]
// %operand@(1, 0): [[3, 4]]
// %operand@(2, 0): [[5, 6]]
// %operand@(3, 0): [[7, 8]]
%result = "stablehlo.collective_broadcast"(%operand) {
  replica_grou<ps = den>se[[2, 1]<] : ten>sor1x2xi64,
  channel_handle = #stablehlo.chan<nel_handlehandle = 0>, type = 0
} : (ten>sor>1x2xi64<) - ten>sor1x2xi64
// %result@(0, 0): [[0, 0]]
// %result@(1, 0): [[5, 6]]
// %result@(2, 0): [[5, 6]]
// %result@(3, 0): [[0, 0]]

collective_permute

Sémantique

Dans chaque groupe de processus de la grille de processus StableHLO, envoie la valeur du tenseur operand du processus source au processus cible et produit un tenseur result.

L'opération divise la grille de processus StableHLO en process_groups, qui est défini comme suit :

• cross_replica(source_target_pairs) si channel_id <= 0.
• cross_partition(source_target_pairs) si channel_id > 0.

result@process est ensuite donné par :

• operand@process_groups[i, 0], s'il existe un i tel que process_groups[i, 1] = process.
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) dans le cas contraire.

Entrées

Sorties

Contraintes

• (C1) dim(source_target_pairs, 1) = 2.
• (C2) is_unique(source_target_pairs[:, 0]).
• (C3) is_unique(source_target_pairs[:, 1]).
• (C4) 0 <= source_target_pairs < N, où N est défini comme suit :
  • num_replicas si cross_replica est utilisé.
  • num_partitions si cross_partition est utilisé.
• (C5) type(result) = type(operand).

Exemples

// num_replicas: 3
// num_partitions: 1
// %operand@(0, 0): [[1, 2], [3, 4]]
// %operand@(1, 0): [[5, 6], [7, 8]]
// %operand@(2, 0): [[9, 10], [11, 12]]
%result = "stablehlo.collective_permute"(%operand) {
  source_target_pai<rs = dense[[0, 1>], [1, 2]<] : ten>sor2x2xi64,
  channel_handle = #stablehlo.chan<nel_handlehandle = 0>, type = 0
}< : (ten>sor>2x2xi64<) - ten>sor2x2xi64
//
// %result@(0, 0): [[0, 0], [0, 0]]
// %result@(1, 0): [[1, 2], [3, 4]]
// %result@(2, 0): [[5, 6], [7, 8]]

 Autres exemples

compare

Sémantique

Effectue une comparaison élément par élément des Tensors lhs et rhs selon comparison_direction et compare_type, et produit un Tensor result.

Les valeurs de comparison_direction et compare_type ont la sémantique suivante :

Pour les types d'éléments booléens et entiers :

• EQ : lhs = rhs.
• NE : lhs != rhs.
• GE : lhs >= rhs.
• GT : lhs > rhs.
• LE : lhs <= rhs.
• LT : lhs < rhs.

Pour les types d'éléments à virgule flottante avec compare_type = FLOAT, l'opération implémente les opérations IEEE-754 suivantes :

• EQ : compareQuietEqual.
• NE : compareQuietNotEqual.
• GE : compareQuietGreaterEqual.
• GT : compareQuietGreater.
• LE : compareQuietLessEqual.
• LT : compareQuietLess.

Pour les types d'éléments à virgule flottante avec compare_type = TOTALORDER, l'opération utilise la combinaison des opérations totalOrder et compareQuietEqual de la norme IEEE-754.

Pour les types d'éléments complexes, la comparaison lexicographique des paires (real, imag) est effectuée à l'aide des comparison_direction et compare_type fournis. L'imposition d'un ordre sur les nombres complexes implique une sémantique surprenante. Nous prévoyons donc de supprimer la prise en charge des nombres complexes lorsque comparison_direction est GE, GT, LE ou LT (#560).

Pour les types quantifiés, effectue dequantize_compare(lhs, rhs, comparison_direction).

Entrées

Sorties

Contraintes

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs).
• (C2) shape(lhs) = shape(rhs) = shape(result).
• (C3) compare_type est défini comme suit :
  • SIGNED si is_signed_integer(element_type(lhs)).
  • UNSIGNED si is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)).
  • FLOAT ou TOTALORDER si is_float(element_type(lhs)).
  • FLOAT si is_complex(element_type(lhs)).

Exemples

// %lhs: [1.0, 3.0]
// %rhs: [1.1, 2.9]
%result = "stablehlo.compare"(%lhs, %rhs) {
  comparison_direction = <#stablehlocomparison_di>rection LT,
  compare_type = <#stablehlocomparison_>type FLOAT
}< : (t>ensor2xf<32, t>ens>or2xf32<) - >tensor2xi1
// %result: [true, false]

 Autres exemples

complexe

Sémantique

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

Entrées

Sorties

Contraintes

• (C1) type(lhs) = type(rhs).
• (C2) shape(result) = shape(lhs).
• (C3) element_type(result) est de type complex<E> où E = element_type(lhs).

Exemples

// %lhs: [1.0, 3.0]
// %rhs: [2.0, 4.0]
%result = "stablehlo.complex"(%lhs, %rhs)< : (t>ensor2xf<64, t>ens>or2xf64<) - tenso<r2x>>complexf64
// %result: [(1.0, 2.0), (3.0, 4.0)]

 Autres exemples

composite

Sémantique

Encapsule une opération composée d'autres opérations StableHLO, en prenant inputs et composite_attributes et en produisant results. La sémantique de l'opération est implémentée par l'attribut decomposition. 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, préférez utiliser custom_call.

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

Entrées

Sorties

Contraintes

• (C1) is_namespaced_op_name(name)
• (C2) is_defined_in_parent_scope(decomposition)
• (C3) types(inputs...) == input_types(decomposition)
• (C4) types(results...) == output_types(decomposition)

Exemples

%results = "stablehlo.composite"(%input0, %input1) {
  name = "my_namespace.my_op",
  composite_attributes = {
    my_attribute = "my_value"
  },
  decomposition = @my_op,
 < ve>rsion = <1 :> i3>2
} : (<ten>sorf32, tensorf32) - tensorf32

 Autres exemples

concatenate

Sémantique

Concatène inputs le long de la dimension dimension dans le même ordre que les arguments fournis et produit un Tensor result. Plus formellement, result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1], où :

1. id = d0 + ... + dk-1 + kd.
2. d est égal à dimension, et d0, ... sont les tailles de la de dimension de inputs.

Entrées

Sorties

Contraintes

• (C1) same(element_type(inputs...)).
• (C2) same(shape(inputs...)), sauf dim(inputs..., dimension).
• (C3) 0 < size(inputs).
• (C4) 0 <= dimension < rank(inputs[0]).
• (C5) element_type(result) = element_type(inputs[0]).
• (C6) shape(result) = shape(inputs[0]), sauf :
  • dim(result, dimension) = dim(inputs[0], dimension) + ....

Exemples

// %input0: [[1, 2], [3, 4], [5, 6]]
// %input1: [[7, 8]]
%result = "stablehlo.concatenate"(%input0, %input1) {
  dimension = 0 : i64
}< : (ten>sor3x2xi<64, ten>sor>1x2xi64<) - ten>sor4x2xi64
// %result: [[1, 2], [3, 4], [5, 6], [7, 8]]

 Autres exemples

constante

Sémantique

Génère un Tensor output à partir d'une constante value.

Entrées

Sorties

Contraintes

• (C1) type(value) = type(output).

Exemples

%output = "stablehlo.constant"() {
  val<ue = dense[[0.0, 1.0], [>2.0, 3.0]<] : ten>sor2x2xf3>2
} : (<) - ten>sor2x2xf32
// %output: [[0.0, 1.0], [2.0, 3.0]]

 Autres exemples

d'effectuer une conversion

Sémantique

Effectue une conversion élément par élément d'un type d'élément à un autre sur le Tensor operand et produit un Tensor result.

Pour les conversions boolean-to-any-supported-type, la valeur false est convertie en zéro et la valeur true est convertie en un. Pour les conversions any-supported-type-to-boolean, une valeur nulle est convertie en false et les valeurs non nulles sont converties en true. Vous trouverez ci-dessous des informations sur le fonctionnement de cette fonctionnalité pour les types complexes.

Pour les conversions d'entier à entier, d'entier à virgule flottante ou de virgule flottante à virgule flottante, si la valeur source peut être représentée exactement dans le type de destination, la valeur résultante est cette représentation exacte. Sinon, le comportement est à déterminer (#180).

Pour les conversions floating-point-to-integer, la partie fractionnaire est tronquée. Si la valeur tronquée ne peut pas être représentée dans le type de destination, le comportement est à déterminer (#180).

Les conversions complexe vers complexe suivent le même comportement que les conversions virgule flottante vers virgule flottante pour convertir les parties réelles et imaginaires.

Pour les conversions complex-to-any-other-type et any-other-type-to-complex, la valeur imaginaire source est ignorée ou la valeur imaginaire de destination est définie sur zéro, respectivement. La conversion de la partie réelle suit les conversions à virgule flottante.

En principe, cette opération pourrait exprimer la déquantification (conversion de Tensors quantifiés en Tensors réguliers), la quantification (conversion de Tensors réguliers en Tensors quantifiés) et la requantification (conversion entre Tensors quantifiés), mais pour le moment, nous avons des opérations dédiées pour cela : uniform_dequantize pour le premier cas d'utilisation et uniform_quantize pour les deuxième et troisième cas d'utilisation. À l'avenir, ces deux opérations pourront être fusionnées dans convert (#1576).

Entrées

Sorties

Contraintes

• (C1) shape(operand) = shape(result).

Exemples

// %operand: [-1, 0, 1]
%result = "stablehlo.convert"(%operand)< : (t>ens>or3xi64<) - tenso<r3x>>complexf64
// %result: [(-1.0, 0.0), (0.0, 0.0), (1.0, 0.0)]

 Autres exemples

convolution

Sémantique

Calcule les produits scalaires entre les fenêtres de lhs et les tranches de rhs, et génère result. Le diagramme suivant montre comment les éléments de result sont calculés à partir de lhs et rhs à l'aide d'un exemple concret.

Plus formellement, considérez la reformulation suivante des entrées en termes de lhs afin de pouvoir exprimer des fenêtres de lhs :

• lhs_window_dimensions = lhs_shape(dim(lhs, input_batch_dimension), dim(rhs, kernel_spatial_dimensions), dim(lhs, input_feature_dimension)).
• lhs_window_strides = lhs_shape(1, window_strides, 1)
• lhs_padding = lhs_shape([0, 0], padding, [0, 0])
• lhs_base_dilations = lhs_shape(1, lhs_dilation, 1)
• lhs_window_dilations = lhs_shape(1, rhs_dilation, 1).

Cette reformulation utilise les fonctions d'assistance suivantes :

• lhs_shape(n, hw, c) = permute([n] + hw + [c], [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]).
• result_shape(n1, hw, c1) = permute([n1] + hw + [c1], [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]).
• permute([j0, j1, ..., jR-1], permutation) = [i0, i1, ..., iR-1] où j[d] = i[permutation[d]].

Si feature_group_count = 1 et batch_group_count = 1, alors pour tous les output_spatial_index dans index_space(dim(result, output_spatial_dimensions...)), result[result_shape(:, output_spatial_index, :)] = dot_product où :

• padding_value = constant(0, element_type(lhs)).
• padded_lhs = pad(lhs, padding_value, lhs_padding[:, 0], lhs_padding[:, 1], lhs_base_dilations - 1)
• lhs_window_start = lhs_shape(0, output_spatial_index, 0) * lhs_window_strides
• lhs_window = slice(padded_lhs, lhs_window_start, lhs_window_start + lhs_window_dimensions, lhs_window_dilations).
• reversed_lhs_window = reverse(lhs_window, [input_spatial_dimensions[dim] for dim in range(size(window_reversal)) if window_reversal[dim] = true]). Cette fonctionnalité semble inutilisée. Nous prévoyons donc de la supprimer à l'avenir (#1181).
• dot_product = dot_general(reversed_lhs_window, rhs, lhs_batching_dimensions=[], lhs_contracting_dimensions=input_spatial_dimensions + [input_feature_dimension], rhs_batching_dimensions=[], rhs_contracting_dimensions=kernel_spatial_dimensions + [kernel_input_feature_dimension]).

Si feature_group_count > 1 :

• lhses = split(lhs, feature_group_count, input_feature_dimension).
• rhses = split(rhs, feature_group_count, kernel_output_feature_dimension)
• results... = convolution(lhses..., rhses..., ..., feature_group_count=1, ...)
• result = concatenate(results, output_feature_dimension).

Si batch_group_count > 1 :

• lhses = split(lhs, batch_group_count, input_batch_dimension).
• rhses = split(rhs, batch_group_count, kernel_output_feature_dimension)
• results... = convolution(lhses..., rhses..., ..., batch_group_count=1, ...).
• result = concatenate(results, output_feature_dimension).

Pour les types quantifiés, effectue dequantize_op_quantize( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs, type(result)).

Pour les types quantifiés hybrides, effectue hybrid_dequantize_then_op( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs).

Entrées

Sorties

Contraintes

• (C1) N = rank(lhs) = rank(rhs).
• (C2) size(window_strides) = N - 2.
• (C3) 0 < window_strides.
• (C4) shape(padding) = [N - 2, 2].
• (C5) size(lhs_dilation) = N - 2.
• (C6) 0 < lhs_dilation.
• (C7) size(rhs_dilation) = N - 2.
• (C8) 0 < rhs_dilation.
• (C9) size(window_reversal) = N - 2.
• (C10) dim(lhs, input_batch_dimension) % batch_group_count = 0.
• (C11) dim(lhs, input_feature_dimension) % feature_group_count = 0.
• (C12) size(input_spatial_dimensions) = N - 2.
• (C13) Étant donné input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension] :
  • is_unique(input_dimensions).
  • 0 <= input_dimensions < N.
• (C14) dim(rhs, kernel_input_feature_dimension) = dim(lhs, input_feature_dimension) / feature_group_count.
• (C15) dim(rhs, kernel_output_feature_dimension) % batch_group_count = 0.
• (C16) dim(rhs, kernel_output_feature_dimension) % feature_group_count = 0.
• (C17) size(kernel_spatial_dimensions) = N - 2.
• (C18) Étant donné kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension] :
  • is_unique(kernel_dimensions).
  • 0 <= kernel_dimensions < N.
• (C19) size(output_spatial_dimensions) = N - 2.
• (C20) Étant donné output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension] :
  • is_unique(output_dimensions).
  • 0 <= output_dimensions < N.
• (C21) 0 < feature_group_count.
• (C22) 0 < batch_group_count.
• (C23) feature_group_count = 1 or batch_group_count = 1.
• (C24) size(precision_config) = 2.
• (C25) dim(result, result_dim) est défini comme suit :
  • dim(lhs, input_batch_dimension) / batch_group_count si result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) si result_dim = output_feature_dimension.
  • num_windows sinon, où :
  • output_spatial_dimensions[spatial_dim] = result_dim.
  • lhs_dim = input_spatial_dimensions[spatial_dim]
  • rhs_dim = kernel_spatial_dimensions[spatial_dim]
  • dilated_input_shape[lhs_dim] = dim(lhs, lhs_dim) = 0 ? 0 : (dim(lhs, lhs_dim) - 1) * lhs_dilation[spatial_dim] + 1
  • padded_input_shape[lhs_dim] = padding[spatial_dim, 0] + dilated_input_shape[lhs_dim] + padding[spatial_dim, 1]
  • dilated_window_shape[lhs_dim] = dim(rhs, rhs_dim) = 0 ? 0 : (dim(rhs, rhs_dim) - 1) * rhs_dilation[spatial_dim] + 1
  • is_empty_window[lhs_dim] = padded_input_shape[lhs_dim] = 0 || dilated_window_shape[lhs_dim] > padded_input_shape[lhs_dim]
  • num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1.
• (C26) rank(result) = N.
• Si l'opération utilise des Tensors non quantifiés :
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Si l'opération utilise des Tensors quantifiés :
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Si is_per_axis_quantized(rhs), alors quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Si is_per_axis_quantized(result), alorsquantization_dimension(result) = output_feature_dimension.
  • Si is_quantized(lhs) :
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Si is_per_tensor_quantized(rhs), alorsis_per_tensor_quantized(result).
  • Si !is_quantized(lhs) :
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Exemples

// %lhs: [[
//        [
//          [1], [2], [5], [6]
//        ],
//        [
//          [3], [4], [7], [8]
//        ],
//        [
//          [10], [11], [14], [15]
//        ],
//        [
//          [12], [13], [16], [17]
//        ]
//      ]]
//
// %rhs: [
//        [[[1]], [[1]], [[1]]],
//        [[[1]], [[1]], [[1]]],
//        [[[1]], [[1]], [[1]]]
//       ]
%result = "stablehlo.convolution"(%lhs, %rhs) {
  window_strid<es = arra>yi64: 4, 4,
  paddi<n>g = dense<0 : ten>sor2x2xi64,
  lhs_dilati<on = arra>yi64: 2, 2,
  rhs_dilati<on = arra>yi64: 1, 1,
  window_revers<al = arrayi1: fa>lse, false,
  // In the StableHLO dialect, dimension numbers are encoded vi<a:
  // `[input >dim<ensions]x[kernel >di>mensions]-[output dimensions]`.
  // "b" is batch dimension, "f" is feature dimension,
  // "i" is input feature dimension, "o" is output feature dimension,
  // "0/1/etc" a<re spatial dimensions.
  d>imension_num>bers = #stablehlo.conv[b, 0, 1, f]x[0, 1, i, o]-[b, 0, 1, f],
  batch_group_count = 1 : i64,
  fea<ture_group_count >= 1 : i64,
 < precision_config> = [#stablehl<oprecision >DEFAULT,< #stablehlo>pre>cision <DEFAULT]
} >: (tensor1x4x4x1xi64, tensor3x3x1x1xi64) - tensor1x2x2x1xi64
// %result: [[
//            [[10], [26]],
//            [[46], [62]]
//          ]]

 Autres exemples

cosinus

Sémantique

Effectue une opération de cosinus au niveau des éléments sur le Tensor operand et génère un Tensor result. En fonction du type d'élément, effectue les actions suivantes :

• Pour les valeurs float : cos à partir de la norme IEEE-754.
• Pour les nombres complexes : cosinus complexe.
• Pour les types quantifiés : dequantize_op_quantize(cosine, operand, type(result)).

Entrées

Sorties

Contraintes

• (C1) baseline_type(operand) = baseline_type(result).

Exemples

// %operand: [
//            [0.0, 1.57079632],       // [0, pi/2]
//            [3.14159265, 4.71238898] // [pi, 3pi/2]
//           ]
%result = "stablehlo.cosine"(%operand)< : (ten>sor>2x2xf32<) - ten>sor2x2xf32
// %result: [[1.0, 0.0], [-1.0, 0.0]]

 Autres exemples

count_leading_zeros

Sémantique

Effectue un décompte élément par élément du nombre de bits zéro de début dans le Tensor operand et produit un Tensor result.

Entrées

Sorties

Contraintes

• (C1) type(operand) = type(result).

Exemples

// %operand: [[0, 1], [128, -1]]
%result = "stablehlo.count_leading_zeros"(%operand)< : (ten>sor>2x2xi64<) - ten>sor2x2xi64
// %result: [[64, 63], [56, 0]]

 Autres exemples

custom_call

Sémantique

Encapsule une opération call_target_name définie par l'implémentation qui prend inputs et called_computations et produit results. has_side_effect, backend_config et api_version peuvent être utilisés pour fournir des métadonnées supplémentaires définies par l'implémentation.

Pour le moment, cette opération contient une collection de métadonnées assez désorganisée qui reflète l'évolution organique de son opération homologue dans le compilateur XLA. À l'avenir, nous prévoyons d'unifier ces métadonnées (#741).

Entrées

Sorties

(Compatibilité XLA avec les GPU) Cibles custom_call spéciales

Il existe trois call_target_name spéciaux liés aux types buffer : CreateBuffer crée un buffer non initialisé, Pin crée un buffer initialisé et Unpin libère un buffer et renvoie le contenu du buffer.

%uninitialized_buffer = "stablehlo.custom_call"() {
  call_target_name = "CreateBuffer",
  api_version> = 4 : <i32,
>} : () - memref4xf64

%initialized_buffer = "stablehlo.custom_call"(%init_value) {
  call_target_name = "Pin&quo<t;,
 > ap>i_versi<on = >4 : i32,
} : (tensor4xf64) - memref4xf64

%dealloc_buffer = "stablehlo.custom_call"(%initialized_buffer) {
  call_target_na<me = >&qu>ot;Unpi<n&quo>t;,
  api_version = 4 : i32,
} : (memref4xf64) - tensor4xf64

Alias

Certaines opérations custom_call peuvent nécessiter qu'une partie des sorties et une partie des opérandes partagent la même mémoire. Cela peut être exprimé via output_operand_aliases. Une représentation de paire d'alias se compose d'une liste d'indices de tuples de sortie représentant la partie de sortie, et d'un operand_index ainsi que d'une liste d'indices de tuples d'opérandes représentant la partie d'opérande. La liste des indices de tuples de sortie ou d'opérandes est vide si le type correspondant n'est pas un type tuple et peut être arbitrairement longue pour un type de tuple arbitrairement imbriqué. C'est semblable à la représentation de l'alias XLA.

La partie de sortie et la partie d'entrée d'une paire d'alias doivent être du même type. Pour les opérations custom_call qui ne sont pas des appels à CreateBuffer, Pin et Unpin, un opérande buffer peut apparaître dans une seule paire d'alias, et une sortie buffer doit apparaître dans une paire d'alias.

Exemples

%results = "stablehlo.custom_call"(%input0) {
  call_target_name = "foo",
  has_side_effect = false,
  backend_config = {bar = 42 : i32},
  api_version = 4 : i32,
  called_computations <= [>@fo>o]
} : <(te>nsorf64) - tensorf64

%updated_buffer = "stablehlo.custom_call"(%buffer) {
  call_target_name = "Update",
  api_version = 4 : i32,
  output_operand_aliases< = [
    #stablehlo.output_operand_aliasoutput_tuple_indices = [],
      operand_ind>ex = 0,
     < oper>and>_tuple_<indic>es = []]
} : (memref4xf64) - memref4xf64

diviser

Sémantique

Effectue une division élément par élément des Tensors de dividende lhs et de diviseur rhs, et produit un Tensor result. En fonction du type d'élément, effectue les actions suivantes :

• Pour les nombres entiers : division entière qui produit le quotient algébrique en supprimant toute partie fractionnaire.
• Pour les valeurs float : division à partir de la norme IEEE-754.
• Pour les nombres complexes : division complexe.
• Pour les types quantifiés :
  • dequantize_op_quantize(divide, lhs, rhs, type(result)).

Entrées

Sorties

Contraintes

• (C1) baseline_type(lhs) = baseline_type(rhs) = baseline_type(result).

Exemples

// %lhs: [17.1, -17.1, 17.1, -17.1]
// %rhs: [3.0, 3.0, -3.0, -3.0]
%result = "stablehlo.divide"(%lhs, %rhs)< : (t>ensor4xf<32, t>ens>or4xf32<) - t>ensor4xf32
// %result: [5.66666651, -5.66666651, -5.66666651, 5.66666651]

 Autres exemples

dot_general

Sémantique

Calcule les produits scalaires entre les tranches de lhs et celles de rhs, et génère un Tensor result.

Plus précisément, result[result_index] = dot_product, où :

• lhs_result_dimensions = [d for d in axes(lhs) and d not in lhs_batching_dimensions and d not in lhs_contracting_dimensions].
• rhs_result_dimensions = [d for d in axes(rhs) and d not in rhs_batching_dimensions and d not in rhs_contracting_dimensions].
• result_batching_index + result_lhs_index + result_rhs_index = result_index où size(result_batching_index) = size(lhs_batching_dimensions), size(result_lhs_index) = size(lhs_result_dimensions) et size(result_rhs_index) = size(rhs_result_dimensions).
• transposed_lhs = transpose(lhs, lhs_batching_dimensions + lhs_result_dimensions + lhs_contracting_dimensions).
• transposed_lhs_slice = slice(transposed_lhs, result_batching_index + result_lhs_index + [:, ..., :])
• reshaped_lhs_slice = reshape(transposed_lhs_slice, dims(lhs, lhs_contracting_dimensions))
• transposed_rhs = transpose(rhs, rhs_batching_dimensions + rhs_result_dimensions + rhs_contracting_dimensions)
• transposed_rhs_slice = slice(transposed_rhs, result_batching_index + result_rhs_index + [:, ..., :])
• reshaped_rhs_slice = reshape(transposed_rhs_slice, dims(rhs, rhs_contracting_dimensions)).
• dot_product = reduce( inputs=[multiply(reshaped_lhs_slice, reshaped_rhs_slice)], init_values=[constant(0, element_type(result))], dimensions=range(size(lhs_contracting_dimensions)), body=lambda x, y: add(x, y)).

Pour les types quantifiés, effectue dequantize_op_quantize( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs, type(result)).

Pour les types quantifiés hybrides, effectue hybrid_dequantize_then_op( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs).

precision_config contrôle le compromis entre vitesse et précision pour les calculs sur les backends d'accélérateur. Il peut s'agir de l'une des valeurs suivantes (pour le moment, la sémantique de ces valeurs d'énumération n'est pas suffisamment spécifiée, mais nous prévoyons de résoudre ce problème dans #755) :

• DEFAULT : calcul le plus rapide, mais approximation la moins précise du nombre d'origine.
• HIGH : calcul plus lent, mais approximation plus précise du nombre d'origine.
• HIGHEST : le calcul le plus lent, mais l'approximation la plus précise du nombre d'origine.

Un DotAlgorithm définit les propriétés principales de l'algorithme utilisé pour implémenter l'opération de produit scalaire, qui définit également la précision. Si les champs d'attribut d'algorithme sont définis, precision_config doit être DEFAULT. DotAlgorithms n'ont pas de valeur par défaut, car les paramètres par défaut sont définis par l'implémentation. Par conséquent, tous les champs de l'algorithme de points peuvent être définis sur None pour spécifier un algorithme de points vide, qui utilisera plutôt la valeur precision_config.

Les champs DotAlgorithm incluent les suivants :

• lhs_precision_type et rhs_precision_type, les précisions auxquelles les côtés gauche et droit de l'opération sont arrondis. Les types de précision sont indépendants des types de stockage des entrées et de la sortie.
• accumulation_type : précision utilisée pour l'accumulation.
• lhs_component_count, rhs_component_count et num_primitive_operations s'appliquent lorsque nous effectuons un algorithme qui décompose les côtés gauche et/ou droit en plusieurs composants et effectue plusieurs opérations de produit scalaire "primitives" sur ces valeurs, généralement pour émuler une précision plus élevée (par exemple, Utiliser le type de données d'intelligence artificielle bfloat16 pour des calculs plus précis : bf16_6x tf32_3x, etc.). Pour les algorithmes sans décomposition, ces valeurs doivent être définies sur 1.
• allow_imprecise_accumulation pour spécifier si l'accumulation dans une précision inférieure est autorisée pour certaines étapes (par exemple, CUBLASLT_MATMUL_DESC_FAST_ACCUM).

Exemples d'attributs DotAlgorithm :

// Inputs are casted to tf32, and then accumulated in f32:
{lhs_precision_type = tf32,
 rhs_precision_type = tf32,
 accumulation_type = f32,
 lhs_component_count = 1,
 rhs_component_count = 1,
 num_primitive_operations = 1,
 allow_imprecise_accumulation = false}


// bf16_6x: each input is decomposed to 3 bf16 components, then 6 dot operations are done on those components, and the result is accumulated in f32.
{lhs_precision_type = bf16,
 rhs_precision_type = bf16,
 accumulation_type = f32,
 lhs_component_count = 3,
 rhs_component_count = 3,
 num_primitive_operations = 6,
 allow_imprecise_accumulation = false}


// Inputs are (casted to) f8e5m2, and we accumulate in f32, but for some steps we may accumulate in lower precision.
{lhs_precision_type = f8e5m2,
 rhs_precision_type = f8e5m2,
 accumulation_type = f32,
 lhs_component_count = 1,
 rhs_component_count = 1,
 num_primitive_operations = 1,
 allow_imprecise_accumulation = true}

Il appartient aux implémentations de décider quelles combinaisons sont prises en charge. En général, il n'est pas garanti que chaque algorithme soit compatible avec chaque type d'accélérateur par le consommateur de StableHLO. Si un algorithme donné n'est pas pris en charge, une erreur doit être générée au lieu de revenir à une alternative. La validation StableHLO fournira la meilleure vérification possible, en empêchant les algorithmes qui ne sont pas connus pour être compatibles avec aucun matériel.

Pour obtenir des exemples de valeurs d'algorithme acceptées, consultez xla_data.proto > Algorithm. La demande 2483 décrit le plan de création d'un document centralisé sur les algorithmes compatibles par backend.

Entrées

Sorties

Contraintes

• (C1) size(lhs_batching_dimensions) = size(rhs_batching_dimensions).
• (C2) size(lhs_contracting_dimensions) = size(rhs_contracting_dimensions).
• (C3) is_unique(lhs_batching_dimensions + lhs_contracting_dimensions).
• (C4) is_unique(rhs_batching_dimensions + rhs_contracting_dimensions).
• (C5) 0 <= lhs_batching_dimensions < rank(lhs).
• (C6) 0 <= lhs_contracting_dimensions < rank(lhs).
• (C7) 0 <= rhs_batching_dimensions < rank(rhs).
• (C8) 0 <= rhs_contracting_dimensions < rank(rhs).
• (C9) dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...).
• (C10) dim(lhs, lhs_contracting_dimensions...) = dim(rhs, rhs_contracting_dimensions...).
• (C11) size(precision_config) = 2.
• (C12) shape(result) = dim(lhs, lhs_batching_dimensions) + dim(lhs, lhs_result_dimensions) + dim(rhs, rhs_result_dimensions).
• Si l'opération utilise des Tensors non quantifiés :
  • (C13) element_type(lhs) = element_type(rhs).
• Si l'opération utilise des Tensors quantifiés :
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C15) zero_points(rhs) = 0.
  • (C16) Si is_per_axis_quantized(rhs), alors quantization_dimension(rhs) n'est pas dans rhs_contracting_dimensions.
  • Si is_quantized(lhs) :
  • (C17) storage_type(lhs) = storage_type(rhs).
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C19) Si is_per_tensor_quantized(rhs), alors is_per_tensor_quantized(result).
  • Si !is_quantized(lhs) :
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result).
• Si !is_empty_algorithm(lhs_precision_type, rhs_precision_type, accumulation_type, lhs_component_count, rhs_component_count, num_primitive_operations allow_imprecise_accumulation) :
  • (C21) precision_config... = DEFAULT.
  • (C22) 0 < lhs_component_count.
  • (C23) 0 < rhs_component_count.
  • (C24) 0 < num_primitive_operations.

Exemples

// %lhs: [
//        [[1, 2],
//         [3, 4]],
//        [[5, 6],
//         [7, 8]]
//       ]
// %rhs: [
//        [[1, 0],
//         [0, 1]],
//        [[1, 0],
//         [0, 1]]
//       ]
%result = "stablehlo.dot_general"(%lhs, %rhs) {
  dot_dimension_numbers = #sta<blehlo.dot
    lhs_batching_dimensions = [0],
    rhs_batching_dimensions = [0],
    lhs_contracting_dimensions = [2],
    rhs_contracting_dimension>s = [1]
  ,
  precision_config = [<#stablehloprecisi>on DEFAULT, <#stablehloprecisi>on DEFAULT],
  algorithm = #stablehlo.dot<_algorithm
    lhs_precision_type = tf32,
    rhs_precision_type = tf32,
    accumulation_type = f32,
    lhs_component_count = 1,
    rhs_component_count = 1,
    num_primitive_operations = 1,
    allow_imprecise_accumulation >= false
  
}< : (tenso>r2x2x2xi<64, tenso>r2x>2x2xi64<) - tenso>r2x2x2xi64
// %result: [
//           [[1, 2],
//            [3, 4]],
//           [[5, 6],
//            [7, 8]]
//          ]

 Autres exemples

dynamic_broadcast_in_dim

Sémantique

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

L'opération accepte également les attributs facultatifs known_expanding_dimensions et known_nonexpanding_dimensions pour exprimer des connaissances statiques sur le comportement d'expansion des dimensions. Si aucune n'est spécifiée, toutes les dimensions sont considérées comme pouvant être développées.

Entrées

Sorties

Contraintes

• (C1) element_type(result) est donné par :
  • element_type(operand), si !is_per_axis_quantized(operand).
  • element_type(operand), sauf si quantization_dimension(operand), scales(operand) et zero_points(operand) diffèrent de quantization_dimension(result), scales(result) et zero_points(result), respectivement.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Pour tous les d dans axes(operand) :
  • dim(operand, d) = 1 ou
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Si is_per_axis_quantized(result) :
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Si dim(operand, quantization_dimension(operand)) = 1, alorsscales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))).
• (C7) size(output_dimensions) = rank(result).
• (C8) is_unique(known_expanding_dimensions + known_nonexpanding_dimensions).
• (C9) 0 <= known_expanding_dimensions < rank(operand).
• (C10) 0 <= known_nonexpanding_dimensions < rank(operand).

Exemples

// %operand: [
//            [1, 2, 3]
//           ]
%operand = stablehlo.constant dense<[[1, 2, 3]]> : tensor<1x3xi64>
%output_dimensions = stablehlo.constant dense<[2, 3, 2]> : tensor<3xi64>
%result = "stablehlo.dynamic_broadcast_in_dim"(%operand, %output_dimensions) {
  broadcast_dimensio<ns = arra>yi64: 2, 1,
  known_expanding_dimensio<ns = a>rrayi64: 0,
  known_nonexpanding_dimensio<ns = a>rrayi64: 1
}< : (ten>sor1x3xi<64, t>ens>or3xi64<) - tenso>r2x3x2xi64
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

 Autres exemples

dynamic_conv

Sémantique

Cette opération est fonctionnellement identique à l'opération convolution, mais le remplissage est spécifié de manière dynamique via padding.

Entrées