Specyfikacja StableHLO

StableHLO to operacja ustawiona dla operacji wysokiego poziomu (HLO) w maszynie (ML). StableHLO działa jako warstwa przenośna między Platformy ML i kompilatory ML: platformy ML służące do tworzenia programów StableHLO są zgodne z kompilatorami ML, które wykorzystują programy StableHLO.

Naszym celem jest uproszczenie i przyspieszenie rozwoju ML przez opracowanie współdziałania różnych platform ML (takich jak TensorFlow, JAX PyTorch) i kompilatorów ML (np. XLA i IREE). W tym celu dokument zawiera specyfikację języka programowania StableHLO.

Specyfikacja składa się z 3 głównych sekcji. Po pierwsze, Sekcja Programy opisuje strukturę programów StableHLO. składające się z funkcji StableHLO, które składają się z operacji StableHLO. W tej strukturze sekcja Operacje określa semantykę poszczególnych operacji. Sekcja Wykonanie zawiera semantykę dla wszystkich i wykonują te operacje w programie. Pamiętaj też, że W sekcji Notation (Zapis) omawiamy sposób zapisu używany w specyfikacji.

Aby wyświetlić specyfikację z poprzedniej wersji StableHLO, otwórz repozytorium w otagowane udostępnienie. Na przykład specyfikacja produktu StableHLO w wersji 0.19.0. Aby wyświetlić zmiany, które wystąpiły przy każdym dodatkowym poziomie wersji StableHLO, zapoznaj się z artykułem log wersji w VhloDialect.td.

Programy

Program ::= {Func}

Programy StableHLO składają się z dowolnej liczby funkcji StableHLO. Poniżej znajdziesz przykładowy program z funkcją @main, która ma 3 dane wejściowe (%image, %weights i %bias) oraz 1 wynik. Treść funkcji ma 6 operacji.

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

Funkcje

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

Funkcje stabilneHLO (nazywane też funkcjami nazwanymi) mają identyfikatora, danych wejściowych/wyjściowych i treści. W przyszłości planujemy wprowadzenie dodatkowych metadanych funkcji, aby uzyskać lepszą zgodność z HLO (#425, #626, #740, 744).

Identyfikatory

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

Identyfikatory StableHLO są podobne do identyfikatorów w wielu programach mają dwie cechy: 1) wszystkie identyfikatory mają sigile, rozróżniają różne typy identyfikatorów, 2) identyfikatory wartości mogą być całkowicie numeryczne, aby uprościć generowanie programów StableHLO.

Typy

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

Typy StableHLO są podzielone na typy wartości (nazywane też typy pierwszej klasy), które reprezentują wartości StableHLO i typy niewartościowe; które opisują inne elementy programu. Typy StableHLO są podobne do typów w w wielu językach programowania, z których charakterystyczną jest specyficzny dla domeny, co skutkuje nietypowymi wynikami (np.typów skalarnych). nie są typami wartości).

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

Typy Tensor reprezentują tensory, czyli tablice wielowymiarowe. Mają kształt oraz typ elementu, gdzie kształt jest liczbą nieujemną lub nieznanych rozmiarów w kolejności rosnącej do odpowiednich wymiary (nazywane również osiami) numerowane od 0 do R-1. liczba wymiarów R jest nazywana rankingiem. Przykład: tensor<2x3xf32> to typ tensora o kształcie 2x3 i typie elementu f32. Ma 2 wymiary (lub, innymi słowy, dwie osie) – wymiar 0 i pierwszy wymiar, których rozmiary to 2 i 3. Jego pozycja to 2.

Kształty mogą być częściowo lub całkowicie nieznane (dynamiczne), np. tensor<?x2xf64> jest częściowo nieznana, a tensor<?x?xf64> jest całkowicie nieznana. Dynamiczna, rozmiary są podawane za pomocą ?. Kształtów nie można uszeregować w kolejności.

Planujemy w przyszłości rozszerzyć zakres typów tensorów rozmiarów i typów elementów, np. aby uwzględnić układy (#629) i rzadkie (#1078).

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

Typy elementów kwantowych reprezentują wartości całkowite typu pamięci masowej w zakres od storage_min do storage_max (włącznie), który odpowiada wartości zmiennoprzecinkowych wyrażenia typu. Dla danej liczby całkowitej: i, odpowiednia wartość zmiennoprzecinkowa f może zostać obliczona jako f = (i - zero_point) * scale, gdzie nazwy scale i zero_point są nazywane parametry kwantyzacji. storage_min i storage_max są opcjonalne w gramatyce, ale mają wartości domyślne min_value(storage_type) oraz max_value(storage_type). Typy elementów poddanych kwantyzacji mają te ograniczenia:

• (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) Jeśli is_empty(quantization_dimension), to size(scales) = 1.
• (C11) 0 <= quantization_dimension.

Obecnie QuantizationScale jest stałą zmiennoprzecinkową, ale jest duże zainteresowanie skalami opartymi na liczbach całkowitych, reprezentowanych za pomocą mnożników i zmiany. W najbliższej przyszłości planujemy to omówić. (nr 1404).

Trwa dyskusja na temat semantyki elementu QuantizationZeroPoint. typu, wartości i tego, czy może być tylko jeden potencjalnie wiele punktów zerowych w kwantowym typie tensora. Na podstawie tej dyskusji, specyfikacja dotycząca zerowej liczby punktów może ulec zmianie w przyszłości (#1405).

Kolejna trwająca dyskusja dotyczy semantyki kategorii QuantizationStorageMin i QuantizationStorageMax, aby określić, czy jakieś ograniczenia powinny być nałożone na te wartości i wartości tensorów kwantowych (#1406).

Na koniec planujemy zbadać reprezentowanie nieznanych skal i zera podobnie jak chcemy badać reprezentowanie nieznanych rozmiarów (#1407).

Kwantowe typy tensorów reprezentują tensory z elementami kwantyzowanymi. Te tensory są dokładnie takie same jak tensory regularne, z tym że ich elementy mają kwantyzowane typy elementów, a nie zwykłe.

W tensorach kwantyzowanych kwantyzacja może być pertensora, co oznacza, że jeden scale i zero_point dla całego tensora lub mogą być na oś, oznacza, że mamy wiele scales i zero_points, po jednej parze na wycinek danego wymiaru: quantization_dimension. Bardziej oficjalnie, w tensorze t z kwantyzacją na oś, powstaje dim(t, quantization_dimension) wycinków z quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :], itd. Wszystkie elementy w i. wycinku używają scales[i] i zero_points[i] jako ich parametry kwantyzacji. Kwantyzowane typy tensorów mają następujące właściwości ograniczenia:

• W przypadku kwantyzacji na centrów:
  • Brak dodatkowych ograniczeń.
• Do kwantyzacji na oś:
  • (C12) quantization_dimension < rank(self).
  • (C13) dim(self, quantization_dimension) = size(scales).

TokenType ::= 'token'

Typy tokenów reprezentują tokeny, tj.tworzone i wykorzystywane wartości nieprzejrzyste. niektóre operacje. Tokeny służą do narzucania kolejności wykonywania operacji zgodnie z opisem w sekcji Wykonanie.

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

Typy krotek reprezentują krotki, czyli listy heterogeniczne. Kropki to starsza wersja funkcję, która istnieje tylko na potrzeby zgodności z HLO. W HLO krotki to: używane do reprezentowania zmiennych danych wejściowych i wyjściowych. W StableHLO zmienne wejściowe i dane wyjściowe są obsługiwane natywnie, a jedynym wykorzystaniem krotek w StableHLO jest kompleksowo odzwierciedla wartość HLO ABI, gdzie np. T, tuple<T> i Wartość tuple<tuple<T>> może się znacznie różnić w zależności od implementacji. Planujemy w przyszłości wprowadzić zmiany w interfejsie HLO ABI. co może umożliwić nam usunięcie typów krotek ze StableHLO (#598).

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

Typy elementów reprezentują elementy typów tensorów. Inaczej niż w przypadku wielu programów W standardzie StableHLO te typy nie są najwyższej klasy. Oznacza to, że Programy StableHLO nie mogą bezpośrednio reprezentować wartości tego typu (w efekcie jest idiomatyczne reprezentuje wartości skalarne typu T za pomocą tensora 0-wymiarowego wartości typu tensor<T>).

• Typ wartości logicznej reprezentuje wartości logiczne true i false.
• Liczba całkowita może być typu ze znakiem (si) lub bez znaku (ui) i zawierać jedną z obsługiwanych szerokości bitów (2, 4, 8, 16, 32 lub 64). Typy siN z podpisem reprezentują wartości całkowite od -2^(N-1) do 2^(N-1)-1 Typy obejmujące i bez znaku uiN reprezentują wartości całkowite od 0 do 2^N-1 włącznie.
• Oto kilka typów liczb zmiennoprzecinkowych:
  • f8E4M3FN i f8E5M2 odpowiadające odpowiednio Kodowanie E4M3 i E5M2 formatu FP8 opisanego w Formaty FP8 do deep learningu.
  • Typy f8E4M3FNUZ i f8E5M2FNUZ odpowiadające parametrom E4M3 i E5M2 kodowania formatów FP8 opisanych w 8-bitowe formaty numeryczne dla głębokich sieci neuronowych.
  • Typ f8E4M3B11FNUZ odpowiadający kodowaniu E4M3 formatów FP8 opisane w HFP8 (HFP8) – trenowanie i wnioskowanie w przypadku głębokich sieci neuronowych.
  • Typ bf16 odpowiadający formatowi bfloat16 opisanemu w BFloat16: tajemnica wysokiej wydajności w Cloud TPU.
  • Typy f16, f32 i f64 odpowiadające odpowiednio binary16 („połowa precyzji”), binary32 („pojedyncza precyzja”) i Formaty binary64 („podwójna precyzja”) opisane w standard IEEE 754.
  • Typ tf32 odpowiada formatowi TensorFloat32 i ma ograniczoną obsługę protokołu StableHLO.
• Typy złożone reprezentują złożone wartości, które mają część rzeczywistą. i część urojona tego samego typu elementu. Obsługiwany kompleks typy to complex<f32> (obie części są typu f32) oraz complex<f64> (obie części są typu f64).

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

Typy funkcji reprezentują zarówno funkcje nazwane, jak i anonimowe. Mają typy danych wejściowych (lista typów po lewej stronie panelu ->) i typy danych wyjściowych (lista typów po prawej stronie ->). W wielu programach w językach, typy funkcji są pierwszej klasy, ale nie w StableHLO.

StringType ::= 'string'

Typ ciągu reprezentuje sekwencje bajtów. Inaczej niż w przypadku wielu programów typ ciągu znaków nie jest pierwszą klasą w StableHLO i jest używany określać statyczne metadane elementów programu.

Operacje

Operacje stabilne HLO (nazywane też operacjami) reprezentują zbiór zamknięty operacji wysokiego poziomu w modelach systemów uczących się. Jak omówiliśmy powyżej, Składnia StableHLO jest mocno inspirowana modelem MLIR, który nie musi jest ergonomiczną alternatywą, ale zapewne najlepiej pasuje do celu StableHLO zapewniając większą interoperacyjność między platformami ML i kompilatorami ML.

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

Operacje stabilne HLO (nazywane też operacjami) mają nazwę, danych wejściowych i wyjściowych oraz podpis. Nazwa składa się z prefiksu stablehlo. i mnemotechnika, która jednoznacznie identyfikuje jedno z obsługiwanych działań. Zobacz poniżej pełną listę obsługiwanych operacji.

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

Operacje wykorzystują dane wejściowe i generują dane wyjściowe. Dane wejściowe są klasyfikowane jako wartości wejściowe (obliczane podczas wykonywania), funkcje wejściowe (podawane statycznie, ponieważ w StableHLO funkcje nie są wartościami pierwszej klasy) oraz atrybutów wejściowych (również statycznych). rodzaj danych wejściowych i wyjściowych, zależy od mnemotechniki, ale też od jej mnemotechniki. Na przykład add funkcja wykorzystuje 2 wartości wejściowe i generuje 1 wartość wyjściową. Dla porównania Operacje select_and_scatter używają 3 wartości wejściowych, 2 funkcji wejściowych i 3 atrybuty wejściowe.

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

Funkcje wejściowe (nazywane też funkcjami anonimowymi) są bardzo są podobne do funkcji nazwanych, z tą różnicą, że: 1) nie mają identyfikatora (więc nazwa „anonimowy”), 2) nie deklarują typów danych wyjściowych (typy danych wyjściowych to wywnioskowane z operacji return w funkcji).

Składnia funkcji wejściowych zawiera obecnie nieużywaną część (zobacz Unused powyżej), który zapewnia zgodność z MLIR. W kontekście MLIR istnieje bardziej ogólne pojęcie „regionów”. które mogą zawierać wiele „bloków” operacji połączonych za pomocą skoków. Te bloki mają identyfikatory, które odpowiadają do środowiska produkcyjnego Unused, aby można je było odróżnić. W StableHLO nie są wykonywane operacje przeskakiwania, więc odpowiadająca mu część składni MLIR to nieużywane (ale nadal jest dostępne).

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

Atrybuty wejściowe mają nazwę i wartość, która jest jedną z obsługiwanych stałe. Stanowią podstawowy sposób określania statycznych metadanych programu . Na przykład operacja concatenate używa atrybutu dimension do pozwala określić wymiar, z którym są połączone jego wartości wejściowe. Podobnie operacja slice wykorzystuje wiele atrybutów, takich jak start_indices i limit_indices aby określić granice używane do wycinania wartości wejściowej.

Obecnie programy StableHLO działające na wolności czasami zawierają atrybuty które nie zostały opisane w tym dokumencie. W przyszłości planujemy wchłania te atrybuty do Opsetu StableHLO lub zabrania ich które pojawiają się w programach StableHLO. W międzyczasie możesz zapoznać się z listą tych kategorii atrybuty:

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

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

Podpis operacji składa się z typów wszystkich wartości wejściowych (listy typów w po lewej stronie elementu ->) oraz typy wszystkich wartości wyjściowych (lista po prawej stronie ->). Ogólnie typy danych wejściowych są nadmiarowe, a typy danych wyjściowych niemal zawsze są nadmiarowe (ponieważ w przypadku w przypadku większości operacji StableHLO typy danych wyjściowych można wywnioskować na podstawie danych wejściowych). Mimo to jest celowo częścią składni StableHLO, która zapewnia zgodność z MLIR.

Poniżej znajdziesz przykład operacji, której mnemotechnika to select_and_scatter. zużywa 3 elementy, wartości wejściowe (%operand, %source i %init_value), 2 funkcje wejściowe i 3 atrybuty wejściowe (window_dimensions, window_strides i padding). Zwróć uwagę, że podpis operacji obejmuje tylko typy wartości wejściowych (ale nie o typach funkcji i atrybutów wejściowych dostępnych w tekście).

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

Stałe

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

stałe StabilHLO mają literał i typ, które razem reprezentują wartość StableHLO. Zasadniczo typ jest częścią składni stałej. Wyjątkiem gdy jest jednoznaczne (np. stała logiczna jednoznacznie ma typ i1, a stała całkowita może mieć kilka typów).

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

Stałe wartości logiczne reprezentują wartości logiczne true i false. Wartość logiczna stałe mają typ i1.

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

Stałe liczby całkowite reprezentują wartości całkowite za pomocą ciągów znaków z ułamkami dziesiętnymi lub w notacji szesnastkowej. Inne podstawy, np. binarnych lub ósemkowych, nie są obsługiwane. Stałe liczby całkowite mają następujące ograniczenia:

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

Stałe zmiennoprzecinkowe odzwierciedlają wartości zmiennoprzecinkowe za pomocą ciągów, które użyj zapisu dziesiętnego lub naukowego. Dodatkowo przy użyciu notacji szesnastkowej można służy do bezpośredniego określania bazowych bitów w formacie zmiennoprzecinkowym odpowiadający jej typ. Stałe zmiennoprzecinkowe mają te ograniczenia:

• (C1) Jeśli używany jest zapis w notacji nieszesnastkowej, is_wellformed(float_literal, float_type)
• (C2) Jeśli używany jest zapis szesnastkowy, size(hexadecimal_digits) = num_bits(float_type) / 4

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

Stałe złożone przedstawiają wartości złożone za pomocą list części rzeczywistych (pierwsza jest część) i część urojona (druga). Przykład: (1.0, 0.0) : complex<f32> oznacza 1.0 + 0.0i, a (0.0, 1.0) : complex<f32> oznacza: 0.0 + 1.0i. Kolejność, są zapisywane w pamięci. Stałe złożone mają te ograniczenia:

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

Stałe Tensor przedstawiają wartości tensorów za pomocą zagnieżdżonych list określonych za pomocą funkcji Zapis NumPy. Na przykład: dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> reprezentuje wartość tensora z następującym mapowaniem indeksów na elementy: {0, 0} => 1, {0, 1} => 2, {0, 2} => 3, {1, 0} => 4, {1, 1} => 5, {1, 2} => 6 Kolejność, w jakiej te elementy są zapisywane w pamięci, to jest bardzo skomplikowana. Stałe tensora mają następujące ograniczenia:

• (C1) has_syntax(tensor_literal, element_type(tensor_type)), gdzie:
  • 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)), gdzie:
  • 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:]).
  • w przeciwnym razie: false.

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

Kwantowe stałe tensora reprezentują skwantyzowane wartości tensora z wykorzystaniem tego samego w postaci stałych tensorów, z elementami określonymi jako stałe ich wartości typu pamięci masowej. Poddane kwantyzacji stałe tensora mają następujące ograniczenia:

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

Literały łańcuchowe składają się z bajtów określonych za pomocą znaków ASCII oraz sekwencje ucieczki. Nie wymagają one kodowania, więc interpretacja tych wartości Bajty są definiowane przez implementację. Literały łańcuchowe mają typ string.

Operacje

brzuch

Semantyka

Wykonuje operację absolutnego punktu widzenia na podstawie elementu na tensorze operand i tworzy result tensora. W zależności od typu elementu wykonuje te działania:

• W przypadku liczb całkowitych ze znakiem: moduł liczby całkowitej.
• W przypadku jednostek zmiennoprzecinkowych: abs w standardzie IEEE-754.
• W przypadku liczb zespolonych: moduł zespolony.
• W przypadku typów kwantowych: dequantize_op_quantize(abs, operand, type(result)).

Dane wejściowe

Wyniki

Ograniczenia

• (C1) shape(result) = shape(operand).
• (C2) Parametr baseline_element_type(result) jest zdefiniowany jako:
  • complex_element_type(element_type(operand)), jeśli is_complex(operand).
  • W przeciwnym razie: baseline_element_type(operand).

Przykłady

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

Więcej przykładów

dodaj

Semantyka

Wykonuje dodawanie elementów z uwzględnieniem 2 tensorów lhs i rhs, tworząc Tensor result. W zależności od typu elementu wykonuje te działania:

• W przypadku wartości logicznych: logiczny LUB.
• W przypadku liczb całkowitych: dodawanie liczby całkowitej.
• W przypadku jednostek zmiennoprzecinkowych: addition w standardzie IEEE-754.
• W przypadku liczb zespolonych: dodawanie zespolone.
• W przypadku typów kwantowych: dequantize_op_quantize(add, lhs, rhs, type(result)).

Dane wejściowe

Wyniki

Ograniczenia

• Jeśli operacja używa tensorów niekwantowych:
  • (C1) type(lhs) = type(rhs) = type(result).
• Jeśli operacja używa tensorów kwantyzowanych:
  • (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) Jeśli is_per_axis_quantized(lhs), to quantization_dimension(lhs) = quantization_dimension(result).
  • (C7) Jeśli is_per_axis_quantized(rhs), to quantization_dimension(rhs) = quantization_dimension(result).

Przykłady

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

Więcej przykładów

after_all

Semantyka

Zapewnia, że operacje tworzące interfejs inputs są wykonywane przed operacji zależnych od result. Wykonanie tej operacji nie przyniesie żadnego efektu. istnieje tylko po to, aby ustalić zależności danych od result do inputs.

Dane wejściowe

Wyniki

Przykłady

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

Więcej przykładów

all_gather

Semantyka

W każdej grupie procesów w siatce procesów StableHLO łączy wartości tensorów operands z każdego procesu w all_gather_dim i daje wynik Tensory: results.

Operacja dzieli siatkę procesów StableHLO na siatkę procesów process_groups, która jest: zdefiniowane w następujący sposób:

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

Następnie w każdym elemencie process_group:

• operands...@receiver = [operand@sender for sender in process_group] za wszystko receiver w: process_group.
• results...@process = concatenate(operands...@process, all_gather_dim) za wszystko process w: process_group.

Dane wejściowe

Wyniki

Ograniczenia

• (C1) 0 <= all_gather_dim < rank(operands...).
• (C2) is_unique(replica_groups).
• (C3) Parametr size(replica_groups) jest zdefiniowany jako:
  • num_replicas, jeśli używana jest właściwość cross_replica.
  • num_replicas, jeśli używana jest właściwość cross_replica_and_partition.
  • num_processes, jeśli używana jest właściwość flattened_ids.
• (C4) 0 <= replica_groups < size(replica_groups).
• (C5) Jeśli use_global_device_ids = true, to channel_id > 0.
• (C6) type(results...) = type(operands...) z wyjątkiem:
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1).

Przykłady

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

Więcej przykładów

all_reduce

Semantyka

W każdej grupie procesów w siatce procesów StableHLO stosuje redukcję funkcji computation na wartości tensorów operands z każdego procesu i generuje tensory results.

Operacja dzieli siatkę procesów StableHLO na siatkę procesów process_groups, która jest: zdefiniowane w następujący sposób:

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

Następnie w każdym elemencie process_group:

• results...@process[result_index] = exec(schedule) dla pewnego drzewa binarnego schedule, gdzie:
  • exec(node) = computation(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value.
• schedule to zdefiniowane w implementacji drzewo binarne, w którym przemierzanie wynosi to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0])).

Dane wejściowe

Wyniki

Ograniczenia

• (C1) is_unique(replica_groups).
• (C2) Parametr size(replica_groups) jest zdefiniowany jako:
  • num_replicas, jeśli używana jest właściwość cross_replica.
  • num_replicas, jeśli używana jest właściwość cross_replica_and_partition.
  • num_processes, jeśli używana jest właściwość flattened_ids.
• (C3) 0 <= replica_groups < size(replica_groups).
• (C4) Jeśli use_global_device_ids = true, to channel_id > 0.
• (C5) computation ma typ (tensor<E>, tensor<E>) -> (tensor<E>), gdzie: is_promotable(element_type(operand), E)
• (C6) shape(results...) = shape(operands...).
• (C7) element_type(results...) = E.

Przykłady

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

Więcej przykładów

all_to_all

Semantyka

W każdej grupie procesów w siatce procesów StableHLO dzieli wartości tensory operands wzdłuż split_dimension na części, rozkłada podział łączą różne elementy, łączą w sobie rozproszone części concat_dimension i tworzy tensory results. Operacja dzieli siatkę procesów StableHLO na siatkę procesów process_groups, która jest: zdefiniowane w następujący sposób:

• cross_replica(replica_groups), jeśli channel_id <= 0.
• cross_partition(replica_groups), jeśli channel_id > 0.

Następnie w każdym elemencie process_group:

• split_parts...@sender = split(operands...@sender, split_count, split_dimension) za wszystkie sender w process_group.
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group], gdzie receiver_index = process_group.index(receiver)
• results...@process = concatenate(scattered_parts...@process, concat_dimension).

Dane wejściowe

Wyniki

Ograniczenia

• (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) Parametr size(replica_groups) jest zdefiniowany jako:
  • num_replicas, jeśli używana jest właściwość cross_replica.
  • num_partitions, jeśli używana jest właściwość cross_partition.
• (C7) 0 <= replica_groups < size(replica_groups).
• (C8) dim(replica_groups, 1) = split_count.
• (C9) type(results...) = type(operands...) oprócz tych, które split_dimension != concat_dimension:
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count.
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count.

Przykłady

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

Więcej przykładów

i

Semantyka

Wykonuje z punktu widzenia elementu ORAZ dwa tensory lhs i rhs, generuje result tensora. W zależności od typu elementu wykonuje te działania:

• W przypadku wartości logicznych: logiczny ORAZ.
• W przypadku liczb całkowitych: bitowe ORAZ.

Dane wejściowe

Wyniki

Ograniczenia

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

Przykłady

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

Więcej przykładów

atan2

Semantyka

Wykonuje operacje atan2 dotyczące elementów na tensorze lhs i rhs, tworząc Tensor result. W zależności od typu elementu wykonuje te działania:

• W przypadku jednostek zmiennoprzecinkowych: atan2 w standardzie IEEE-754.
• W przypadku liczb zespolonych: atan2 zespolona.
• W przypadku typów kwantowych: dequantize_op_quantize(atan2, lhs, rhs, type(result)).

Dane wejściowe

Wyniki

Ograniczenia

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

Przykłady

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

Więcej przykładów

batch_norm_grad

Semantyka

Oblicza gradienty kilku danych wejściowych z propagacją wsteczną batch_norm_training od grad_output i tworzy grad_operand, grad_scale i grad_offset tensorów. Bardziej oficjalnie operacja ta można wyrazić jako istniejących operacji StableHLO przy użyciu składni Pythona w następujący sposób:

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

W przypadku typów kwantowych 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))

Dane wejściowe

Wyniki

Ograniczenia

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, mean, variance, grad_output, grad_operand, grad_scale i grad_offset mają takie same wartości baseline_element_type.
• (C3) operand, grad_output i grad_operand mają taki sam kształt.
• (C4) scale, mean, variance, grad_scale i grad_offset mają taki sam kształt.
• (C5) size(scale) = dim(operand, feature_index).

Przykłady

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

Semantyka

Normalizuje tensor operand we wszystkich wymiarach oprócz feature_index i generuje tensor result. Bardziej formalnie można wyrazić jako dekompozycję na istniejące operacje StableHLO używając składni Pythona w ten sposób:

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)

W przypadku typów kwantowych 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))

Dane wejściowe

Wyniki

Ograniczenia

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, mean, variance i result mają ten sam 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).

Przykłady

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

Semantyka

Oblicza średnią i wariancję dla wszystkich wymiarów oprócz feature_index i normalizuje tensor operand tworzący output, batch_mean i batch_var tensorów. Bardziej oficjalnie operacja ta może zostać przedstawiona jako do istniejących operacji StableHLO przy użyciu składni Pythona jako następujące:

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

W przypadku typów kwantowych 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))

Dane wejściowe

Wyniki

Ograniczenia

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, batch_mean, batch_var i output – ten sam 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).

Przykłady

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

Semantyka

Wykonuje operację bitcastową na tensorze operand i tworzy tensor result gdzie bity całego tensora operand są ponownie interpretowane przy użyciu funkcji typ tensora result.

Bardziej oficjalnie, biorąc pod uwagę E = element_type(operand), E' = element_type(result), i R = rank(operand):

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

Funkcja bits zwraca określoną wartość w pamięci i jej działanie jest zdefiniowana implementacją, ponieważ dokładne reprezentowanie tensorów jest są zdefiniowane, a dokładna reprezentacja typów elementów jest a także kilka sposobów wdrożenia.

Dane wejściowe

Wyniki

Ograniczenia

• (C1) Dane E = is_quantized(operand) ? storage_type(operand) : element_type(operand), E' = is_quantized(result) ? storage_type(result) : element_type(result) i R = rank(operand):
  .
  • Jeśli num_bits(E') = num_bits(E), shape(result) = shape(operand).
  • Jeśli num_bits(E') < num_bits(E):
  • rank(result) = R + 1.
  • dim(result, i) = dim(operand, i) za wszystkie 0 <= i < R.
  • dim(result, R) * num_bits(E') = num_bits(E).
  • Jeśli num_bits(E') > num_bits(E):
  • rank(result) = R - 1.
  • dim(result, i) = dim(operand, i) za wszystkie 0 <= i < R.
  • dim(operand, R - 1) * num_bits(E) = num_bits(E').
• (C2) Jeśli is_complex(operand) or is_complex(result), to is_complex(operand) and is_complex(result)

Przykłady

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

Więcej przykładów

broadcast_in_dim

Semantyka

Rozszerza wymiary lub pozycję tensora wejściowego przez zduplikowanie danych w tensorze operand i tworzy tensor result. Bardziej oficjalnie, result[result_index] = operand[operand_index], gdzie dla wszystkich d w axes(operand):

• operand_index[d] = 0, jeśli dim(operand, d) = 1.
• W przeciwnym razie: operand_index[d] = result_index[broadcast_dimensions[d]].

Dane wejściowe

Wyniki

Ograniczenia

• (C1) element_type(result) otrzymuje:
  • element_type(operand), jeśli !is_per_axis_quantized(operand).
  • element_type(operand) oprócz quantization_dimension(operand), scales(operand) i zero_points(operand) mogą różnić się od quantization_dimension(result), scales(result) i zero_points(result) w odniesieniu do danych osobowych.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Dla wszystkich d w axes(operand):
  • dim(operand, d) = 1 lub
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Jeśli is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Jeśli dim(operand, quantization_dimension(operand)) = 1, to scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result)))

Przykłady

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

Więcej przykładów

etui

Semantyka

Generuje dane wyjściowe w wyniku wykonania dokładnie 1 funkcji z funkcji branches w zależności od wartości index. Więcej formalnie: result = selected_branch() gdzie:

• selected_branch = branches[index], jeśli 0 <= index < size(branches).
• W przeciwnym razie: selected_branch = branches[-1].

Dane wejściowe

Wyniki

Ograniczenia

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

Przykłady

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

Więcej przykładów

Cbrt

Semantyka

Wykonuje operację pierwiastka sześciennego z uwzględnieniem elementów na tensorze operand i tworzy Tensor result. W zależności od typu elementu wykonuje te działania:

• W przypadku jednostek zmiennoprzecinkowych: rootn(x, 3) w standardzie IEEE-754.
• W przypadku liczb zespolonych: pierwiastek sześcienny zespolony.
• W przypadku typów kwantowych: dequantize_op_quantize(cbrt, operand, type(result))

Dane wejściowe

Wyniki

Ograniczenia

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

Przykłady

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

Więcej przykładów

Ceil

Semantyka

Wykonuje ruch tasorowy tensora operand z uwzględnieniem elementów i tworzy tensor result. Implementuje operację roundToIntegralTowardPositive z standardu IEEE-754 specyfikacji. W przypadku typów kwantowych dequantize_op_quantize(ceil, operand, type(result))

Dane wejściowe

Wyniki

Ograniczenia

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

Przykłady

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

Więcej przykładów

Cholesky

Semantyka

Oblicza rozkład Cholesky'ego grupy macierzy.

Mówiąc bardziej oficjalnie, z okazji wszystkich i w index_space(result). result[i0, ..., iR-3, :, :] to rozkład Cholesky'ego a[i0, ..., iR-3, :, :], w postaci dowolnego z dolnych trójkątów (jeśli lower to true) lub macierz trójkątną górną (jeśli lower to false). wartości wyjściowe w trójkącie przeciwległym, tj. ścisłego górnego trójkąta lub odpowiednio dolnego trójkąta są definiowane na potrzeby implementacji.

Jeśli istnieje i, gdzie macierz wejściowa nie jest dodatnią precyzją hermitańską to zachowanie jest niezdefiniowane.

W przypadku typów kwantowych dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result))

Dane wejściowe

Wyniki

Ograniczenia

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

Przykłady

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

ograniczać (zakres)

Semantyka

Łączy każdy element tensora operand między wartością minimalną a maksymalną i tworzy tensor result. result[result_index] = minimum(maximum(operand[result_index], min_element), max_element). gdzie min_element = rank(min) = 0 ? min[] : min[result_index], max_element = rank(max) = 0 ? max[] : max[result_index] W przypadku typów kwantowych wykonuje dequantize_op_quantize(clamp, min, operand, max, type(result)).

Nakładanie kolejności na liczby zespolone wymaga zaskakującej semantyki, więc w przyszłości planujemy wycofać obsługę liczb zespolonych dla tej operacji (#560).

Dane wejściowe

Wyniki

Ograniczenia

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

Przykłady

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

Więcej przykładów

collective_broadcast

Semantyka

W ramach każdej grupy procesów w siatce procesów StableHLO wyślij wartość klucza operand tensor z procesu źródłowego do procesów docelowych i utwórz Tensor result.

Operacja dzieli siatkę procesów StableHLO na siatkę procesów process_groups, która jest: zdefiniowane w następujący sposób:

• cross_replica(replica_groups), jeśli channel_id <= 0.
• cross_partition(replica_groups), jeśli channel_id > 0.

Później wartość result@process jest określana przez:

• operand@process_groups[i, 0], jeśli istnieje i, dzięki któremu proces jest w aplikacji process_groups[i].
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) w przeciwnym razie.

Dane wejściowe

Wyniki

Ograniczenia

• (C1) is_unique(replica_groups).
• (C2) 0 <= replica_groups < N, gdzie N jest zdefiniowany jako:
  • num_replicas, jeśli używana jest właściwość cross_replica.
  • num_partitions, jeśli używana jest właściwość cross_partition.
• (C3) type(result) = type(operand).

Przykłady

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

Semantyka

W każdej grupie procesów w siatce procesów StableHLO wysyła wartość klucza Tensor operand z procesu źródłowego do procesu docelowego tworzy Tensor result.

Operacja dzieli siatkę procesów StableHLO na siatkę procesów process_groups, która jest: zdefiniowane w następujący sposób:

• cross_replica(source_target_pairs), jeśli channel_id <= 0.
• cross_partition(source_target_pairs), jeśli channel_id > 0.

Później wartość result@process jest określana przez:

• operand@process_groups[i, 0], jeśli występuje i taki jak process_groups[i, 1] = process
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) w przeciwnym razie.

Dane wejściowe

Wyniki

Ograniczenia

• (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, gdzie N jest zdefiniowany jako:
  • num_replicas, jeśli używana jest właściwość cross_replica.
  • num_partitions, jeśli używana jest właściwość cross_partition.
• (C5) type(result) = type(operand).

Przykłady

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

Więcej przykładów

porównaj

Semantyka

Wykonuje porównanie tensorów lhs i rhs z uwzględnieniem elementów zgodnie z comparison_direction i compare_type. Generuje tensor result.

Wartości comparison_direction i compare_type mają następujące właściwości semantyka:

W przypadku elementów z wartościami logicznymi i liczbami całkowitymi:

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

W przypadku typów elementów zmiennoprzecinkowych z parametrem compare_type = FLOAT operator działa następujące operacje IEEE-754:

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

W przypadku elementów zmiennoprzecinkowych typu compare_type = TOTALORDER parametr op korzysta z kombinacji operacji totalOrder i compareQuietEqual z IEEE-754.

W przypadku złożonych typów elementów leksykograficzne porównanie par (real, imag) jest następujące została wykonana za pomocą podanych wartości comparison_direction i compare_type. Nakładanie kolejności na liczby zespolone wymaga zaskakującej semantyki, więc w przyszłości planujemy wycofać obsługę liczb zespolonych gdy comparison_direction to GE, GT, LE lub LT (#560).

Dotyczy typów kwantowych. wykonuje dequantize_compare(lhs, rhs, comparison_direction).

Dane wejściowe

Wyniki

Ograniczenia

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs).
• (C2) shape(lhs) = shape(rhs) = shape(result).
• (C3) Parametr compare_type jest zdefiniowany jako:
  • SIGNED, jeśli is_signed_integer(element_type(lhs)).
  • UNSIGNED, jeśli is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)).
  • FLOAT lub TOTALORDER, jeśli is_float(element_type(lhs)).
  • FLOAT, jeśli is_complex(element_type(lhs)).

Przykłady

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

Więcej przykładów

Złożone

Semantyka

Konwertuje elementy na podstawie wartości złożonej na podstawie pary wartości rzeczywistych i wartości urojonych lhs i rhs, tworząc tensor result.

Dane wejściowe

Wyniki

Ograniczenia

• (C1) type(lhs) = type(rhs).
• (C2) shape(result) = shape(lhs).
• (C3) element_type(result) ma typ complex<E>, gdzie: E = element_type(lhs)

Przykłady

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

Więcej przykładów

wieloskładnikowa

Semantyka

Obejmuje operację składającą się z innych operacji StableHLO, biorąc inputs i composite_attributes oraz produkując results. semantyka operacji jest implementowana przez atrybut decomposition. composite – przeciwnicy można zastąpić własnym programem bez zmiany programu semantyka. W przypadkach, gdy wbudowanie dekompozycji nie zapewnia takiego samego semantyka operacji, preferuj użycie parametru custom_call.

Pole version (wartość domyślna to 0) jest używane do określania, kiedy trzeba zmienić semantykę.

Dane wejściowe

Wyniki

Ograniczenia

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

Przykłady

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

Więcej przykładów

połączyć

Semantyka

Łączy element inputs wzdłuż wymiaru dimension w tej samej kolejności jak podana argumentów i generuje tensor result. Bardziej oficjalnie, result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1], gdzie:

1. id = d0 + ... + dk-1 + kd.
2. Parametr d jest równy dimension, a d0 to d rozmiary wymiarów z inputs.

Dane wejściowe

Wyniki

Ograniczenia

• (C1) same(element_type(inputs...)).
• (C2) same(shape(inputs...)) oprócz 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]) oprócz:
  • dim(result, dimension) = dim(inputs[0], dimension) + ....

Przykłady

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

Więcej przykładów

stała

Semantyka

Generuje tensor output ze stałej wartości value.

Dane wejściowe

Wyniki

Ograniczenia

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

Przykłady

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

Więcej przykładów

dokonają konwersji

Semantyka

Przeprowadza konwersję z jednego typu elementu na inny Tensor operand tworzy tensor result.

W przypadku konwersji typu boolean-to-any-supported-type wartość false to jest przeliczana na zero, a wartość true jest konwertowana na 1. Dla: any-supported-type-to-boolean, wartość 0 jest konwertowana na false i wartości różne od zera są konwertowane na true. Zobacz poniżej, jak to działa działa w przypadku typów złożonych.

w przypadku konwersji, w których występują całkowite liczby zmiennoprzecinkowe i liczba zmiennoprzecinkowa na liczbę całkowitą; lub floating-point-to-floating-point, jeśli wartością źródłową może być dokładnie reprezentowana w typie miejsca docelowego, wynikową wartością będzie dokładnie to, reprezentacja. W przeciwnym razie działanie jest podawane do ustalenia. (#180).

W przypadku konwersji obejmujących liczby floating-point-to-integer część ułamkowa to obcięte. Jeśli obciętej wartości nie można zapisać w typie miejsca docelowego, zachowanie należy do ustalenia (#180).

Konwersja, w której występuje konwersja kompleksowe i złożone, jest taki sam jak w przypadku konwersji konwersji floating-point-to-floating-point na rzeczywiste i części urojonych.

W przypadku konwersji typu complex-to-any-other-type i complex-to-any-other-type wartość urojona źródła jest ignorowana lub wartość urojona miejsca docelowego jest wyzerować odpowiednio. Konwersja części rzeczywistej następuje po konwersji zmiennoprzecinkowych.

Zasadniczo operacja ta może wyrażać dekwantyzację (konwersję z kwantyzowane tensory do tensorów regularnych), kwantyzacja (konwersja ze zwykłych tensory do tensorów kwantyzowanych) i rekwantyzacji (konwersja między kwantyzowanymi danymi tensory), ale obecnie prowadzone są specjalne działania, uniform_dequantize dla pierwszego przypadku użycia i uniform_quantize dla w drugim i trzecim przypadku użycia. W przyszłości te 2 operacje mogą zostać scalone na convert (nr 1576).

Dane wejściowe

Wyniki

Ograniczenia

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

Przykłady

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

Więcej przykładów

splot

Semantyka

Oblicza iloczyn skalarny między oknami o długości lhs i wycinkach rhs, uzyskując wynik result Ten diagram pokazuje, jak elementy w result są obliczane na podstawie lhs i rhs za pomocą konkretnego przykładu.

Rozważ bardziej formalne dostosowanie danych wejściowych w kontekście: lhs aby móc tworzyć okna dla trybu 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).

Ta zmiana kadrowania wykorzystuje te funkcje pomocnicze:

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

Jeśli feature_group_count = 1 i batch_group_count = 1, to dla wszystkich output_spatial_index w: index_space(dim(result, output_spatial_dimensions...)), result[result_shape(:, output_spatial_index, :)] = dot_product, gdzie:

• 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]) Funkcja ta jest prawdopodobnie nieużywana, dlatego planujemy usunąć w przyszłości go (#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]).

Jeśli 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).

Jeśli 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)

W przypadku typów kwantowych funkcja 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)).

W przypadku typów hybrydowych kwantyzowana wartość wylicza 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).

Dane wejściowe

Wyniki

Ograniczenia

• (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) Podany 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) Podany 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) Podany 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) Parametr dim(result, result_dim) jest zdefiniowany jako:
  • dim(lhs, input_batch_dimension) / batch_group_count, jeśli result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension), jeśli result_dim = output_feature_dimension.
  • num_windows w przeciwnym razie, gdzie:
  • 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.
• Jeśli operacja używa tensorów niekwantowych:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Jeśli operacja używa tensorów kwantyzowanych:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Jeśli is_per_axis_quantized(rhs), a następnie quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Jeśli is_per_axis_quantized(result), quantization_dimension(result) = output_feature_dimension
  • Jeśli is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Jeśli is_per_tensor_quantized(rhs), is_per_tensor_quantized(result)
  • Jeśli !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Przykłady

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

Więcej przykładów

cosinus

Semantyka

Wykonuje operacje cosinusowe związane z elementami na tensorze operand i tworzy Tensor result. W zależności od typu elementu wykonuje te działania:

• W przypadku jednostek zmiennoprzecinkowych: cos w standardzie IEEE-754.
• W przypadku liczb zespolonych: cosinus zespolony.
• W przypadku typów kwantowych: dequantize_op_quantize(cosine, operand, type(result)).

Dane wejściowe

Wyniki

Ograniczenia

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

Przykłady

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

Więcej przykładów

count_leading_zeros

Semantyka

Wykonuje liczbę bitów na początku wiersza operand według elementu tensora i tworzy tensor result.

Dane wejściowe

Wyniki

Ograniczenia

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

Przykłady

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

Więcej przykładów

custom_call

Semantyka

Obejmuje zdefiniowaną przez implementację operację call_target_name, która wykonuje inputs i called_computations oraz tworzy results. has_side_effect, backend_config i api_version mogą posłużyć do podania dodatkowych informacji metadanych zdefiniowanych w implementacji.

W tej chwili ta operacja obejmuje dość nieuporządkowany zbiór które odzwierciedlają naturalną ewolucję ich odpowiednika kompilator XLA. W przyszłości planujemy ujednolicenie tych metadanych. (#741).

Dane wejściowe

Wyniki

Przykłady

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

dzielenie

Semantyka

Dzieli dzielną lhs i dzielnik tensorów rhs z uwzględnieniem elementów i generuje tensor result. W zależności od typu elementu wykonuje te działania:

• W przypadku liczb całkowitych: dzielenie całkowity, które daje iloraz algebraiczny z dowolnym odrzucono część ułamkową.
• W przypadku jednostek zmiennoprzecinkowych: division w standardzie IEEE-754.
• W przypadku liczb zespolonych: dzielenie zespolone.
• W przypadku typów kwantowych:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)).

Dane wejściowe

Wyniki

Ograniczenia

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

Przykłady

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

Więcej przykładów

dot_general

Semantyka

Oblicza iloczyn skalarny między wycinkami lhs i rhs, uzyskując wynik Tensor result.

Więcej formalnie, result[result_index] = dot_product, gdzie:

• 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 gdzie size(result_batching_index) = size(lhs_batching_dimensions), size(result_lhs_index) = size(lhs_result_dimensions) i 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))

W przypadku typów kwantowych funkcja 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)).

W przypadku typów hybrydowych kwantyzowana wartość wylicza 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 kontroluje kompromis między szybkością a dokładnością w backendach akceleratora. Może to być jedna z tych wartości (na stronie , semantyka tych wartości wyliczeniowych jest niedostateczna, ale planujemy zająć się tym 755):

• DEFAULT: najszybsze obliczenie, ale najmniej dokładne przybliżenie oryginalny numer.
• HIGH: wolniejsze obliczenia, ale dokładniejsze przybliżenie oryginalny numer.
• HIGHEST: najwolniejsze obliczanie, ale najdokładniejsze przybliżenie oryginalny numer.

DotAlgorithm określa główne właściwości algorytmu używanego do implementacji operacji kropkowej, która określa też dokładność. Jeśli atrybut algorytmu pola, precision_config musi mieć wartość DEFAULT. DotAlgorithms nie mają wartości domyślnej, ponieważ parametry domyślne to implementacja zdefiniowano jego definicję. Dlatego wszystkie pola algorytmu punktowego mogą mieć wartość None, aby określić algorytm z pustymi punktami, który zamiast tego używa wartości precision_config.

Pola DotAlgorithm obejmują:

• lhs_precision_type i rhs_precision_type, wartości precyzji, które po lewej stronie i Wartość RHS operacji jest zaokrąglana do wartości. Typy dokładności są niezależne od typu pamięci masowej danych wejściowych i wyjściowych.
• accumulation_type precyzja używana do agregacji.
• lhs_component_count, rhs_component_count i num_primitive_operations mają zastosowanie przy algorytmie rozkładającym LHS i/lub RHS na ma wiele komponentów i wykonuje wiele elementów podstawowych operacje kropkowe na tych wartości – zwykle w celu emulacji większej dokładności (np. Wykorzystanie typu danych bfloat16 wykorzystujące sztuczną inteligencję do bardziej precyzyjnych obliczeń: bf16_6x tf32_3x itp.). W przypadku algorytmów bez rozkładu wartości te powinna mieć wartość 1.
• allow_imprecise_accumulation, aby określić mniejszą precyzję agregacji jest dozwolony w przypadku niektórych kroków (np. CUBLASLT_MATMUL_DESC_FAST_ACCUM).

Przykładowe atrybuty 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}

To od implementacji zależy, które kombinacje są obsługiwane. W nie gwarantujemy, że każdy algorytm będzie obsługiwany typu akceleratora przez konsumenta StableHLO. Jeśli dany algorytm nie jest można zgłosić błąd, a nie . Weryfikacja StableHLO zapewnia najlepszą weryfikację, co zapobiega algorytmom, o których nie wiemy, że są obsługiwane na jakimkolwiek sprzęcie.

Zobacz xla_data.proto > Algorithm dla niektórych obsługiwanych wartości algorytmu. Zgłoszenie nr 2483 przedstawia plan utworzenia scentralizowany dokument na temat obsługiwanych algorytmów w backendzie.

Dane wejściowe

Wyniki

Ograniczenia

• (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).
• Jeśli operacja używa tensorów niekwantowych:
  • (C13) element_type(lhs) = element_type(rhs).
• Jeśli operacja używa tensorów kwantyzowanych:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C15) zero_points(rhs) = 0.
  • (C16) Jeśli is_per_axis_quantized(rhs), quantization_dimension(rhs) nie znajduje się w grupie rhs_contracting_dimensions.
  • Jeśli is_quantized(lhs):
  • (C17) storage_type(lhs) = storage_type(rhs).
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C19) Jeśli is_per_tensor_quantized(rhs), is_per_tensor_quantized(result)
  • Jeśli !is_quantized(lhs):
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result).
• Jeśli !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.

Przykłady

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

Więcej przykładów

dynamic_broadcast_in_dim

Semantyka

Operacja ta jest funkcjonalna identyczna jak broadcast_in_dim op, ale kształt wyniku jest określany dynamicznie przez output_dimensions.

Operacja akceptuje też opcjonalne atrybuty known_expanding_dimensions, known_non_expanding_dimensions w celu wyrażenia statycznej wiedzy o rozwijaniu wymiarów. Jeśli jej nie określisz, przyjmuje się, że wszystkie wymiary mogą się rozwijać.

Dane wejściowe

Wyniki

Ograniczenia

• (C1) element_type(result) otrzymuje:
  • element_type(operand), jeśli !is_per_axis_quantized(operand).
  • element_type(operand) oprócz quantization_dimension(operand), scales(operand) i zero_points(operand) mogą różnić się od quantization_dimension(result), scales(result) i zero_points(result) w odniesieniu do danych osobowych.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Dla wszystkich d w axes(operand):
  • dim(operand, d) = 1 lub
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Jeśli is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Jeśli dim(operand, quantization_dimension(operand)) = 1, to 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_non_expanding_dimensions).
• (C9) 0 <= known_expanding_dimensions < rank(operand).
• (C10) 0 <= known_non_expanding_dimensions < rank(operand).

Przykłady

// %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_non_expanding_dimensions = array<i64: 1>
} : (tensor<1x3xi64>, tensor<3xi64>) -> tensor<2x3x2xi64>
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

Więcej przykładów

dynamic_conv

Semantyka

Operacja ta jest funkcjonalna identyczna jak convolution , ale dopełnienie jest określane dynamicznie w parametrze padding.

Dane wejściowe

Wyniki

Ograniczenia

• (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) Podany 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) Podany 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) Podany 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) Parametr dim(result, result_dim) jest zdefiniowany jako:
  • dim(lhs, input_batch_dimension) / batch_group_count, jeśli result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension), jeśli result_dim = output_feature_dimension.
  • num_windows w przeciwnym razie, gdzie:
  • 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.
• Jeśli operacja używa tensorów niekwantowych:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Jeśli operacja używa tensorów kwantyzowanych:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Jeśli is_per_axis_quantized(rhs), a następnie quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Jeśli is_per_axis_quantized(result), quantization_dimension(result) = output_feature_dimension
  • Jeśli is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Jeśli is_per_tensor_quantized(rhs), is_per_tensor_quantized(result)
  • Jeśli !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Przykłady

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