StableHLO-Spezifikation

StableHLO ist ein Satz von High-Level-Operationen (HLO) in Modellen für maschinelles Lernen (ML). StableHLO dient als Portabilitätsschicht zwischen verschiedenen ML-Frameworks und ML-Compilern: ML-Frameworks, die StableHLO-Programme generieren, sind mit ML-Compilern kompatibel, die StableHLO-Programme verwenden.

Unser Ziel ist es, die ML-Entwicklung zu vereinfachen und zu beschleunigen, indem wir die Interoperabilität zwischen verschiedenen ML-Frameworks (z. B. TensorFlow, JAX und PyTorch) und ML-Compilern (z. B. XLA und IREE) verbessern. Zu diesem Zweck enthält dieses Dokument eine Spezifikation für die Programmiersprache StableHLO.

Diese Spezifikation enthält drei Hauptabschnitte. Im Abschnitt Programme wird zuerst die Struktur von StableHLO-Programmen beschrieben, die aus StableHLO-Funktionen bestehen, die wiederum aus StableHLO-Vorgängen bestehen. Innerhalb dieser Struktur wird im Abschnitt Ops die Semantik der einzelnen Vorgänge angegeben. Der Abschnitt Ausführung enthält die Semantik für alle diese Vorgänge, die gemeinsam in einem Programm ausgeführt werden. Im Abschnitt Notation wird die in der gesamten Spezifikation verwendete Notation erläutert.

Öffnen Sie das Repository für den gewünschten getaggten Release, um die Spezifikation einer früheren StableHLO-Version anzusehen. Beispiel: StableHLO-Version 0.19.0. Informationen zu den Änderungen bei jeder Minor-Version von StableHLO finden Sie im Versionsprotokoll in VhloDialect.td.

Programme

Program ::= {Func}

StableHLO-Programme bestehen aus einer beliebigen Anzahl von StableHLO-Funktionen. Unten sehen Sie ein Beispielprogramm mit einer Funktion @main mit drei Eingaben (%image, %weights und %bias) und einer Ausgabe. Der Hauptteil der Funktion hat 6 Operationen.

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>) -> ()
}

Funktionen

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

StableHLO-Funktionen (auch benannte Funktionen genannt) haben eine Kennung, Eingaben/Ausgaben und einen Body. Künftig planen wir, zusätzliche Metadaten für Funktionen einzuführen, um eine bessere Kompatibilität mit HLO zu erreichen (#425, #626, #740, #744).

IDs

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

StableHLO-Kennungen ähneln den Kennungen in vielen Programmiersprachen, mit zwei Besonderheiten: 1) Alle Kennungen haben Sigils, die verschiedene Arten von Kennzeichnungen unterscheiden, 2) Wertkennungen können vollständig numerisch sein, um die Erstellung von StableHLO-Programmen zu vereinfachen.

Typen

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

StableHLO-Typen werden in Werttypen kategorisiert, die auch First-Class-Typen genannt werden. Sie stellen stabileHLO-Werte dar und Typen ohne Wert, die andere Programmelemente beschreiben. StableHLO-Typen ähneln den Typen vieler Programmiersprachen. Die Hauptbesonderheit ist die domänenspezifische Natur von StableHLO, die zu einigen ungewöhnlichen Ergebnissen führt (z. B. sind Skalartypen keine Werttypen).

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

Tensor-Typen stellen Tensoren dar, also mehrdimensionale Arrays. Sie haben eine Form und einen Elementtyp, wobei eine Form nicht negative oder unbekannte Dimensionsgrößen in aufsteigender Reihenfolge der entsprechenden Abmessungen (auch Achsen genannt) darstellt, die von 0 bis R-1 nummeriert sind. Die Anzahl der Dimensionen R wird als Rang bezeichnet. tensor<2x3xf32> ist beispielsweise ein Tensortyp mit der Form 2x3 und dem Elementtyp f32. Es hat zwei Dimensionen (oder zwei Achsen): die Nullte Dimension und die Erste Dimension, deren Größe 2 und 3 ist. Sein Rang ist 2.

Formen können teilweise oder vollständig unbekannt (dynamisch) sein, z. B. tensor<?x2xf64> teilweise unbekannt und tensor<?x?xf64> vollständig unbekannt. Die Größe dynamischer Dimensionen wird durch ein ? dargestellt. Die Rangfolge von Formen kann nicht aufgehoben werden.

In Zukunft möchten wir Tensortypen über Dimensionsgrößen und Elementtypen hinaus erweitern, z. B. um Layouts (#629) und Sparsity (#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

Quantisierte Elementtypen stellen Ganzzahlwerte eines Speichertyps im Bereich von storage_min bis storage_max (einschließlich) dar, die Gleitkommawerten eines ausgedrückten Typs entsprechen. Für einen gegebenen Ganzzahlwert i kann der entsprechende Gleitkommawert f als f = (i - zero_point) * scale berechnet werden, wobei scale und zero_point als Quantisierungsparameter bezeichnet werden. storage_min und storage_max sind in der Grammatik optional, haben aber die Standardwerte min_value(storage_type) und max_value(storage_type). Quantisierte Elementtypen unterliegen den folgenden Einschränkungen:

• (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) Wenn is_empty(quantization_dimension), dann size(scales) = 1.
• (C11) 0 <= quantization_dimension.

Derzeit ist QuantizationScale eine Gleitkommakonstante. Es besteht jedoch ein großes Interesse an ganzzahlbasierten Skalen, die durch Multiplikatoren und Verschiebungen dargestellt werden. Wir planen, dies in naher Zukunft zu untersuchen (#1404).

Es gibt eine laufende Diskussion über die Semantik von QuantizationZeroPoint, einschließlich des Typs, der Werte und der Frage, ob es in einem quantisierten Tensortyp nur einen oder potenziell mehrere Nullpunkte geben kann. Basierend auf den Ergebnissen dieser Diskussion kann sich die Spezifikation für Nullpunkte in Zukunft ändern (#1405).

Eine weitere laufende Diskussion betrifft die Semantik von QuantizationStorageMin und QuantizationStorageMax, um festzustellen, ob für diese Werte und die Werte quantisierter Tensoren Einschränkungen gelten sollten (#1406).

Außerdem möchten wir die Darstellung unbekannter Skalen und Nullpunkte untersuchen, ähnlich wie wir die Darstellung unbekannter Dimensionsgrößen untersuchen (#1407).

Quantisierte Tensortypen stellen Tensoren mit quantisierten Elementen dar. Diese Tensoren sind genau wie normale Tensoren, mit der Ausnahme, dass ihre Elemente quantisierte Elementtypen anstelle von regulären Elementtypen haben.

Bei quantisierten Tensoren kann die Quantisierung pro Tensor erfolgen, d. h., es gibt ein scale und ein zero_point für den gesamten Tensor. Sie kann aber auch pro Achse erfolgen, d. h., es gibt mehrere scales und zero_points, jeweils ein Paar pro Scheibe einer bestimmten Dimension quantization_dimension. Formal gibt es in einem Tensor t mit Achsenquantisierung dim(t, quantization_dimension)-Segmente des quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :] usw. Alle Elemente im i-ten Segment verwenden scales[i] und zero_points[i] als Quantisierungsparameter. Für quantisierte Tensortypen gelten die folgenden Einschränkungen:

• Für die Quantisierung pro Tensor:
  • Es gelten keine zusätzlichen Einschränkungen.
• Für die Quantisierung pro Achse:
  • (C12) quantization_dimension < rank(self).
  • (C13) dim(self, quantization_dimension) = size(scales).

TokenType ::= 'token'

Tokentypen repräsentieren Tokens, also nicht-transparente Werte, die von einigen Vorgängen erzeugt und verbraucht werden. Mithilfe von Tokens wird die Ausführungsreihenfolge von Vorgängen festgelegt, wie im Abschnitt Ausführung beschrieben.

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

Tupeltypen stellen Tupel dar, also heterogene Listen. Tupel sind eine ältere Funktion, die nur zur Kompatibilität mit HLO vorhanden ist. In HLO werden Tupel verwendet, um variadische Ein- und Ausgaben darzustellen. In StableHLO werden variadische Eingaben und Ausgaben nativ unterstützt. Die einzige Verwendung von Tupeln in StableHLO besteht darin, die HLO-ABI umfassend darzustellen, wobei sich z. B. T, tuple<T> und tuple<tuple<T>> je nach Implementierung erheblich unterscheiden können. Wir planen in Zukunft Änderungen am HLO ABI, die es uns ermöglichen könnten, Tupeltypen aus StableHLO zu entfernen (#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'

Elementtypen stellen Elemente von Tensortypen dar. Im Gegensatz zu vielen Programmiersprachen sind diese Typen in StableHLO keine Erstklassigen. Das bedeutet, dass Werte dieser Typen in StableHLO-Programmen nicht direkt dargestellt werden können. Daher ist es üblich, Skalarwerte vom Typ T mit 0-dimensionalen Tensorwerten vom Typ tensor<T> darzustellen.

• Der Boolesche Typ stellt die booleschen Werte true und false dar.
• Ganzzahltypen können vorzeichenbehaftet (si) oder vorzeichenlos (ui) sein und haben eine der unterstützten Bitbreiten (2, 4, 8, 16, 32 oder 64). Vorzeichenbehaftete siN-Typen stellen Ganzzahlwerte von -2^(N-1) bis 2^(N-1)-1 dar, vorzeichenlose uiN-Typen Ganzzahlwerte von 0 bis 2^N-1.
• Gleitkommatypen können einen der folgenden Werte haben:
  • f8E3M4, f8E4M3 und f8E5M2 sind 8‑Bit-Gleitkommazahlen gemäß IEEE-754-Konventionen.
  • Die Typen f8E4M3FN und f8E5M2 entsprechen den E4M3- und E5M2-Codierungen des FP8-Formats, die in FP8-Formate für Deep Learning beschrieben sind.
  • f8E4M3FNUZ- und f8E5M2FNUZ-Typen, die der Codierung E4M3 und E5M2 der FP8-Formate entsprechen, die unter Numerische 8-Bit-Formate für neuronale Deep-Learning-Netzwerke beschrieben werden.
  • f8E4M3B11FNUZ-Typ, der der E4M3-Codierung der FP8-Formate entspricht, die unter Hybrid-8-Bit-Gleitkommazahl (HFP8) Training und Inferenz für neuronale Deep-Learning-Netzwerke beschrieben werden.
  • bf16-Typ, der dem bfloat16-Format entspricht, das in BFloat16: Das Geheimnis der hohen Leistung auf Cloud TPUs beschrieben wird.
  • Die Typen f16, f32 und f64 entsprechen den Formaten binary16 („halbe Genauigkeit“), binary32 („einfache Genauigkeit“) und binary64 („doppelte Genauigkeit“), die im IEEE 754-Standard beschrieben sind.
  • Der Typ tf32 entspricht dem TensorFloat32-Format und wird in StableHLO nur eingeschränkt unterstützt.
  • f4E2M1FN, f6E2M3FN, f6E3M2FN und f8E8M0FNU MX-Typen (Microscaling) gemäß der OCP-Spezifikation für Microscaling-Formate
• Komplexe Typen stellen komplexe Werte dar, die einen reellen Teil und einen imaginären Teil desselben Elementtyps haben. Unterstützte komplexe Typen sind complex<f32> (beide Teile vom Typ f32) und complex<f64> (beide Teile vom Typ f64).

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

Funktionstypen können sowohl benannte als auch anonyme Funktionen sein. Sie haben Eingabetypen (die Liste der Typen links von ->) und Ausgabetypen (die Liste der Typen rechts von ->). In vielen Programmiersprachen sind Funktionstypen erstklassig, aber nicht in StableHLO.

StringType ::= 'string'

Der Stringtyp steht für Bytefolgen. Anders als in vielen Programmiersprachen ist der Stringtyp in StableHLO nicht erstklassig und wird nur zum Angeben statischer Metadaten für Programmelemente verwendet.

Vorgänge

StableHLO-Vorgänge (auch Ops genannt) stellen eine geschlossene Gruppe von Hochebenen-Vorgängen in Modellen für maschinelles Lernen dar. Wie bereits erwähnt, ist die StableHLO-Syntax stark von MLIR inspiriert. MLIR ist nicht unbedingt die innovativste Alternative, eignet sich aber wohl am besten für das Ziel von StableHLO, mehr Interoperabilität zwischen ML-Frameworks und ML-Compilern zu schaffen.

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

StableHLO-Vorgänge (auch Ops genannt) haben einen Namen, Eingaben/Ausgaben und eine Signatur. Der Name besteht aus dem Präfix stablehlo. und einer Mnemonik, die eine der unterstützten Vorgänge eindeutig identifiziert. Unten finden Sie eine vollständige Liste aller unterstützten Vorgänge.

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

Vorgänge verbrauchen Eingaben und erzeugen Ausgaben. Eingaben werden in Eingabewerte (während der Ausführung berechnet), Eingabefunktionen (statisch bereitgestellt, da Funktionen in StableHLO keine Erstklassigen Werte sind) und Eingabeattribute (auch statisch bereitgestellt) unterteilt. Die Art der Eingaben und Ausgaben, die von einer Operation verbraucht und erzeugt werden, hängt von ihrer Mnemonic ab. Der Operator add benötigt beispielsweise zwei Eingabewerte und liefert einen Ausgabewert. Im Vergleich dazu verbraucht die Operation select_and_scatter 3 Eingabewerte, 2 Eingabefunktionen und 3 Eingabeattribute.

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

Eingabefunktionen (auch anonyme Funktionen genannt) ähneln benannten Funktionen, mit folgenden Ausnahmen: 1) Sie haben keine Kennung (daher der Name „anonym“). 2) Sie deklarieren keine Ausgabetypen. Diese werden aus der return-Operation innerhalb der Funktion abgeleitet.

Die Syntax für Eingabefunktionen enthält einen derzeit nicht verwendeten Teil (siehe Unused-Produktion oben), der aus Gründen der Kompatibilität mit MLIR vorhanden ist. In MLIR gibt es das allgemeinere Konzept von „Regionen“, die mehrere „Blöcke“ von Anweisungen haben können, die über Sprunganweisungen miteinander verbunden sind. Diese Blöcke haben IDs, die der Unused-Produktion entsprechen, damit sie voneinander unterschieden werden können. StableHLO hat keine Sprungbefehle. Daher wird der entsprechende Teil der MLIR-Syntax nicht verwendet, ist aber weiterhin vorhanden.

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

Eingabeattribute haben einen Namen und einen Wert, der eine der unterstützten Konstanten ist. Sie sind die primäre Möglichkeit, statische Metadaten für Programmelemente anzugeben. Bei der Operation concatenate wird beispielsweise das Attribut dimension verwendet, um die Dimension anzugeben, entlang derer die Eingabewerte zusammengefügt werden. Ähnlich werden bei der slice-Operation mehrere Attribute wie start_indices und limit_indices verwendet, um die Grenzen anzugeben, die zum Schneiden des Eingabewerts verwendet werden.

Derzeit enthalten StableHLO-Programme in der Praxis manchmal Attribute, die in diesem Dokument nicht beschrieben werden. Künftig werden wir diese Attribute entweder in den StableHLO-Opset aufnehmen oder ihre Verwendung in StableHLO-Programmen verbieten. In der Zwischenzeit finden Sie hier eine Liste dieser Attribute:

• layout (#629).
• mhlo.frontend_attributes (#628).
• mhlo.sharding (#619)
• output_operand_aliases (#740).
• Standortmetadaten (#594)

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

Die Operatorsignatur besteht aus den Typen aller Eingabewerte (die Liste der Typen auf der linken Seite von ->) und den Typen aller Ausgabewerte (die Liste der Typen auf der rechten Seite von ->). Streng genommen sind die Eingabetypen redundant und die Ausgabetypen fast immer auch, da bei den meisten StableHLO-Operatoren die Ausgabetypen aus den Eingaben abgeleitet werden können. Die Operatorsignatur ist jedoch bewusst Teil der StableHLO-Syntax, um mit MLIR kompatibel zu sein.

Unten sehen Sie ein Beispiel für eine Operation mit der mnemonischen Taste select_and_scatter. Sie nimmt drei Eingabewerte (%operand, %source und %init_value), zwei Eingabefunktionen und drei Eingabeattribute (window_dimensions, window_strides und padding) auf. Beachten Sie, dass die Signatur der Operation nur die Typen der Eingabewerte enthält, aber nicht die Typen der Eingabefunktionen und Attribute, die inline angegeben werden.

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

Konstanten

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

StableHLO-Konstanten haben ein Literal und einen Typ, die zusammen einen StableHLO-Wert darstellen. Generell ist der Typ Teil der konstanten Syntax, es sei denn, er ist eindeutig. Eine boolesche Konstante hat z. B. den Typ i1, während eine Ganzzahlkonstante mehrere mögliche Typen haben kann.

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

Boolesche Konstanten repräsentieren die booleschen Werte true und false. Boolesche Konstanten haben den Typ i1.

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

Ganzzahlkonstanten stellen Ganzzahlwerte über Strings in Dezimal- oder Hexadezimalschreibweise dar. Andere Basen, z. B. binär oder oktal, werden nicht unterstützt. Für Ganzzahlkonstanten gelten die folgenden Einschränkungen:

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

Gleitkommakonstanten stellen Gleitkommawerte über Strings dar, die die Dezimal- oder wissenschaftliche Schreibweise verwenden. Außerdem können Sie mit der Hexadezimalschreibweise die zugrunde liegenden Bits direkt im Gleitkommaformat des entsprechenden Typs angeben. Für Gleitkommakonstanten gelten die folgenden Einschränkungen:

• (C1) Wenn keine hexadezimal Schreibweise verwendet wird, is_wellformed(float_literal, float_type).
• (C2) Bei Verwendung der Hexadezimalschreibweise: size(hexadecimal_digits) = num_bits(float_type) / 4.

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

Komplexe Konstanten stellen komplexe Werte mithilfe von Listen aus einem reellen Teil (an erster Stelle) und einem imaginären Teil (an zweiter Stelle) dar. (1.0, 0.0) : complex<f32> steht beispielsweise für 1.0 + 0.0i und (0.0, 1.0) : complex<f32> für 0.0 + 1.0i. Die Reihenfolge, in der diese Teile dann im Arbeitsspeicher gespeichert werden, ist implementierungsabhängig. Für komplexe Konstanten gelten die folgenden Einschränkungen:

• (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

Tensorkonstanten stellen Tensorwerte mit verschachtelten Listen dar, die über die NumPy-Notation angegeben werden. dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> steht beispielsweise für einen Tensorwert mit der folgenden Zuordnung von Indizes zu Elementen: {0, 0} => 1, {0, 1} => 2, {0, 2} => 3, {1, 0} => 4, {1, 1} => 5, {1, 2} => 6. Die Reihenfolge, in der diese Elemente dann im Arbeitsspeicher gespeichert werden, ist von der Implementierung abhängig. Für Tensorkonstanten gelten die folgenden Einschränkungen:

• (C1) has_syntax(tensor_literal, element_type(tensor_type)), wobei:
  • 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)), wobei gilt:
  • 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:]).
  • Andernfalls false.

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

Quantisierte Tensorkonstanten stellen quantisierte Tensorwerte mit derselben Notation wie Tensorkonstanten, wobei Elemente als Konstanten ihres Speichertyps angegeben sind. Für quantisierte Tensorkonstanten gelten die folgenden Einschränkungen:

• (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))

Stringliterale bestehen aus Bytes, die mit ASCII-Zeichen und Escape-Sequenzen angegeben werden. Sie sind unabhängig von der Codierung, sodass die Interpretation dieser Bytes von der Implementierung abhängt. Stringliterale haben den Typ string.

Operativer Betrieb

abs

Semantik

Führt einen elementweisen Absolutwert-Vorgang auf den operand-Tensor aus und erzeugt einen result-Tensor. Je nach Elementtyp gilt Folgendes:

• Bei vorzeichenbehafteten Ganzzahlen: Ganzzahlmodul.
• Für Gleitkommazahlen: abs von IEEE-754.
• Bei komplexen Zahlen: Komplexmodulus.
• Für quantisierte Typen: dequantize_op_quantize(abs, operand, type(result)).

Eingaben

Ausgaben

Einschränkungen

• (C1) shape(result) = shape(operand).
• (C2) baseline_element_type(result) ist definiert als:
  • complex_element_type(element_type(operand)), wenn is_complex(operand).
  • Andernfalls baseline_element_type(operand).

Beispiele

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

 Weitere Beispiele

Hinzufügen

Semantik

Führt die elementweise Addition von zwei Tensoren lhs und rhs aus und erzeugt einen result-Tensor. Je nach Elementtyp gilt Folgendes:

• Für boolesche Werte: Logisches OR.
• Für Ganzzahlen: Addition von Ganzzahlen.
• Für Gleitkommazahlen: addition von IEEE-754.
• Bei komplexen Zahlen: komplexe Addition.
• Für quantisierte Typen: dequantize_op_quantize(add, lhs, rhs, type(result)).

Eingaben

Ausgaben

Einschränkungen

• Wenn für die Operation nicht quantisierte Tensoren verwendet werden:
  • (C1) type(lhs) = type(rhs) = type(result).
• Wenn für den Vorgang quantisierte Tensoren verwendet werden:
  • (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) Wenn is_per_axis_quantized(lhs), dann quantization_dimension(lhs) = quantization_dimension(result).
  • (C7) Wenn is_per_axis_quantized(rhs), dann quantization_dimension(rhs) = quantization_dimension(result).

Beispiele

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

 Weitere Beispiele

after_all

Semantik

Sorgt dafür, dass die Vorgänge, die die inputs erzeugen, vor Vorgängen ausgeführt werden, die von result abhängig sind. Die Ausführung dieses Vorgangs hat keine Auswirkungen. Er dient nur dazu, Datenabhängigkeiten von result zu inputs herzustellen.

Eingaben

Ausgaben

Beispiele

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

Weitere Beispiele

all_gather

Semantik

Verkettet innerhalb jeder Prozessgruppe im StableHLO-Prozess-Raster die Werte der operands-Tensoren aus jedem Prozess entlang all_gather_dim und erzeugt results-Tensoren.

Dabei wird das StableHLO-Prozess-Raster in process_groups aufgeteilt, was so definiert ist:

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

Gehen Sie anschließend in jedem process_group so vor:

• operands...@receiver = [operand@sender for sender in process_group] für alle receiver in process_group
• results...@process = concatenate(operands...@process, all_gather_dim) für alle process in process_group

Eingaben

Ausgaben

Einschränkungen

• (C1) 0 <= all_gather_dim < rank(operands...).
• (C2) is_unique(replica_groups).
• (C3) size(replica_groups) wird so definiert:
  • num_replicas, wenn cross_replica verwendet wird.
  • num_replicas, wenn cross_replica_and_partition verwendet wird.
  • num_processes, wenn flattened_ids verwendet wird.
• (C4) 0 <= replica_groups < size(replica_groups).
• (C5) Wenn use_global_device_ids = true, dann channel_id > 0.
• (C6) type(results...) = type(operands...), mit folgenden Ausnahmen:
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1).

Beispiele

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

 Weitere Beispiele

all_reduce

Semantik

Innerhalb jeder Prozessgruppe im StableHLO-Prozess-Raster wird eine Reduzierungsfunktion computation auf die Werte der operands-Tensoren aus jedem Prozess angewendet und es werden results-Tensoren erzeugt.

Dabei wird das StableHLO-Prozess-Raster in process_groups aufgeteilt, was so definiert ist:

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

Gehen Sie anschließend in jedem process_group so vor:

• results...@process[result_index] = exec(schedule) für einen Binärbaum schedule wobei:
  • exec(node) = computation(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value.
• schedule ist ein implementierungsdefinierter Binärbaum, dessen sequenzielle Durchlaufreihenfolge to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0])) ist.

Eingaben

Ausgaben

Einschränkungen

• (C1) is_unique(replica_groups).
• (C2) size(replica_groups) ist definiert als:
  • num_replicas, wenn cross_replica verwendet wird.
  • num_replicas, wenn cross_replica_and_partition verwendet wird.
  • num_processes, wenn flattened_ids verwendet wird.
• (C3) 0 <= replica_groups < size(replica_groups).
• (C4) Wenn use_global_device_ids = true, dann channel_id > 0.
• (C5) computation hat den Typ (tensor<E>, tensor<E>) -> (tensor<E>), wobei is_promotable(element_type(operand), E).
• (C6) shape(results...) = shape(operands...).
• (C7) element_type(results...) = E.

Beispiele

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

 Weitere Beispiele

all_to_all

Semantik

Teilt innerhalb jeder Prozessgruppe im StableHLO-Prozess-Raster die Werte der operands-Tensoren entlang von split_dimension in Teile auf, verteilt die aufgeteilten Teile auf die Prozesse, verkettet die verteilten Teile entlang von concat_dimension und erzeugt results-Tensoren. Der Vorgang teilt das StableHLO-Prozessraster in process_groups auf, was so definiert ist:

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

Gehen Sie anschließend in jedem process_group so vor:

• split_parts...@sender = split(operands...@sender, split_count, split_dimension) für alle sender in process_group.
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group], wobeireceiver_index = process_group.index(receiver).
• results...@process = concatenate(scattered_parts...@process, concat_dimension).

Eingaben

Ausgaben

Einschränkungen

• (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) wird so definiert:
  • num_replicas, wenn cross_replica verwendet wird.
  • num_partitions, wenn cross_partition verwendet wird.
• (C7) 0 <= replica_groups < size(replica_groups).
• (C8) dim(replica_groups, 1) = split_count.
• (C9) type(results...) = type(operands...), außer wenn split_dimension != concat_dimension:
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count.
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count.

Beispiele

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

 Weitere Beispiele

und

Semantik

Führt die elementweise UND-Verknüpfung der beiden Tensoren lhs und rhs durch und erzeugt einen result-Tensor. Gehen Sie je nach Elementtyp so vor:

• Für boolesche Werte: Logisches AND.
• Für Ganzzahlen: Bitweises AND.

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

Weitere Beispiele

atan2

Semantik

Führt eine elementweise atan2-Operation auf den Tensoren lhs und rhs aus und erzeugt einen result-Tensor. Je nach Elementtyp gilt Folgendes:

• Für Gleitkommazahlen: atan2 von IEEE-754.
• Für komplexe Zahlen: komplexer Atan2
• Für quantisierte Typen: dequantize_op_quantize(atan2, lhs, rhs, type(result)).

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

Weitere Beispiele

batch_norm_grad

Semantik

Berechnet die Gradienten mehrerer Eingaben von batch_norm_training, die von grad_output zurückgepropagiert werden, und erzeugt grad_operand-, grad_scale- und grad_offset-Tensoren. Formeller ausgedrückt kann dieser Vorgang als Zerlegung in vorhandene StableHLO-Vorgänge mit der folgenden Python-Syntax ausgedrückt werden:

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

Bei quantisierten Typen wird 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)) ausgeführt.

Eingaben

Ausgaben

Einschränkungen

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, mean, variance, grad_output, grad_operand, grad_scale und grad_offset haben dieselbe baseline_element_type.
• (C3) operand, grad_output und grad_operand haben dieselbe Form.
• (C4) scale, mean, variance, grad_scale und grad_offset haben die gleiche Form.
• (C5) size(scale) = dim(operand, feature_index).

Beispiele

// %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

Semantik

Normalisiert den operand-Tensor in allen Dimensionen mit Ausnahme der feature_index-Dimension und erzeugt einen result-Tensor. Formaler kann dieser Vorgang als Zerlegung vorhandener StableHLO-Vorgänge unter Verwendung der Python-Syntax so ausgedrückt werden:

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)

Bei quantisierten Typen wird 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)) ausgeführt.

Eingaben

Ausgaben

Einschränkungen

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, mean, variance und result haben dieselbe 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).

Beispiele

// %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

Semantik

Mittelwert und Varianz werden für alle Dimensionen außer der feature_index-Dimension berechnet. Der operand-Tensor wird normalisiert, wodurch die Tensoren output, batch_mean und batch_var entstehen. Formeller ausgedrückt kann dieser Vorgang als Zerlegung in vorhandene StableHLO-Vorgänge mit der Python-Syntax so ausgedrückt werden:

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

Bei quantisierten Typen wird 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)) ausgeführt.

Eingaben

Ausgaben

Einschränkungen

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, batch_mean, batch_var und output haben dieselbe 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).

Beispiele

// %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

Semantik

Führt eine Bitcast-Operation für den operand-Tensor aus und erzeugt einen result-Tensor, bei dem die Bits des gesamten operand-Tensors mit dem Typ des result-Tensors neu interpretiert werden.

Formeller gesprochener Text bei E = element_type(operand), E' = element_type(result) und R = rank(operand):

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

bits gibt die In-Memory-Darstellung eines bestimmten Werts zurück. Das Verhalten ist implementierungsabhängig, da die genaue Darstellung von Tensoren und Elementtypen implementierungsabhängig ist.

Eingaben

Ausgaben

Einschränkungen

• (C1) Für E = is_quantized(operand) ? storage_type(operand) : element_type(operand), E' = is_quantized(result) ? storage_type(result) : element_type(result) und R = rank(operand):
  • Wenn num_bits(E') = num_bits(E), shape(result) = shape(operand).
  • Wenn num_bits(E') < num_bits(E):
  • rank(result) = R + 1.
  • dim(result, i) = dim(operand, i) für alle 0 <= i < R.
  • dim(result, R) * num_bits(E') = num_bits(E).
  • Wenn num_bits(E') > num_bits(E):
  • rank(result) = R - 1.
  • dim(result, i) = dim(operand, i) für alle 0 <= i < R.
  • dim(operand, R - 1) * num_bits(E) = num_bits(E').
• (C2) Wenn is_complex(operand) or is_complex(result), dannis_complex(operand) and is_complex(result).

Beispiele

// %operand: 0x0123456789ABCDEF
%result = "stablehlo.bitcast_convert"(%operand) : (tensor<f64>) -> tensor<4xf16>
// %result: [0xCDEF, 0x89AB, 0x4567, 0x0123] // little-endian representation

 Weitere Beispiele

broadcast_in_dim

Semantik

Erweitert die Dimensionen und/oder den Rang eines Eingabetensors, indem die Daten im operand-Tensor dupliziert werden. Es wird ein result-Tensor erzeugt. Formal gilt: result[result_index] = operand[operand_index], wo für alle d in axes(operand) Folgendes gilt:

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

Eingaben

Ausgaben

Einschränkungen

• (C1) element_type(result) ist gegeben durch:
  • element_type(operand), wenn !is_per_axis_quantized(operand).
  • element_type(operand), wobei sich quantization_dimension(operand), scales(operand) und zero_points(operand) von quantization_dimension(result), scales(result) und zero_points(result) unterscheiden können.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Für alle d in axes(operand):
  • dim(operand, d) = 1 oder
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Wenn is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Wenn dim(operand, quantization_dimension(operand)) = 1, dann scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))).

Beispiele

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

 Weitere Beispiele

Supportanfrage

Semantik

Erzeugt die Ausgabe aus der Ausführung genau einer Funktion aus branches, je nach Wert von index. Formeller ausgedrückt: result = selected_branch() wobei:

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

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

 Weitere Beispiele

cbrt

Semantik

Führt eine elementweise Kubikwurzel-Operation auf dem operand-Tensor aus und erzeugt einen result-Tensor. Je nach Elementtyp gilt Folgendes:

• Für Gleitkommazahlen: rootn(x, 3) aus IEEE-754.
• Für komplexe Zahlen: komplexe kubische Wurzel.
• Für quantisierte Typen: dequantize_op_quantize(cbrt, operand, type(result))

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

 Weitere Beispiele

Ceil

Semantik

Führt eine elementweise Aufrundung des operand-Tensors aus und erzeugt einen result-Tensor. Implementiert den roundToIntegralTowardPositive-Vorgang aus der IEEE-754-Spezifikation. Bei quantisierten Typen wird dequantize_op_quantize(ceil, operand, type(result)) ausgeführt.

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

 Weitere Beispiele

Cholesky

Semantik

Berechnet die Cholesky-Zerlegung einer Reihe von Matrizen.

Formeller ausgedrückt: Für alle i in index_space(result) ist result[i0, ..., iR-3, :, :] eine Cholesky-Zerlegung von a[i0, ..., iR-3, :, :] in Form einer unter- oder obertriangularen Matrix (wenn lower true oder false ist). Die Ausgabewerte im gegenüberliegenden Dreieck, d. h. im strikten oberen oder unteren Dreieck, sind implementierungsabhängig.

Wenn es i gibt, bei dem die Eingabematrix keine hermitische positiv-definierte Matrix ist, ist das Verhalten nicht definiert.

Bei quantisierten Typen wird dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)) ausgeführt.

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

einschränken

Semantik

Hiermit wird jedes Element des operand-Tensors zwischen einem Mindest- und einem Höchstwert begrenzt und ein result-Tensor erzeugt. Formal gilt: result[result_index] = minimum(maximum(operand[result_index], min_element), max_element), wobei min_element = rank(min) = 0 ? min[] : min[result_index] und max_element = rank(max) = 0 ? max[] : max[result_index] sind. Bei quantisierten Typen wird dequantize_op_quantize(clamp, min, operand, max, type(result)) ausgeführt.

Die Sortierung komplexer Zahlen führt zu überraschenden Ergebnissen. Daher werden wir die Unterstützung für komplexe Zahlen für diesen Vorgang in Zukunft entfernen (#560).

Eingaben

Ausgaben

Einschränkungen

• (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).

Beispiele

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

 Weitere Beispiele

collective_broadcast

Semantik

Senden Sie innerhalb jeder Prozessgruppe im StableHLO-Prozessraster den Wert des Tensors operand vom Quellprozess an die Zielprozesse und erzeugen Sie einen result-Tensor.

Der Vorgang teilt das StableHLO-Prozessraster in process_groups auf, was so definiert ist:

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

Danach ergibt sich result@process aus:

• operand@process_groups[i, 0], wenn es eine i gibt, sodass sich der Prozess in process_groups[i] befindet.
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) Andernfalls

Eingaben

Ausgaben

Einschränkungen

• (C1) is_unique(replica_groups).
• (C2) 0 <= replica_groups < N, wobei N so definiert ist:
  • num_replicas, wenn cross_replica verwendet wird.
  • num_partitions, wenn cross_partition verwendet wird.
• (C3) type(result) = type(operand).

Beispiele

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

Semantik

Sendet innerhalb jeder Prozessgruppe im StableHLO-Prozess-Raster den Wert des operand-Tensors vom Quellprozess an den Zielprozess und erzeugt einen result-Tensor.

Dabei wird das StableHLO-Prozess-Raster in process_groups aufgeteilt, was so definiert ist:

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

Danach ergibt sich result@process aus:

• operand@process_groups[i, 0], wenn ein i vorhanden ist, sodass process_groups[i, 1] = process.
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) Andernfalls

Eingaben

Ausgaben

Einschränkungen

• (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, wobei N so definiert ist:
  • num_replicas, wenn cross_replica verwendet wird.
  • num_partitions, wenn cross_partition verwendet wird.
• (C5) type(result) = type(operand).

Beispiele

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

 Weitere Beispiele

vergleichen

Semantik

Führt einen elementweisen Vergleich der Tensoren lhs und rhs gemäß comparison_direction und compare_type durch und erzeugt einen result-Tensor.

Die Werte von comparison_direction und compare_type haben die folgende Semantik:

Für boolesche und Ganzzahlelementtypen:

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

Für Gleitkommaelementtypen mit compare_type = FLOAT implementiert der Operator die folgenden IEEE-754-Operationen:

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

Für Gleitkomma-Elementtypen mit compare_type = TOTALORDER wird die Kombination aus totalOrder- und compareQuietEqual-Vorgängen aus IEEE-754 verwendet.

Bei komplexen Elementtypen wird der lexikografische Vergleich von (real, imag)-Paaren anhand der angegebenen comparison_direction- und compare_type-Werte durchgeführt. Die Sortierung komplexer Zahlen führt zu überraschenden Ergebnissen. Daher werden wir in Zukunft die Unterstützung für komplexe Zahlen entfernen, wenn comparison_direction = GE, GT, LE oder LT ist (#560).

Für quantisierte Typen wird dequantize_compare(lhs, rhs, comparison_direction) ausgeführt.

Eingaben

Ausgaben

Einschränkungen

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs).
• (C2) shape(lhs) = shape(rhs) = shape(result).
• (C3) compare_type ist definiert als:
  • SIGNED, wenn is_signed_integer(element_type(lhs)).
  • UNSIGNED wenn is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)).
  • FLOAT oder TOTALORDER, wenn is_float(element_type(lhs)).
  • FLOAT, wenn is_complex(element_type(lhs)).

Beispiele

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

 Weitere Beispiele

Komplex

Semantik

Führt eine elementweise Umwandlung in einen komplexen Wert aus einem Paar reeller und imaginärer Werte, lhs und rhs, aus und erzeugt einen result-Tensor.

Eingaben

Ausgaben

Einschränkungen

• (C1) type(lhs) = type(rhs).
• (C2) shape(result) = shape(lhs).
• (C3) element_type(result) hat den Typ complex<E>, wobei E = element_type(lhs).

Beispiele

// %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)]

 Weitere Beispiele

Zusammengesetzt

Semantik

Kapselt einen Vorgang ein, der aus anderen StableHLO-Vorgängen besteht, wobei inputs und composite_attributes als Eingabe und results als Ausgabe dienen. Die Semantik der Operation wird durch das Attribut decomposition implementiert. Die composite-Operation kann durch ihre Dekomposition ersetzt werden, ohne die Programmsemantik zu ändern. Wenn die Einbettung der Dekomposition nicht dieselbe Operatorsemantik bietet, verwenden Sie custom_call.

Das Feld version (Standardwert: 0) gibt an, wann sich die Semantik eines Composites ändert.

Eingaben

Ausgaben

Einschränkungen

• (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)

Beispiele

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

 Weitere Beispiele

concatenate

Semantik

Verkettet inputs entlang der Dimension dimension in derselben Reihenfolge wie die angegebenen Argumente und erzeugt einen result-Tensor. Formeller ausgedrückt: result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1], wobei:

1. id = d0 + ... + dk-1 + kd.
2. d ist gleich dimension und d0 sind die d. Dimensionsgrößen von inputs.

Eingaben

Ausgaben

Einschränkungen

• (C1) same(element_type(inputs...)).
• (C2) same(shape(inputs...)), mit Ausnahme von 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]), mit folgenden Ausnahmen:
  • dim(result, dimension) = dim(inputs[0], dimension) + ....

Beispiele

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

 Weitere Beispiele

Konstante

Semantik

Erzeugt einen output-Tensor aus einer Konstante value.

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

Weitere Beispiele

eine Conversion ausführen

Semantik

Führt eine elementweise Konvertierung von einem Elementtyp in einen anderen mit dem operand-Tensor durch und erzeugt einen result-Tensor.

Bei Conversions vom Typ boolean-to-any-supported-type wird der Wert false in null und der Wert true in eins konvertiert. Bei Konvertierungen vom Typ any-supported-type-to-boolean wird ein Nullwert in false und Werte, die nicht null sind, in true konvertiert. Unten finden Sie weitere Informationen dazu, wie das bei komplexen Typen funktioniert.

Bei Conversions vom Typ Ganzzahl zu Ganzzahl, Ganzzahl zu Gleitkomma oder Gleitkomma zu Gleitkomma wird der Ergebniswert genau so dargestellt, wie der Quellwert im Zieltyp dargestellt werden kann. Andernfalls ist das Verhalten noch nicht festgelegt (#180).

Bei Conversions vom Typ Gleitkomma in Ganzzahl wird der Bruchteil abgeschnitten. Wenn der abgeschnittene Wert nicht im Zieltyp dargestellt werden kann, ist das Verhalten noch nicht festgelegt (#180).

Bei der Konvertierung von komplexen Zahlen in komplexe Zahlen werden reelle und imaginäre Teile wie bei der Konvertierung von Gleitkommazahlen in Gleitkommazahlen konvertiert.

Bei Umwandlungen vom Typ Komplex in einen beliebigen anderen Typ und Beliebiger anderer Typ in komplex wird der Imaginäre Wert der Quelle bzw. des Ziels ignoriert. Die Umwandlung des reellen Teils erfolgt nach den Regeln für Gleitkommaumwandlungen.

Im Prinzip könnte dieser Vorgang die Entquantisierung (Umwandlung von quantisierten Tensoren in reguläre Tensoren), die Quantisierung (Umwandlung von regulären Tensoren in quantisierte Tensoren) und die Neuquantisierung (Umwandlung zwischen quantisierten Tensoren) ausdrücken. Derzeit haben wir jedoch spezielle Vorgänge dafür: uniform_dequantize für den ersten Anwendungsfall und uniform_quantize für den zweiten und dritten Anwendungsfall. In Zukunft werden diese beiden Vorgänge möglicherweise in convert zusammengeführt (#1576).

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

// %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)]

Weitere Beispiele

Faltung

Semantik

Hiermit werden Punktprodukte zwischen Fenstern von lhs und Scheiben von rhs berechnet und result ausgegeben. Das folgende Diagramm zeigt anhand eines konkreten Beispiels, wie Elemente in result aus lhs und rhs berechnet werden.

Formeller ausgedrückt: Betrachten Sie die folgenden Umformulierungen der Eingaben in Bezug auf lhs, um Zeitfenster für lhs angeben zu können:

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

Für diese Neuausrichtung werden die folgenden Hilfsfunktionen verwendet:

• 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], wobei j[d] = i[permutation[d]].

Wenn feature_group_count = 1 und batch_group_count = 1, dann für alle output_spatial_index in index_space(dim(result, output_spatial_dimensions...)) result[result_shape(:, output_spatial_index, :)] = dot_product, wobei:

• 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]). Dieses Feature wird anscheinend noch nicht verwendet, deshalb planen wir, es in Zukunft zu entfernen (#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]).

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

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

Für quantisierte Typen wird 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)) ausgeführt.

Für hybrid quantisierte Typen wird 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) ausgeführt.

Eingaben

Ausgaben

Einschränkungen

• (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) Gegeben ist 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) Gegeben ist kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]:
  • is_unique(kernel_dimensions).
  • 0 <= kernel_dimensions < N.
• (C19) size(output_spatial_dimensions) = N - 2.
• (C20) Gegeben ist 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) wird so definiert:
  • dim(lhs, input_batch_dimension) / batch_group_count, wenn result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension), wenn result_dim = output_feature_dimension.
  • num_windows, andernfalls gilt:
  • 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.
• Wenn für die Operation nicht quantisierte Tensoren verwendet werden:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Wenn für den Vorgang quantisierte Tensoren verwendet werden:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Wenn is_per_axis_quantized(rhs), dann quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Wenn is_per_axis_quantized(result), dann quantization_dimension(result) = output_feature_dimension.
  • Wenn is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Wenn is_per_tensor_quantized(rhs), dann is_per_tensor_quantized(result).
  • Wenn !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Beispiele

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

 Weitere Beispiele

Kosinus

Semantik

Führt eine elementweise Kosinus-Operation auf den Tensor operand aus und erzeugt einen result-Tensor. Je nach Elementtyp gilt Folgendes:

• Für Gleitkommazahlen: cos von IEEE-754.
• Bei komplexen Zahlen: komplexer Kosinus
• Für quantisierte Typen: dequantize_op_quantize(cosine, operand, type(result)).

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

 Weitere Beispiele

leading_zeros_count

Semantik

Führt eine elementweise Zählung der Anzahl der führenden Nullbits im operand-Tensor aus und erzeugt einen result-Tensor.

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

Weitere Beispiele

custom_call

Semantik

Kapselt einen implementierungsdefinierten Vorgang call_target_name ein, der inputs und called_computations als Eingaben nimmt und results als Ausgabe liefert. has_side_effect, backend_config und api_version können verwendet werden, um zusätzliche implementierungsdefinierte Metadaten anzugeben.

Derzeit enthält dieser Vorgang eine ziemlich unorganisierte Sammlung von Metadaten, die die organische Entwicklung des entsprechenden Vorgangs im XLA-Compiler widerspiegelt. Wir planen, diese Metadaten in Zukunft zu vereinheitlichen (#741).

Eingaben

Ausgaben

Beispiele

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

Dividieren

Semantik

Führt eine elementweise Division des Dividenden-Tensors lhs durch den Divisor-Tensor rhs aus und erzeugt einen result-Tensor. Je nach Elementtyp gilt Folgendes:

• Für Ganzzahlen: Ganzzahldivision, die den algebraischen Quotienten mit verworfenen Bruchteilen erzeugt.
• Für Gleitkommazahlen: division von IEEE-754.
• Für komplexe Zahlen: komplexe Division.
• Für quantisierte Typen:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)).

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

Weitere Beispiele

dot_general

Semantik

Berechnet Punktprodukte zwischen Scheiben von lhs und Scheiben von rhs und erzeugt einen result-Tensor.

Formeller result[result_index] = dot_product, wobei Folgendes gilt:

• 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 wobei size(result_batching_index) = size(lhs_batching_dimensions), size(result_lhs_index) = size(lhs_result_dimensions) und 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)).

Für quantisierte Typen wird 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)) ausgeführt.

Führt für hybrid quantisierte Typen 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) aus.

precision_config steuert den Kompromiss zwischen Geschwindigkeit und Genauigkeit bei Berechnungen in Accelerator-Back-Ends. Dies kann einer der folgenden sein (derzeit ist die Semantik dieser Enum-Werte zu niedrig angegeben, aber wir planen, dies in #755 zu beheben):

• DEFAULT: Schnellste Berechnung, aber ungenaueste Näherung an die ursprüngliche Zahl.
• HIGH: langsamere Berechnung, aber genauere Annäherung an die ursprüngliche Zahl.
• HIGHEST: Die langsamste Berechnung, aber die genaueste Annäherung an die ursprüngliche Zahl.

Ein DotAlgorithm definiert die Haupteigenschaften des Algorithmus, der zur Implementierung der Punktoperation verwendet wird, und definiert auch die Genauigkeit. Wenn die Attributfelder des Algorithmus festgelegt sind, muss precision_config DEFAULT sein. DotAlgorithms haben keinen Standardwert, da die Standardparameter von der Implementierung definiert werden. Daher können alle Felder für den Punktalgorithmus auf None gesetzt werden, um einen leeren Punktalgorithmus anzugeben, für den stattdessen der Wert precision_config verwendet wird.

Zu den DotAlgorithm-Feldern gehören:

• lhs_precision_type und rhs_precision_type: die Genauigkeiten, auf die die rechte und linke Seite der Operation gerundet werden. Die Genauigkeitstypen sind unabhängig von den Speichertypen der Ein- und Ausgabe.
• accumulation_type ist die für die Akkumulation verwendete Genauigkeit.
• lhs_component_count, rhs_component_count und num_primitive_operations werden angewendet, wenn ein Algorithmus die linke und/oder rechte Seite in mehrere Komponenten zerlegt und mehrere „primitive“ Punktprodukte auf diese Werte ausführt, in der Regel um eine höhere Genauigkeit zu emulieren (z. B. Der bfloat16-Datentyp für künstliche Intelligenz für Berechnungen mit höherer Genauigkeit: bf16_6x, tf32_3x usw.). Bei Algorithmen ohne Zerlegung sollten diese Werte auf 1 gesetzt werden.
• allow_imprecise_accumulation, um anzugeben, ob für einige Schritte eine Akkumulation mit niedrigerer Genauigkeit zulässig ist (z.B. CUBLASLT_MATMUL_DESC_FAST_ACCUM).

Beispiele für DotAlgorithm-Attribute:

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

Welche Kombinationen unterstützt werden, hängt von der Implementierung ab. Im Allgemeinen kann nicht garantiert werden, dass jeder Algorithmus von jedem Beschleunigertyp vom Nutzer des StableHLO unterstützt wird. Wenn ein bestimmter Algorithmus nicht unterstützt wird, sollte ein Fehler ausgegeben werden, anstatt auf eine Alternative zurückzugreifen. Die StableHLO-Überprüfung bietet eine Best-Effort-Überprüfung und verhindert Algorithmen, die bekanntermaßen auf keiner Hardware unterstützt werden.

Einige unterstützte Algorithmuswerte finden Sie unter xla_data.proto > Algorithm. In Ticket 2483 wird der Plan zum Erstellen einer zentralen Dokumentation zu den unterstützten Algorithmen nach Backend beschrieben.

Eingaben

Ausgaben

Einschränkungen

• (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).
• Wenn für die Operation nicht quantisierte Tensoren verwendet werden:
  • (C13) element_type(lhs) = element_type(rhs).
• Wenn für den Vorgang quantisierte Tensoren verwendet werden:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C15) zero_points(rhs) = 0.
  • (C16) Wenn is_per_axis_quantized(rhs), dann quantization_dimension(rhs) nicht in rhs_contracting_dimensions.
  • Wenn is_quantized(lhs):
  • (C17) storage_type(lhs) = storage_type(rhs).
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C19) Wenn is_per_tensor_quantized(rhs), dannis_per_tensor_quantized(result).
  • Wenn !is_quantized(lhs):
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result).
• Wenn !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.

Beispiele

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

 Weitere Beispiele

dynamic_broadcast_in_dim

Semantik

Dieser Vorgang ist funktional mit der Operation broadcast_in_dim identisch. Die Ergebnisform wird jedoch dynamisch über output_dimensions angegeben.

Der Vorgang akzeptiert auch die optionalen Attribute known_expanding_dimensions und known_nonexpanding_dimensions, um statisches Wissen über das Maximierungsverhalten von Dimensionen auszudrücken. Wenn Sie keine Angabe machen, wird davon ausgegangen, dass alle Dimensionen möglicherweise erweitert werden.

Eingaben

Ausgaben

Einschränkungen

• (C1) element_type(result) ist gegeben durch:
  • element_type(operand), wenn !is_per_axis_quantized(operand).
  • element_type(operand), wobei sich quantization_dimension(operand), scales(operand) und zero_points(operand) von quantization_dimension(result), scales(result) und zero_points(result) unterscheiden können.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Für alle d in axes(operand):
  • dim(operand, d) = 1 oder
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Wenn is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Wenn dim(operand, quantization_dimension(operand)) = 1, dann 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).

Beispiele

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

 Weitere Beispiele

dynamic_conv

Semantik

Funktional ist dieser Vorgang mit dem Vorgang Faltung identisch. Das Padding wird jedoch dynamisch über padding angegeben.

Eingaben

Ausgaben

Einschränkungen

• (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) Gegeben ist 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) Gegeben ist kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]:
  • is_unique(kernel_dimensions).
  • 0 <= kernel_dimensions < N.
• (C19) size(output_spatial_dimensions) = N - 2.
• (C20) Gegeben ist 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) wird so definiert:
  • dim(lhs, input_batch_dimension) / batch_group_count, wenn result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension), wenn result_dim = output_feature_dimension.
  • num_windows, andernfalls gilt:
  • 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.
• Wenn für die Operation nicht quantisierte Tensoren verwendet werden:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Wenn für den Vorgang quantisierte Tensoren verwendet werden:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Wenn is_per_axis_quantized(rhs), dann quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Wenn is_per_axis_quantized(result), dann quantization_dimension(result) = output_feature_dimension.
  • Wenn is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Wenn is_per_tensor_quantized(rhs), dann is_per_tensor_quantized(result).
  • Wenn !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Beispiele

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