Specyfikacja StableHLO

StableHLO to zestaw operacji wysokiego poziomu (HLO) w modelach systemów uczących się. StableHLO działa jako warstwa przenoszenia między różnymi platformami ML i kompilatorami ML: platformy ML tworzące programy StableHLO są zgodne z kompilatorami ML, które korzystają z programów StableHLO.

Naszym celem jest uproszczenie i przyspieszenie programowania ML przez zapewnienie większej interoperacyjności między różnymi platformami ML (np. TensorFlow, JAX i PyTorch) i kompilatorami ML (np. XLA i IREE). W tym celu zamieściliśmy specyfikację języka programowania StableHLO.

Specyfikacja zawiera 3 główne sekcje. Sekcja Programy opisuje strukturę programów StableHLO, które składają 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ę wszystkich tych operacji wykonywanych razem w programie. W sekcji Notacja omawiamy też zapis stosowany w całej specyfikacji.

Aby wyświetlić specyfikację z poprzedniej wersji StableHLO, otwórz repozytorium w oznaczonym wydaniu. Na przykład specyfikacja StableHLO w wersji 0.19.0. Aby wyświetlić zmiany wprowadzone w każdej małej wersji StableHLO, zapoznaj się z logiem wersji w pliku VhloDialect.td.

Programy

Program ::= {Func}

Programy StableHLO składają się z dowolnej liczby funkcji StableHLO. Poniżej znajduje się przykładowy program z funkcją @main, która ma 3 parametry wejściowe (%image, %weights i %bias) oraz 1 wyjście. Treść funkcji zawiera 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 StableHLO (nazywane też funkcjami nazwanymi) mają identyfikator, wejścia/wyjścia i ciało. W przyszłości planujemy wprowadzić dodatkowe metadane funkcji, aby zwiększyć 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 językach programowania, ale mają 2 szczególne cechy: 1) wszystkie identyfikatory mają sygnatury, które odróżniają różne rodzaje 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ż typami najwyższej klasy), które reprezentują wartości StableHLO, oraz typy bez wartości, które opisują inne elementy programu. Typy StableHLO są podobne do typów w wielu językach programowania, a ich główną osobliwością jest specyfika danej domeny, która powoduje pewne nietypowe wyniki (np. typy skalarne nie są typami wartości).

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

Typy Tensor reprezentują tensory, czyli tablice wielowymiarowe. Mają one kształt i typ elementu, gdzie kształt reprezentuje nieujemne lub nieznane rozmiary wymiarów w kolejności rosnącej odpowiadających im wymiarów (zwanych też osiami), których numery wahają się od 0 do R-1. Liczba wymiarów R to ranking. Na przykład tensor<2x3xf32> to typ tensora o kształcie 2x3 i typie elementu f32. Ma 2 wymiary (czyli 2 osi) – 0 i 1, których rozmiary wynoszą odpowiednio 2 i 3. Jego pozycja to 2.

Kształty mogą być częściowo lub całkowicie nieznane (dynamiczne), np. tensor<?x2xf64> jest częściowo nieznany, a tensor<?x?xf64> jest całkowicie nieznany. Dynamiczne rozmiary wymiarów są reprezentowane za pomocą ?. Kształty nie mogą być niesklasyfikowane.

W przyszłości planujemy rozszerzyć typy tensorów poza rozmiary wymiarów i typy elementów, np. o układy (#629) i 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

Typy elementów z kwantyzacją reprezentują wartości całkowite typu magazynowania w zakresie od storage_min do storage_max (łącznie) odpowiadające wartościom zmiennoprzecinkowym typu wyrażonego. Dla danej wartości całkowitej i odpowiadającą wartość zmiennoprzecinkową f można obliczyć jako f = (i - zero_point) * scale, gdzie scale i zero_point to parametry kwantyzacji. Parametry storage_min i storage_max są opcjonalne w gramatyce, ale mają wartości domyślne odpowiednio min_value(storage_type) i max_value(storage_type). Elementy typu „kwantowany” 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 istnieje duże zainteresowanie skalami opartymi na liczbach całkowitych, reprezentowanymi przez mnożniki i przesunięcia. W najbliższej przyszłości planujemy wziąć to pod uwagę (#1404).

Trwa dyskusja na temat semantyki funkcji QuantizationZeroPoint, w tym jej typu, wartości i tego, czy w skwantyzowanym typie tensora może istnieć tylko jeden punkt zerowy, czy może on mieć większą liczbę punktów zerowych. Na podstawie wyników tej dyskusji specyfikacja dotycząca punktów 0 może w przyszłości ulec zmianie (#1405).

Kolejna trwająca dyskusja dotyczy semantyki QuantizationStorageMin i QuantizationStorageMax, aby określić, czy należy nałożyć jakieś ograniczenia na te wartości i wartości skończonych tensorów (#1406).

Na koniec planujemy zgłębić temat przedstawiania nieznanych skal i punktów zerowych w sposób podobny do tego, w jaki planujemy przedstawiać nieznane rozmiary wymiarów (#1407).

Typy kwantyzowanych tensorów reprezentują tensory z kwantyzowanymi elementami. Te tensory są dokładnie takie same jak zwykłe tensory, z tą różnicą, że ich elementy mają kwantowane typy elementów zamiast zwykłych typów elementów.

W przypadku skończonych tensorów kwantyzowanie może być na poziomie tensora, co oznacza, że dla całego tensora jest jeden element scale i zero_point, lub na poziomie osi, co oznacza, że dla danej osi danego wymiaru quantization_dimension jest wiele elementów scales i zero_points. Bardziej formalnie, w tensorze t z kwantyzacją na osi występują dim(t, quantization_dimension) krojenia quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :] itd. Wszystkie elementy w itym krojeniu używają wartości scales[i] i zero_points[i] jako parametrów kwantowania. Typy zaokrąglonych tensorów mają te ograniczenia:

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

TokenType ::= 'token'

Typy tokenów to tokeny, czyli nieprzezroczyste wartości tworzone i wykorzystywane przez niektóre operacje. Tokeny są używane do określania kolejności wykonywania operacji zgodnie z opisem w sekcji Wykonywanie.

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

Typy tupla reprezentują tuple, czyli listy niejednorodne. Kropki to starsza funkcja, która zapewnia zgodność z HLO. W HLO tuple służą do reprezentowania zmiennych danych wejściowych i wyjściowych. W StableHLO obsługiwane są natywny wejścia i wyjścia, a jedynym zastosowaniem tupletów w StableHLO jest kompleksowe reprezentowanie HLO ABI, w którym np. T, tuple<T> i tuple<tuple<T>> mogą się znacznie różnić w zależności od konkretnej implementacji. W przyszłości planujemy wprowadzić zmiany w interfejsie HLO ABI, które mogą pozwolić nam usunąć typy tuple z StableHLO (#598).

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

Typy elementów reprezentują elementy typu tensora. W odróżnieniu od wielu języków programowania te typy nie są typu first class w StableHLO. Oznacza to, że programy StableHLO nie mogą bezpośrednio przedstawiać wartości tych typów (w efekcie są to wartości skalarne typu T z wartościami 0-wymiarowymi tensora typu tensor<T>).

• Typ logiczny reprezentuje wartości logiczne true i false.
• Typy całkowite mogą być typu ze znakiem (si) lub bez znaku (ui) i mieć jedną z obsługiwanych szerokości bitów (2, 4, 8, 16, 32 lub 64). Podpisane typy siN reprezentują liczby całkowite z zakresu od -2^(N-1) do 2^(N-1)-1 włącznie, a bez znaku uiN – wartości całkowite z zakresu od 0 do 2^N-1.
• Typy zmiennoprzecinkowe mogą być następujące:
  • f8E3M4, f8E4M3 i f8E5M2 8-bitowe liczby zmiennoprzecinkową zgodnie z konwencjami IEEE-754.
  • Typy f8E4M3FN i f8E5M2 odpowiadające odpowiednio kodowaniom E4M3 i E5M2 formatu FP8 opisanego w artykule Formaty FP8 do uczenia głębokiego.
  • Typy f8E4M3FNUZ i f8E5M2FNUZ odpowiadające kodowaniom E4M3 i E5M2 formatów FP8 opisanych w artykule 8-bitowe formaty liczbowe na potrzeby głębokich sieci neuronowych.
  • Typ f8E4M3B11FNUZ odpowiadający kodowaniu E4M3 formatów FP8 opisanych w artykule Hybrid 8-bit Floating Point (HFP8) Training and Inference for Deep Neural Networks (Trenowanie i wykonywanie wnioskowania w sieciach głębokich neuronowych przy użyciu hybrydowego formatu 8-bitowego liczby zmiennoprzecinkowej (HFP8)).
  • Typ bf16 odpowiadający formatowi bfloat16 opisanemu w artykule BFloat16: sekret wysokiej wydajności w Cloud TPU.
  • Typy f16, f32 i f64 odpowiadające odpowiednio formatom binary16 („półpełnej precyzji”), binary32 („pełnej precyzji”) i binary64 („podwójnej precyzji”) opisanym w standardzie IEEE 754.
  • Typ tf32 odpowiada formatowi TensorFloat32 i ma ograniczoną obsługę w ramach StableHLO.
  • Typy f4E2M1FN, f6E2M3FN, f6E3M2FN i f8E8M0FNU MX (mikroskalowanie) opisane w specyfikacji formatów mikroskalowania OCP.
• Typy zespolone reprezentują wartości zespolone, które mają część rzeczywistą i część urojoną tego samego typu elementu. Obsługiwane złożone typy to complex<f32> (obie części są typu f32) i 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 wejścia (lista typów po lewej stronie ->) i typy wyjścia (lista typów po prawej stronie ->). W wielu językach programowania typy funkcji są typu first class, ale nie w StableHLO.

StringType ::= 'string'

Typ ciągu tekstowego reprezentuje sekwencje bajtów. W przeciwieństwie do wielu języków programowania typ ciągu znaków nie jest pierwszą klasą w StableHLO i służy tylko do określania statycznych metadanych elementów programu.

Operacje

Operacje StableHLO (nazywane też operacjami) stanowią zamknięty zbiór operacji wysokiego poziomu w modelach uczenia maszynowego. Jak już wspomnieliśmy, składnia StableHLO jest mocno inspirowana MLIR, co niekoniecznie jest najbardziej ergonomiczną alternatywą, ale prawdopodobnie najlepiej pasuje do celu StableHLO, którym jest zwiększenie interoperacyjności między frameworkami i kompilatorami uczenia maszynowego.

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

Operacje StableHLO (nazywane też operacjami) mają nazwę, dane wejściowe/wyjściowe i podpis. Nazwa składa się z prefiksu stablehlo. i mnemotechniki, która jednoznacznie identyfikuje jedno z obsługiwanych operacji. Pełną listę wszystkich obsługiwanych operacji znajdziesz poniżej.

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ą podzielone na wartości wejściowe (obliczone podczas wykonywania), funkcje wejściowe (podawane statycznie, ponieważ w StableHLO funkcje nie są wartościami najwyższej klasy) oraz atrybuty wejściowe (również podawane statycznie). Rodzaj danych wejściowych i wyjściowych zużywanych oraz wytwarzanych przez operatora zależy od jego skrótu. Na przykład funkcja add op pobiera 2 wartości wejściowe i zwraca 1 wartość wyjściową. Natomiast operator select_and_scatter wymaga 3 wartości wejściowych, 2 funkcji wejściowych i 3 atrybutów wejściowych.

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

Funkcje wejściowe (nazywane też funkcjami anonimowymi) są bardzo podobne do funkcji nazwanych, z tym że: 1) nie mają identyfikatora (stąd nazwa „anonimowa”) i 2) nie deklarują typów danych wyjściowych (typy danych są wywnioskowane z opcji return w ramach funkcji).

Składnia funkcji wejściowych zawiera obecnie nieużywaną część (patrz Unused produkcja powyżej), która jest tam ze względu na zgodność z MLIR. W MLIR występuje bardziej ogólna koncepcja „regionów”, które mogą mieć wiele „bloków” operacji połączonych ze sobą za pomocą operacji przeskoku. Te bloki mają identyfikatory odpowiadające Unused produkcji, dzięki czemu można je odróżnić od siebie. StableHLO nie obsługuje wykonywania operacji przejść, więc odpowiadająca mu część składni MLIR nie jest używana (ale nadal jest).

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

Atrybuty wejściowe mają nazwę i wartość, która jest jedną z obsługiwanych stałych. Stanowią one podstawowy sposób określania metadanych statycznych elementów programu. Na przykład operator concatenate używa atrybutu dimension do określania wymiaru, wzdłuż którego wartości wejściowe są łączone. Podobnie operacja slice używa wielu 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 w środowisku naturalnym mogą czasami zawierać atrybuty, które nie są opisane w tym dokumencie. W przyszłości planujemy uwzględnić te atrybuty w opcji StableHLO lub zabronić ich wykorzystywania w programach StableHLO. Oto lista tych atrybutów:

• 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 (lista typów po lewej stronie ->) oraz typów wszystkich wartości wyjściowych (lista typów po prawej stronie ->). Ściśle mówiąc, typy wejścia są zbędne, a typy wyjścia prawie zawsze są zbędne (ponieważ w przypadku większości operacji StableHLO typy wyjścia można wywnioskować z danych wejściowych). Mimo to op signature jest celowo częścią składni StableHLO, aby zapewnić zgodność z MLIR.

Poniżej znajduje się przykład operacji, której skrót to select_and_scatter. Przyjmuje 3 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 typy funkcji i atrybutów wejściowych podane 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 StableHLO mają literał i typ, które razem reprezentują wartość StableHLO. Typ jest zwykle częścią składni stałej, z wyjątkiem sytuacji, gdy jest jednoznaczny (np. stała logiczna ma jednoznacznie typ i1, podczas gdy stała całkowita może mieć kilka możliwych typów).

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

Stałe logiczne reprezentują wartości logiczne true i false. Stałe logiczne 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 liczbowe reprezentują wartości liczb całkowitych za pomocą ciągów tekstowych, które używają zapisu dziesiętnego lub szesnastkowego. Inne systemy liczbowe, np. binarny czy octal, nie są obsługiwane. Stałe liczbowe są objęte tymi ograniczeniami:

• (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 reprezentują wartości zmiennoprzecinkowe za pomocą ciągów znaków, które używają zapisu dziesiętnego lub wykładniczego. Dodatkowo przy użyciu notacji szesnastkowej można bezpośrednio określać bazowe bity w formacie zmiennoprzecinkowym odpowiedniego typu. Stałe zmiennoprzecinkowe są objęte tymi ograniczeniami:

• (C1) Jeśli używany jest zapis w formacie innym niż szesnastkowy, is_wellformed(float_literal, float_type).
• (C2) Jeśli używana jest notacja szesnastkowa, size(hexadecimal_digits) = num_bits(float_type) / 4.

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

Stałe zespolone reprezentują wartości zespolone za pomocą list części rzeczywistej (pojawia się jako pierwsza) i części urojonej (pojawia się jako druga). Na 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ść, w jakiej te części są przechowywane w pamięci, jest zdefiniowana przez implementację. 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 tensora reprezentują wartości tensora za pomocą zagnieżdżonych list określonych za pomocą notacji NumPy. Na przykład dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> reprezentuje wartość tensora z tym 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ą przechowywane w pamięci, jest zdefiniowana przez implementację. 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) '>'

Zwartości skonwertowanego tensora reprezentują wartości skonwertowanego tensora za pomocą tej samej notacji co stałe tensora, a ich elementy są określone jako stałe ich typu magazynu. 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 i sekwencji ucieczki. Są one niezależne od kodowania, więc interpretacja tych bajtów jest zdefiniowana przez implementację. Literały łańcuchowe mają typ string.

Operacje

abs

Semantyka

Wykonuje elementarną operację abs na tensorze operand i tworzy tensor result. W zależności od typu elementu:

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

Dane wejściowe

Wyniki

Ograniczenia

• (C1) shape(result) = shape(operand).
• (C2) baseline_element_type(result) zdefiniowano jako:
  • complex_element_type(element_type(operand)) jeżeli is_complex(operand).
  • baseline_element_type(operand) w innych przypadkach.

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 dwóch tensorów lhs i rhs i tworzy tensor result. W zależności od typu elementu:

• W przypadku wartości logicznych: operator logiczny LUB.
• W przypadku liczb całkowitych: dodawanie liczb całkowitych.
• Dla liczb zmiennoprzecinkowych: addition z IEEE-754.
• W przypadku liczb zespolonych: dodawanie zespolone.
• W przypadku typów skwantowanych: dequantize_op_quantize(add, lhs, rhs, type(result)).

Dane wejściowe

Wyniki

Ograniczenia

• Jeśli operacja używa niespakowanych tensorów:
  • (C1) type(lhs) = type(rhs) = type(result).
• Jeśli operacja używa kwantowanych tensorów:
  • (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 generujące inputs są wykonywane przed operacjami zależnymi od result. Wykonywanie tej operacji nie powoduje żadnych zmian. Służy ona tylko do ustanowienia zależności danych z poziomu 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 konkatenuje wartości tensorów operands z każdego procesu wzdłuż wymiaru all_gather_dim i tworzy tensory results.

Operacja dzieli siatkę procesu StableHLO na process_groups, która jest zdefiniowana w ten sposób:

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

Następnie w każdym wierszu process_group:

• operands...@receiver = [operand@sender for sender in process_group] za wszystkie receiver w process_group.
• results...@process = concatenate(operands...@process, all_gather_dim) za wszystkie 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 ramach każdej grupy procesów w siatce procesów StableHLO stosuje funkcję redukcji computation do wartości tensorów operands z każdego procesu i generuje tensory results.

Operacja dzieli siatkę procesu StableHLO na process_groups, która jest zdefiniowana w ten sposób:

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

Następnie w każdym wierszu process_group:

• results...@process[result_index] = exec(schedule) dla niektórych drzew binarnych schedule gdzie:
  • exec(node) = computation(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value.
• schedule to drzewo binarne zdefiniowane przez implementację, którego traversal w kolejności to to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0])).

Dane wejściowe

Wyniki

Ograniczenia

• (C1) is_unique(replica_groups).
• (C2) size(replica_groups) zdefiniowano 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 ramach każdej grupy procesów w siatce procesów StableHLO dzieli wartości tensorów operands wzdłuż split_dimension na części, rozprasza te części między procesami, konkatenuje rozproszone części wzdłuż concat_dimension i tworzy tensory results. Operacja dzieli siatkę procesu StableHLO na process_groups, która jest zdefiniowana w ten sposób:

• cross_replica(replica_groups) jeżeli channel_id <= 0.
• cross_partition(replica_groups) jeżeli channel_id > 0.

Następnie w każdym wierszu process_group:

• split_parts...@sender = split(operands...@sender, split_count, split_dimension)dla wszystkich 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) 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, tworzy intensywność result. W zależności od typu elementu:

• W przypadku wartości logicznych: operator logiczny OR.
• 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 elementową operację atan2 na tensorach lhs i rhs, tworząc tensor result. W zależności od typu elementu:

• Dla liczb zmiennoprzecinkowych: atan2 z IEEE-754.
• W przypadku liczb zespolonych: complex atan2.
• W przypadku typów skwantowanych: 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 wejść batch_norm_training z backpropagatinggrad_output i tworzy tensory grad_operand, grad_scale i grad_offset. Bardziej formalnie ta operacja może zostać wyrażona jako dekompozycja istniejących operacji StableHLO z użyciem składni Pythona:

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 skwantowanych wykonuje działanie 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ą te same wartości baseline_element_type.
• (C3) operand, grad_output i grad_operand mają ten sam kształt.
• (C4) Znaczniki scale, mean, variance, grad_scale i grad_offset mają ten 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 z wyjątkiem wymiaru feature_index i tworzy tensor result. Bardziej formalnie tę operację można wyrazić jako dekompozycję istniejących operacji StableHLO za pomocą 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 wybiera 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ą te same ustawienia 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ę we wszystkich wymiarach oprócz wymiaru feature_index oraz normalizuje tensor operand, tworząc tensory output, batch_mean i batch_var. Bardziej formalnie ta operacja może być wyrażona jako dekompozycja istniejących operacji StableHLO z użyciem składni Pythona:

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 skwantowanych wykonuje działanie 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 mają 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ę bitcast na tensorze operand i tworzy tensor result, w którym bity całego tensora operand są ponownie interpretowane przy użyciu typu tensora result.

Bardziej formalnie, jeśli 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 reprezentację danej wartości w pamięci, a jej działanie jest definiowane przez implementację, ponieważ dokładna reprezentacja tensorów jest zdefiniowana przez implementację, a dokładna reprezentacja typów elementów również jest zdefiniowana w implementacji.

Dane wejściowe

Wyniki

Ograniczenia

• (C1) Przy założeniu, że 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) dla wszystkich 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

Rozwija wymiary lub rangę wejściowego tensora przez powielanie danych w tensorze operand i tworzenie tensora result. Formalnie: result[result_index] = operand[operand_index] dla wszystkich d w ramach axes(operand):

• operand_index[d] = 0 jeżeli dim(operand, d) = 1.
• operand_index[d] = result_index[broadcast_dimensions[d]] w innych przypadkach.

Dane wejściowe

Wyniki

Ograniczenia

• (C1) element_type(result) otrzymuje:
  • element_type(operand), jeśli !is_per_axis_quantized(operand).
  • element_type(operand), z tym że quantization_dimension(operand), scales(operand) i zero_points(operand) mogą się różnić od quantization_dimension(result), scales(result) i zero_points(result) odpowiednio.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) W przypadku 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. Bardziej formalnie: result = selected_branch()gdzie:

• selected_branch = branches[index] jeżeli 0 <= index < size(branches).
• selected_branch = branches[-1] w innych przypadkach.

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 elementarną operację pierwiastka sześciennego 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.
• Liczby zespolone: pierwiastek sześcienny zespolony.
• W przypadku typów skwantowanych: 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

Przeprowadza element po elemencie zaokrąglenie w górę tensora operand i generuje tensor result. Realizuje operację roundToIntegralTowardPositive według specyfikacji IEEE-754. W przypadku typów skwantowanych wykonuje działanie 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 Choleskiego dla zbioru macierzy.

W bardziej formalnym ujęciu dla wszystkich i w index_space(result) result[i0, ..., iR-3, :, :] jest rozkładem Cholesky'ego a[i0, ..., iR-3, :, :] w postaci dolnej macierzy trójkątnej (jeśli lower jest true) lub górnej macierzy trójkątnej (jeśli lower jest false). Wartości wyjściowe w trójkącie przeciwległym, tj. odpowiednio ścisły górny trójkąt lub odpowiednio ścisły trójkąt dolny, są definiowane na podstawie implementacji.

Jeśli istnieje element i, w którym macierz wejściowy nie jest macierzową o dodatnie dodatnim, działanie jest niezdefiniowane.

W przypadku typów skwantowanych wykonuje działanie 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ą, tworząc tensor result. Bardziej formalnie: 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 skwantowanych wykonuje działanie dequantize_op_quantize(clamp, min, operand, max, type(result)).

Narzucanie kolejności liczbom zespolonym wiąże się z nieoczywistą semantyką, dlatego w przyszłości planujemy usunąć obsługę liczb zespolonych w przypadku 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 każdej grupie procesów w siatce procesów StableHLO wyślij wartość tensora operand z procesu źródłowego do procesów docelowych i utwórz tensor result.

Ta operacja dzieli siatkę procesów StableHLO na siatkę procesów process_groups, która jest zdefiniowana w ten sposób:

• cross_replica(replica_groups) jeżeli channel_id <= 0.
• cross_partition(replica_groups) jeżeli channel_id > 0.

Następnie result@process jest obliczana jako:

• operand@process_groups[i, 0], jeśli istnieje i, który określa, że proces znajduje się w regionie process_groups[i].
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result))w innym przypadku.

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ść tensora operand z procesu źródłowego do procesu docelowego i tworzy tensor result.

Operacja dzieli siatkę procesu StableHLO na process_groups, która jest zdefiniowana w ten sposób:

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

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

• operand@process_groups[i, 0], jeśli istnieje i takie, że process_groups[i, 1] = process.
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result))

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

Porównuje elementy tensorów lhs i rhs zgodnie z definicjami comparison_direction i compare_type, tworząc tensor result.

Wartości comparison_direction i compare_type mają następującą semantykę:

W przypadku typów elementów wartości logicznych i liczb całkowitych:

• 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 compare_type = FLOAT operator op implementuje te operacje IEEE-754:

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

W przypadku typów elementów zmiennoprzecinkowych z wartością compare_type = TOTALORDER operator używa kombinacji operacji totalOrder i compareQuietEqual z normy IEEE-754.

W przypadku złożonych typów elementów przeprowadzane jest porównanie leksykograficzne par (real, imag) za pomocą podanych wartości comparison_direction i compare_type. Narzucanie kolejności liczbom zespolonym wiąże się z nieoczywistą semantyką, dlatego w przyszłości planujemy usunąć obsługę liczb zespolonych, gdy comparison_direction = GE, GT, LE lub LT (#560).

W przypadku typów poddanych kwantyzacji wykonuje działanie 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) compare_type jest zdefiniowany jako:
  • SIGNED jeżeli is_signed_integer(element_type(lhs)).
  • UNSIGNED jeżeli is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)).
  • FLOAT lub TOTALORDER, jeśli is_float(element_type(lhs)).
  • FLOAT jeżeli 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

Przeprowadza konwersję element po elemencie na wartość zespoloną z pary wartości rzeczywistych i urojonych, lhs i rhs, oraz generuje 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

Zawiera operację złożoną z innych operacji StableHLO, przyjmując inputs i composite_attributes oraz zwraca results. Semantyka operacji jest implementowana przez atrybut decomposition. Opcję composite można zastąpić jej rozkładem bez zmiany semantyki programu. Jeśli wstawienie dekompozycji w kod źródłowy nie zapewnia tej samej semantyki op, użyj custom_call.

Pole version (domyślnie 0) służy do informowania o zmianie semantyki elementu złożonego.

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

konkatenacja

Semantyka

Łączy elementy inputs wzdłuż wymiaru dimension w takim samym porządku jak podane argumenty i tworzy tensor result. Bardziej formalnie: result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1], gdzie:

1. id = d0 + ... + dk-1 + kd.
2. d jest równy dimension, a d0, ... to drozmiary inputs.

Dane wejściowe

Wyniki

Ograniczenia

• (C1) same(element_type(inputs...)).
• (C2) same(shape(inputs...)) z wyjątkiem 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

Tworzy tensor output na podstawie stałej 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ę elementów z jednego typu na inny w tensorze operand i tworzy tensor result.

W przypadku konwersji boolean-to-any-supported-type wartość false jest zamieniana na 0, a wartość true na 1. W przypadku konwersji typu any-supported-type-to-boolean wartość 0 jest konwertowana na false, a wartości inne niż zero – na true. Poniżej znajdziesz informacje o tym, jak to działa w przypadku typów złożonych.

W przypadku konwersji zawierających liczby całkowite na liczbę całkowitą, liczbę zmiennoprzecinkową lub zmiennoprzecinkową na liczbę zmiennoprzecinkową wartość źródłową może być dokładnie reprezentowana w typie miejsca docelowego. W przeciwnym razie zachowanie jest nieokreślone (#180).

W przypadku konwersji zawierających floating-point-to-integer część ułamkowa jest skracana. Jeśli skrócona wartość nie może być przedstawiona w przypadku typu miejsca docelowego, zachowanie jest nieokreślone (#180).

Konwersja z typu zespolonego na typ zespolony działa tak samo jak konwersja z typu zmiennoprzecinkowego na typ zmiennoprzecinkowy w przypadku konwersji części rzeczywistej i części urojonej.

W przypadku konwersji z typu złożonego na dowolny inny typ i z dowolnego innego typu na złożony wartość wyobrażona źródła jest ignorowana lub docelowa wartość wyobrażona jest ustawiana na 0. Konwersja części rzeczywistej następuje zgodnie z konwersją na liczby zmiennoprzecinkowe.

Zasadniczo ta operacja może wyrażać dekwantyzację (przekształcanie tensorów regularnych w tensory regularne), kwantyzację (konwersja z tensorów regularnych na tensory kwantyzowane) i rekwantyzację (przekształcanie między kwantyzowanymi procesorami), ale obecnie mamy dla nich specjalne operacje – uniform_dequantize dla pierwszego przypadku użycia i uniform_quantize w drugim i trzecim przypadku użycia. W przyszłości te 2 operacje mogą zostać połączone w convert (#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

splotu

Semantyka

Oblicza iloczyny skalarne między oknami lhs i wycinkami rhs oraz generuje result. Na diagramie poniżej widać, jak elementy w elementach result są obliczane na podstawie elementów lhs i rhs.

Bardziej formalnie rozważ następujące zmiany w danych wejściowych w postaci elementu lhs, aby można było tworzyć przedziały czasu 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]). Ta funkcja wydaje się nieużywana, więc planujemy ją usunąć w przyszłości (#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 wybiera 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) Zgodnie z 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) Zgodnie z 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) jest zdefiniowany jako:
  • dim(lhs, input_batch_dimension) / batch_group_count jeżeli result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) jeżeli 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 kwantowanych tensorów:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Jeśli is_per_axis_quantized(rhs), to quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Jeśli is_per_axis_quantized(result), to 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), to 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 elementarną operację cosinusa na tensorze operand i tworzy tensor result. W zależności od typu elementu:

• Dla liczb zmiennoprzecinkowych: cos z 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 element po elemencie zliczanie liczby początkowych zer w tensorze operandi 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

Zawiera zdefiniowaną przez implementację operację call_target_name, która przyjmuje argumenty inputs i called_computations oraz zwraca results. Atrybuty has_side_effect, backend_config i api_version mogą służyć do udostępniania dodatkowych metadanych zdefiniowanych przez implementację.

Obecnie ta operacja zawiera dość nieuporządkowaną kolekcję metadanych, która odzwierciedla organiczną ewolucję jej odpowiednika w kompilatorze XLA. W przyszłości planujemy ujednolicić te metadane (#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

Wykonuje element po elemencie dzielenie tensorów dzielnika lhs i dzielenia rhs oraz zwraca tensor result. W zależności od typu elementu:

• W przypadku liczb całkowitych: dzielenie liczb całkowitych, które zwraca iloraz algebraiczny z pominięciem części ułamkowej.
• Dla liczb zmiennoprzecinkowych: division z IEEE-754.
• W przypadku liczb zespolonych: dzielenie zespolone.
• W przypadku typów skwantowanych:
  • 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 tensor result.

Bardziej 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_indexgdzie 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 skwantowanych wykonuje operację 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 hybrydowych typów kwantyzacji wykonuje działanie 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ą obliczeń na backendzie akceleratora. Może to być jedna z tych wartości (obecnie semantyka tych wartości jest niewystarczająco sprecyzowana, ale planujemy rozwiązać ten problem w bładze zgłoszenia 755):

• DEFAULT: najszybsze obliczenia, ale najmniej dokładne przybliżenie do pierwotnego wyniku.
• HIGH: wolniejsze obliczenia, ale dokładniejsze przybliżenie do pierwotnej wartości.
• HIGHEST: najwolniejsze obliczenia, ale najbardziej dokładne przybliżenie do pierwotnej wartości.

DotAlgorithm definiuje główne właściwości algorytmu używanego do implementacji operacji kropki, która określa też dokładność. Jeśli pola atrybutów algorytmu są ustawione, precision_config musi mieć wartość DEFAULT. DotAlgorithms nie mają wartości domyślnej, ponieważ parametry domyślne są definiowane przez implementację. Dlatego wszystkie pola algorytmu kropki mogą być ustawione na None, aby określić pusty algorytm kropki, który zamiast tego użyje wartości precision_config.

Pola DotAlgorithm:

• lhs_precision_type i rhs_precision_type, dokładność, do której zaokrąglane są wartości po lewej i po prawej stronie operacji. Typy dokładności są niezależne od typów magazynowania danych wejściowych i wyjściowych.
• accumulation_type dokładność użyta do skumulowania.
• lhs_component_count, rhs_component_count i num_primitive_operations stosuje się, gdy używamy algorytmu, który rozkłada lewą lub prawą stronę na kilka komponentów i wykonuje na tych wartościach wiele „prostych” operacji dot – zwykle w celu emulacji większej precyzji (np. Korzystanie z typu danych bfloat16 do obliczeń o większej precyzji: bf16_6x tf32_3x itp.). W przypadku algorytmów bez dekompozycji te wartości powinny wynosić 1.
• allow_imprecise_accumulation, aby określić, czy akumulacja z mniejszą dokładnością jest dozwolona 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. Ogólnie nie ma gwarancji, że każdy algorytm jest obsługiwany przez każdy typ akceleratora przez użytkownika StableHLO. Jeśli dany algorytm nie jest obsługiwany, należy zgłosić błąd zamiast korzystać z alternatywnego rozwiązania. Weryfikacja StableHLO zapewni najlepszą weryfikację, zapobiegając działaniu algorytmów, które nie są obsługiwane na żadnym sprzęcie.

Niektóre obsługiwane wartości algorytmu znajdziesz w sekcji xla_data.proto > Algorithm. zgłoszenie #2483 zawiera plan stworzenia scentralizowanego dokumentu na temat obsługiwanych algorytmów na zapleczu.

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 niespakowanych tensorów:
  • (C13) element_type(lhs) = element_type(rhs).
• Jeśli operacja używa kwantowanych tensorów:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C15) zero_points(rhs) = 0.
  • (C16) Jeśli is_per_axis_quantized(rhs), to quantization_dimension(rhs) nie jest w 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), to 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

Ta operacja jest pod względem funkcjonalnym identyczna z operacją broadcast_in_dim, ale kształt wyniku jest określany dynamicznie za pomocą parametru output_dimensions.

Operacja akceptuje też opcjonalne atrybuty known_expanding_dimensions i known_nonexpanding_dimensions, które służą do wyrażania statycznej wiedzy o zachowaniu rozszerzania wymiarów. Jeśli nie są określone, przyjmuje się, że wszystkie wymiary mogą się rozszerzać.

Dane wejściowe

Wyniki

Ograniczenia

• (C1) element_type(result) otrzymuje:
  • element_type(operand), jeśli !is_per_axis_quantized(operand).
  • element_type(operand), z tym że quantization_dimension(operand), scales(operand) i zero_points(operand) mogą się różnić od quantization_dimension(result), scales(result) i zero_points(result) odpowiednio.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) W przypadku 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_nonexpanding_dimensions).
• (C9) 0 <= known_expanding_dimensions < rank(operand).
• (C10) 0 <= known_nonexpanding_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_nonexpanding_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

Ta operacja jest pod względem funkcjonalnym identyczna z operacją convolution, ale wypełnienie jest podawane dynamicznie za pomocą 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) Zgodnie z 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) Zgodnie z 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) jest zdefiniowany jako:
  • dim(lhs, input_batch_dimension) / batch_group_count jeżeli result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) jeżeli 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 kwantowanych tensorów:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Jeśli is_per_axis_quantized(rhs), to quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Jeśli is_per_axis_quantized(result), to 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), to 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]]
//          ]]