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 fonctionne comme une couche de portabilité entre différents frameworks de ML 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 interopérabilité plus poussée 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 consistent en des fonctions StableHLO qui, à leur tour, consistent en des opérations StableHLO. Au sein de cette structure, la section Ops spécifie la sémantique des opérations individuelles. La section Execution fournit une sémantique pour toutes ces opérations exécutées ensemble dans un programme. Enfin, la section Notation décrit la notation utilisée dans l'ensemble de la spécification.

Pour afficher les spécifications d'une version précédente de StableHLO, ouvrez le dépôt à la version taguée de votre choix. Par exemple, la spécification StableHLO v0.19.0. Pour afficher les modifications apportées à chaque version 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 6 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) possèdent 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 comportent des sigils 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
NonValueType ::= TensorElementType | QuantizedTensorElementType | FunctionType | StringType

Les types StableHLO sont classés en types de valeurs (également appelés types de première classe) qui représentent les valeurs StableHLO et en types non de valeur qui décrivent d'autres éléments de 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 de Tensor représentent des Tensors, c'est-à-dire des tableaux multidimensionnels. Ils 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é classement. Par exemple, tensor<2x3xf32> est un type de tenseur avec la forme 2x3 et le 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 2 et 3. Son classement est de 2.

Les formes peuvent être partiellement ou totalement inconnues (dynamiques), par exemple, tensor<?x2xf64> est partiellement inconnu et tensor<?x?xf64> est totalement inconnu. Les tailles de dimension dynamique sont représentées à l'aide d'un ?. Les formes ne peuvent pas être désclassées.

À l'avenir, nous prévoyons d'étendre les types de tenseurs au-delà des tailles de dimension et des types d'éléments, par exemple pour inclure les mises en page (numéro 629) et la sparsité (numéro 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 les valeurs entières d'un type de stockage compris entre storage_min et storage_max (inclus) qui correspondent aux 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 comme 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 leurs valeurs par défaut sont respectivement min_value(storage_type) et max_value(storage_type). Les types d'éléments quantifiés présentent les 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 les échelles basées sur des entiers, représentées par des multiplicateurs et des décalages, suscitent un vif intérêt. Nous prévoyons d'explorer cette question prochainement (#1404).

La sémantique de QuantizationZeroPoint est en cours de discussion, y compris concernant le type et les valeurs, et sur la possibilité de n'avoir qu'un seul ou plusieurs points zéro dans un type de Tensor quantifié. D'après les résultats de cette discussion, les spécifications concernant les points nuls peuvent changer à l'avenir (numéro 1405).

Une autre discussion en cours concerne la sémantique de QuantizationStorageMin et QuantizationStorageMax pour déterminer si des contraintes doivent être imposées à ces valeurs et aux valeurs des tenseurs quantifiés (numéro 1406).

Enfin, nous prévoyons d'explorer la représentation d'échelles et de points zéro inconnus, comme nous prévoyons de le faire pour la représentation de tailles de dimension inconnues (numéro 1407).

Les types de tenseurs quantifiés représentent des tenseurs dont les éléments sont 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, et non des types d'éléments standards.

Dans les tenseurs quantifiés, la quantification peut être par Tensor, ce qui signifie qu'elle peut avoir une valeur scale et zero_point pour l'ensemble du Tensor ou par axe, ce qui signifie qu'elle peut avoir plusieurs scales et zero_points, une paire par tranche d'une dimension quantization_dimension particulière. Plus formellement, dans un tenseur t avec quantification par axe, il existe dim(t, quantization_dimension) tranches de quantization_dimension : t[:, ..., 0, ..., :], t[:, ..., 1, ..., :], etc. Tous les éléments de la ie tranche utilisent scales[i] et zero_points[i] comme paramètres de quantification. Les types de Tensors quantifiés présentent les contraintes suivantes:

• Pour la quantification par tenseur :
  • 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 l'ordre d'exécution aux opérations, comme décrit dans la section Exécution.

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

Les types de tuple représentent des tupels, c'est-à-dire des listes hétérogènes. Les tupels sont une ancienne fonctionnalité qui n'existe que pour la compatibilité avec HLO. Dans HLO, les tupels sont utilisés pour représenter les entrées et les sorties variadiques. Dans StableHLO, les entrées et sorties variadiques sont prises en charge en mode natif, et le seul usage des tupels dans StableHLO est de représenter de manière exhaustive l'ABI HLO, par exemple, T, tuple<T> et tuple<tuple<T>> peuvent être sensiblement différents en fonction d'une implémentation particulière. À l'avenir, nous prévoyons d'apporter des modifications à l'ABI HLO, ce qui nous permettra peut-être de supprimer les types de tuple de StableHLO (numéro 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 de 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 courant de représenter les valeurs scalaires de type T avec des valeurs de tenseur à dimension 0 de type tensor<T>).

• Le type booléen représente les valeurs booléennes true et false.
• Les types d'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 les valeurs d'entiers de -2^(N-1) à 2^(N-1)-1 inclus, et les types uiN non signés représentent les valeurs d'entiers de 0 à 2^N-1 inclus.
• Les types à virgule flottante peuvent être l'un des éléments suivants :
  • f8E3M4, f8E4M3 et f8E5M2 : 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écrit dans la section Formats FP8 pour l'apprentissage profond.
  • Types f8E4M3FNUZ et f8E5M2FNUZ correspondant aux encodages E4M3 et E5M2 des formats FP8 décrits dans la section Formats numériques 8 bits pour les réseaux de neurones profonds.
  • Type f8E4M3B11FNUZ correspondant à l'encodage E4M3 des formats FP8 décrits dans la section Entraînement et inférence à virgule flottante hybride 8 bits (HFP8) pour les réseaux de neurones profonds.
  • Type de bf16 correspondant au format bfloat16 décrit dans la section BFloat16: le secret des hautes performances sur les Cloud TPU.
  • Types f16, f32 et f64 correspondant respectivement aux formats binary16 ("demi-précision"), binary32 ("précision simple") et binary64 ("double précision") décrits dans la norme IEEE 754.
  • Le type tf32 correspond au format TensorFloat32 et n'est pas entièrement compatible avec StableHLO.
  • Types MX (microscaling) f4E2M1FN, f6E2M3FN, f6E3M2FN et f8E8M0FNU décrits dans la spécification des formats de microscaling OCP.
• Les types complexes représentent des valeurs complexes ayant une partie réelle et une partie imaginaire du même type d'élément. Les types complexes compatibles 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 fonctions représentent à la fois 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 fonction sont de première classe, mais pas dans StableHLO.

StringType ::= 'string'

Le type de chaîne 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 opérations) 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 est fortement inspirée de MLIR, qui n'est pas nécessairement l'alternative la plus ergonomique, mais qui est sans doute la plus adaptée à l'objectif de StableHLO, qui est de créer une interopérabilité accrue entre les frameworks ML et les compilateurs ML.

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

Les opérations StableHLO (également appelées opérations) 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 de toutes les opérations prises en charge.

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 en 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 3 valeurs d'entrée, 2 fonctions d'entrée et 3 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 similaires aux fonctions nommées, à l'exception des points suivants : 1) elles n'ont pas d'identifiant (d'où le nom "anonyme"), 2) elles ne déclarent pas de types de sortie (les types de sortie sont inférés à partir 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" pouvant comporter plusieurs "blocs" d'opérations reliés entre eux via des jump ops. Ces blocs ont des ID qui correspondent à la production Unused afin de pouvoir les distinguer les uns des autres. StableHLO ne comporte pas d'opérations de saut. Par conséquent, la partie correspondante de la syntaxe MLIR n'est pas utilisée (mais elle est toujours là).

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 via 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 découper la valeur d'entrée.

Pour le moment, les programmes StableHLO dans la nature contiennent parfois des attributs qui ne sont pas décrits dans ce document. À l'avenir, nous prévoyons d'intégrer ces attributs dans l'ensemble d'opérations StableHLO ou de les interdire 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 lieu (#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 ->). À strictement parler, les types d'entrée sont redondants, et les types de sortie sont presque toujours redondants également (car pour la plupart des opérations StableHLO, les types de sortie peuvent être déduits des entrées). Néanmoins, la signature d'opération fait délibérément partie de la syntaxe StableHLO pour assurer la compatibilité avec MLIR.

Vous trouverez ci-dessous un exemple d'opération dont l'expression mnémotechnique est select_and_scatter. Elle 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 comportent un littéral et un type qui, ensemble, représentent une valeur StableHLO. En général, le type fait partie de la syntaxe de la constante, sauf lorsqu'il est sans ambiguïté (par exemple, une constante booléenne a un type i1 sans ambiguïté, 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 via des chaînes qui utilisent la notation décimale ou hexadécimale. Les autres bases, telles que les bases binaires ou octales, ne sont pas acceptées. Les constantes entières présentent les 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 une notation décimale ou scientifique. De plus, la notation hexadécimale peut être utilisée pour spécifier directement les bits sous-jacents au format à virgule flottante du type correspondant. Les constantes à virgule flottante présentent les contraintes suivantes:

• (C1) Si une notation non hexadécimale est utilisée, is_wellformed(float_literal, float_type).
• (C2) Si vous utilisez la notation hexadécimale, 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 d'une partie réelle (qui vient en premier) et d'une partie imaginaire (qui vient en deuxième). 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 présentent les 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 de tenseur représentent les valeurs de tenseur à l'aide de listes imbriquées spécifiées via la notation NumPy. Par exemple, dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> représente une valeur de Tensor avec le mappage suivant entre les index et les é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 de tenseur 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 les valeurs de tenseur quantifié à l'aide de la même notation que les constantes de tenseur, avec des éléments spécifiés en tant que constantes de leur type de stockage. Les constantes de tenseur quantifiées présentent les 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 chaînes littérales sont constituées d'octets spécifiés à l'aide de caractères ASCII et de séquences d'échappement. Ils sont indépendants de l'encodage, de sorte que l'interprétation de ces octets est définie par mise en œuvre. Les littéraux de chaîne sont de type string.

Opérations

abs

Sémantique

Effectue une opération abs élément par élément sur le tenseur operand et produit un tenseur result. En fonction du type d'élément, effectue les opérations suivantes :

• Pour les entiers signés : module d'entier.
• Pour les nombres à virgule flottante : abs d'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).
  • baseline_element_type(operand) dans les autres cas.

Exemples

// %operand: [-2, 0, 2]
%result = "stablehlo.abs"(%operand) : (tensor<3xi32>) -> tensor<3xi32>
// %result: [2, 0, 2]

 Autres exemples

add

Sémantique

Effectue l'addition par élément de deux Tensors lhs et rhs, et génère un Tensor result. En fonction du type d'élément, effectue les opérations suivantes :

• Pour les valeurs booléennes : opérateur logique "OR".
• Pour les entiers: addition d'entiers.
• Pour les nombres à virgule flottante : addition d'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 tenseurs non linéarisables :
  • (C1) type(lhs) = type(rhs) = type(result).
• Si l'opération utilise des tenseurs 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) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[6, 8], [10, 12]]

Autres exemples

after_all

Sémantique

Assure que les opérations produisant le inputs sont exécutées avant toutes les opérations qui dépendent 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, !stablehlo.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 tenseurs operands de chaque processus le long de all_gather_dim et produit des tenseurs results.

L'opération divise la grille de processus StableHLO en process_groups, qui est définie 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 l'ensemble des process de 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_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  // channel_id = 0
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
  // use_global_device_ids = false
} : (tensor<2x2xi64>, tensor<2x2xi64>) -> (tensor<2x4xi64>, tensor<2x4xi64>)
// %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 tenseurs operands de chaque processus et produit des tenseurs results.

L'opération divise la grille de processus StableHLO en process_groups, qui est définie 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 schedule où :
  • exec(node) = computation(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value.
• schedule est une arborescence binaire définie par l'implémentation dont le balayage 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(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  // channel_id = 0
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
  // use_global_device_ids = false
} : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>, tensor<4xi64>)
// %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, il divise les valeurs des Tensors operands le long de split_dimension en plusieurs parties, répartit les parties divisées entre les processus, concatène les parties dispersées le long de concat_dimension et produit des Tensors results. L'opération divise la grille de processus StableHLO en process_groups, qui est définie 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_groups = dense<[[0, 1]]> : tensor<1x2xi64>
  // channel_id = 0
} : (tensor<2x4xi64>, tensor<2x4xi64>) -> (tensor<4x2xi64>, tensor<4x2xi64>)
// %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 une opération AND élément par élément sur deux tenseurs lhs et rhs, et produit un tenseur result. En fonction du type d'élément, effectue les opérations suivantes :

• Pour les valeurs booléennes : "AND" logique
• Pour les 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) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[1, 2], [3, 0]]

Autres exemples

atan2

Sémantique

Effectue une opération atan2 élément par élément sur le tenseur lhs et rhs, et génère un tenseur result. En fonction du type d'élément, effectue les opérations suivantes :

• Pour les nombres à virgule flottante : atan2 d'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) : (tensor<3xf64>, tensor<3xf64>) -> tensor<3xf64>
// %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 des 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, exécute 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
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>,
     tensor<2x2x2xf64>) -> (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>)
// %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 des 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
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>) -> tensor<2x2x2xf64>
// %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 tenseur operand pour produire les tenseurs 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
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>) ->
    (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>)
// %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 de cast de bits sur le tenseur operand et produit un tenseur result où les bits de l'ensemble du tenseur operand sont réinterprétés à l'aide du type du tenseur 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), alors bits(result[i0, ..., iR-1]) = bits(operand[i0, ..., iR-1]).

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

Entrées

Sorties

Contraintes

• (C1) Étant donnés 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), alors is_complex(operand) and is_complex(result).

Exemples

// %operand: 0x0123456789ABCDEF
%result = "stablehlo.bitcast_convert"(%operand) : (tensor<f64>) -> tensor<4xf16>
// %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 du 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), à l'exception de quantization_dimension(operand), scales(operand) et zero_points(operand), qui peuvent différer de quantization_dimension(result), scales(result) et zero_points(result), respectivement, dans le cas contraire.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Pour tous les d de 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, alors scales(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_dimensions = array<i64: 2, 1>
} : (tensor<1x3xi32>) -> tensor<2x3x2xi32>
// %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 de branches en fonction de la valeur de index. Plus formellement, 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, %result_branch0) : (tensor<2xi64>, tensor<2xi64>) -> ()
}, {
  "stablehlo.return"(%result_branch1, %result_branch1) : (tensor<2xi64>, tensor<2xi64>) -> ()
}) : (tensor<i32>) -> (tensor<2xi64>, tensor<2xi64>)
// %result0: [1, 1]
// %result1: [1, 1]

 Autres exemples

cbrt

Sémantique

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

• Pour les nombres à virgule flottante : rootn(x, 3) d'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) : (tensor<4xf64>) -> tensor<4xf64>
// %result: [0.0, 1.0, 2.0, 3.0]

 Autres exemples

ceil

Sémantique

Effectue le plafond par élément du tenseur operand et produit un tenseur result. Implémente l'opération roundToIntegralTowardPositive 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) : (tensor<5xf32>) -> tensor<5xf32>
// %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 Cholesky de a[i0, ..., iR-3, :, :], sous la forme d'une matrice triangulaire inférieure (si lower est true) ou triangulaire 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 est indé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
} : (tensor<3x3xf32>) -> tensor<3x3xf64>
// %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 tenseur operand entre une valeur minimale et maximale, et génère un tenseur 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, exécute 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 (numéro 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) : (tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<3xi32>
// %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éfinie comme suit :

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

Ensuite, result@process est donné par :

• operand@process_groups[i, 0] si un i existe de sorte que le processus se trouve dans process_groups[i].
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) dans les autres cas.

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_groups = dense<[[2, 1]]> : tensor<1x2xi64>,
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor1x2xi64>) -> tensor<1x2xi64>
// %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éfinie comme suit :

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

Ensuite, result@process est 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 les autres cas.

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_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>,
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor<2x2xi64>) -> tensor<2x2xi64>
//
// %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 par élément des tenseurs lhs et rhs selon comparison_direction et compare_type, et produit un tenseur 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 méthodes comparison_direction et compare_type fournies. 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 (numéro 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 = #stablehlo<comparison_direction LT>,
  compare_type = #stablehlo<comparison_type FLOAT>
} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xi1>
// %result: [true, false]

 Autres exemples

complexe

Sémantique

Effectue une conversion par élément en une 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) : (tensor<2xf64>, tensor<2xf64>) -> tensor<2xcomplex<f64>>
// %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, prenant inputs et composite_attributes, générant ainsi 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. Si l'intégration de la décomposition ne fournit pas la même sémantique d'opération, privilégiez custom_call.

Le champ version (par défaut 0) permet d'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,
  version = 1 : i32
} : (tensor<f32>, tensor<f32>) -> tensor<f32>

 Autres exemples

concatenate

Sémantique

Chaîne inputs le long de la dimension dimension dans le même ordre que les arguments donnés et produit un tenseur 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, etc. 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
} : (tensor<3x2xi64>, tensor<1x2xi64>) -> tensor<4x2xi64>
// %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"() {
  value = dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32>
} : () -> tensor<2x2xf32>
// %output: [[0.0, 1.0], [2.0, 3.0]]

Autres exemples

d'effectuer une conversion

Sémantique

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

Pour les conversions booléen vers n'importe quel type compatible, la valeur false est convertie en zéro, et la valeur true 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. Pour en savoir plus sur le fonctionnement de cette fonctionnalité pour les types complexes, consultez la section ci-dessous.

Pour les conversions impliquant des nombres entier en entier, entier en virgule flottante ou de type virgule flottante en virgule flottante, si la valeur source peut être exactement représentée dans le type de destination, la valeur du résultat correspond à cette représentation exacte. Sinon, le comportement est à déterminer (numéro 180).

Pour les conversions impliquant une 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 (numéro 180).

Les conversions complexe-complexe suivent le même comportement que les conversions virgule flottante-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 mise à zéro, respectivement. La conversion de la partie réelle suit les conversions à virgule flottante.

En principe, cette opération pourrait exprimer la déquantisation (conversion des tenseurs quantiques en tenseurs réguliers), la quantification (conversion des tenseurs réguliers en tenseurs quantiques) et la requantisation (conversion entre tenseurs quantiques), mais pour le moment, nous disposons d'opérations dédiées à cet effet : uniform_dequantize pour le premier cas d'utilisation et uniform_quantize pour le deuxième et le 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) : (tensor<3xi64>) -> tensor<3xcomplex<f64>>
// %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 produit 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 le recadrage suivant 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).

Ce recadrage 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 hybrides quantifiés, exécute 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) Compte kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension] donné :
  • is_unique(kernel_dimensions).
  • 0 <= kernel_dimensions < N.
• (C19) size(output_spatial_dimensions) = N - 2.
• (C20) Avec 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 tenseurs non linéarisables :
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Si l'opération utilise des tenseurs 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), alors quantization_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), alors is_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_strides = array<i64: 4, 4>,
  padding = dense<0> : tensor<2x2xi64>,
  lhs_dilation = array<i64: 2, 2>,
  rhs_dilation = array<i64: 1, 1>,
  window_reversal = array<i1: false, false>,
  // In the StableHLO dialect, dimension numbers are encoded via:
  // `[<input dimensions>]x[<kernel dimensions>]->[output dimensions]`.
  // "b" is batch dimension, "f" is feature dimension,
  // "i" is input feature dimension, "o" is output feature dimension,
  // "0/1/etc" are spatial dimensions.
  dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
  batch_group_count = 1 : i64,
  feature_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>) -> tensor<1x2x2x1xi64>
// %result: [[
//            [[10], [26]],
//            [[46], [62]]
//          ]]

 Autres exemples

cosinus

Sémantique

Effectue une opération cosinus élément par élément sur le tenseur operand et produit un tenseur result. En fonction du type d'élément, effectue les opérations suivantes :

• Pour les nombres à virgule flottante : cos d'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) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[1.0, 0.0], [-1.0, 0.0]]

 Autres exemples

count_leading_zeros

Sémantique

Effectue un comptage par élément du nombre de bits à zéro au début du tenseur operand et produit un tenseur result.

Entrées

Sorties

Contraintes

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

Exemples

// %operand: [[0, 1], [128, -1]]
%result = "stablehlo.count_leading_zeros"(%operand) : (tensor<2x2xi64>) -> tensor<2x2xi64>
// %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 équivalente dans le compilateur XLA. À l'avenir, nous prévoyons d'unifier ces métadonnées (numéro 741).

Entrées

Sorties

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 = [@foo]
} : (tensor<f64>) -> tensor<f64>

diviser

Sémantique

Effectue une division élément par élément des tenseurs lhs et rhs du dividende et du diviseur, et produit un tenseur result. En fonction du type d'élément, effectue les opérations suivantes :

• Pour les entiers: division entière qui produit le quotient algébrique, en ignorant toute partie fractionnaire.
• Pour les nombres à virgule flottante : division d'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) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
// %result: [5.66666651, -5.66666651, -5.66666651, 5.66666651]

 Autres exemples

dot_general

Sémantique

Calcule les produits scalaires entre des tranches de lhs et des tranches de rhs, et produit un tenseur result.

Plus formellement, 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 de quantification hybride, 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'un des éléments suivants (pour le moment, la sémantique de ces valeurs d'énumération est sous-spécifiée, mais nous prévoyons de résoudre ce problème dans la version #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: calcul le plus lent, mais approximation la plus précise du nombre d'origine.

Un DotAlgorithm définit les principales propriétés de l'algorithme utilisé pour implémenter l'opération de point, qui définit également la précision. Si les champs d'attribut de l'algorithme sont définis, precision_config doit être DEFAULT. DotAlgorithms n'a 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 d'algorithme de point peuvent être définis sur None pour spécifier un algorithme de point vide, qui utilisera plutôt la valeur precision_config.

Les champs DotAlgorithm incluent les suivants :

• lhs_precision_type et rhs_precision_type, précisions auxquelles le côté gauche et le côté droit de l'opération sont arrondis. Les types de précision sont indépendants des types de stockage des entrées et des sorties.
• 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 le membre gauche et/ou le membre droit en plusieurs composants et effectue plusieurs opérations de produit scalaire "primaires" sur ces valeurs, généralement pour émuler une précision plus élevée (par exemple, Utiliser le type de données d'IA bfloat16 pour les calculs à plus haute précision : 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 avec 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}

C'est aux implémentations de décider des combinaisons acceptées. 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 compatible, une erreur doit être générée au lieu de passer à une autre option. La validation StableHLO fournit une validation du meilleur effort, ce qui empêche les algorithmes qui ne sont pas connus pour être compatibles avec aucun matériel.

Consultez xla_data.proto > Algorithm pour connaître certaines valeurs d'algorithme acceptées. La demande n° 2483 décrit le plan visant à créer 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 tenseurs non linéarisables :
  • (C13) element_type(lhs) = element_type(rhs).
• Si l'opération utilise des tenseurs 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 = #stablehlo.dot<
    lhs_batching_dimensions = [0],
    rhs_batching_dimensions = [0],
    lhs_contracting_dimensions = [2],
    rhs_contracting_dimensions = [1]
  >,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision 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
  >
} : (tensor<2x2x2xi64>, tensor<2x2x2xi64>) -> tensor<2x2x2xi64>
// %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 valeur n'est spécifiée, toutes les dimensions sont supposées pouvoir s'étendre.

Entrées

Sorties

Contraintes

• (C1) element_type(result) est donné par :
  • element_type(operand), si !is_per_axis_quantized(operand).
  • element_type(operand), à l'exception de quantization_dimension(operand), scales(operand) et zero_points(operand), qui peuvent différer de quantization_dimension(result), scales(result) et zero_points(result), respectivement, dans le cas contraire.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Pour tous les d de 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 la valeur est dim(operand, quantization_dimension(operand)) = 1, alors scales(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_dimensions = array<i64: 2, 1>,
  known_expanding_dimensions = array<i64: 0>,
  known_nonexpanding_dimensions = array<i64: 1>
} : (tensor<1x3xi64>, tensor<3xi64>) -> tensor<2x3x2xi64>
// %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 de convolution, mais le remplissage est spécifié de manière dynamique via padding.

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) Compte kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension] donné :
  • is_unique(kernel_dimensions).
  • 0 <= kernel_dimensions < N.
• (C19) size(output_spatial_dimensions) = N - 2.
• (C20) Avec 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 tenseurs non linéarisables :
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Si l'opération utilise des tenseurs 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), alors quantization_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), alors is_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]]]
//        ]
// %padding: [[1, 1],
//            [1, 1]]
%result = "stablehlo.dynamic_conv"(%lhs, %rhs, %padding) {
  window_strides = array<i64: 4, 4>,
  lhs_dilation = array<i64: 2, 2>,
  rhs_dilation = array<i64: 1, 1>,
  window_reversal = array<i1: false, false>,
  dimension_numbers = #stablehlo.conv<raw
    input_batch_dimension = 0,
    input_feature_dimension = 3,
    input_spatial_dimensions = [0, 1],
    kernel_input_feature_dimension = 2,
    kernel_output_feature_dimension = 3,
    kernel_spatial_dimensions = [0, 1],
    output_batch_dimension = 0,
    output_feature_dimension = 3,
    output_spatial_dimensions = [1, 2]
  >,
  feature_group_count = 1 : i64,
  batch_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>, tensor<2x2xi64>) -> tensor<1x2x2x1xi64>
// %result: [[
//            [[1], [5]],
//            [[10], [14]]
//          ]]