Specyfikacja StableHLO

StableHLO to operacja ustawiona na potrzeby operacji wysokiego poziomu w modelach systemów uczących się. StableHLO działa jako warstwa przenośności między różnymi platformami ML i kompilatorami ML: platformy systemów uczących się, które tworzą programy StableHLO, są zgodne z kompilatorami ML, które używają programów StableHLO.

Naszym celem jest uproszczenie i przyspieszenie tworzenia systemów uczących się przez zapewnienie większej interoperacyjności między różnymi platformami ML (np. TensorFlow, JAX i PyTorch) a kompilatorami ML (np. XLA i IREE). W tym dokumencie znajdziesz specyfikację języka programowania StableHLO.

Specyfikacja obejmuje 3 główne sekcje. Pierwsza sekcja Programs (Programy) opisuje strukturę programów StableHLO, które składają się z funkcji StableHLO, które same składają się z operacji StableHLO. Sekcja Operacje w tej strukturze określa semantykę poszczególnych operacji. Sekcja Wykonanie zawiera opis semantyki wszystkich operacji wykonywanych razem w ramach programu. Na koniec w sekcji Notacja omówiono zapisy używane w danej specyfikacji.

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 wejścia (%image, %weights i %bias) oraz 1 wyjście. 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 stabilnych HLO (nazywane też funkcjami nazwanymi) mają identyfikator, dane wejściowe/wyjściowe oraz treść. W przyszłości planujemy wprowadzić dodatkowe metadane funkcji, aby uzyskać większą 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 stabilnej HLO są podobne do identyfikatorów w wielu językach programowania. Występują 2 charakterystyczne cechy: 1) wszystkie identyfikatory zawierają sigil, które rozróżniają różne rodzaje identyfikatorów; 2) identyfikatory wartości mogą być w pełni liczbowe, aby uprościć generowanie programów StableHLO.

Typy

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

Typy stabilny HLO są podzielone na typy wartości (nazywane też typami pierwszej klasy), które reprezentują wartości w systemie StableHLO i typy innych niż wartości, które opisują inne elementy programu. Typy StableHLO są podobne do typów w wielu językach programowania, a główną cechą jest specyficzna dla StableHLO charakter charakterystyczny dla domeny, co prowadzi do nietypowych wyników (np. typy skalarne nie są typami wartości).

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

Typy tensorów reprezentują tensory, tj. tablice wielowymiarowe. Mają kształt i typ elementu, gdzie kształt oznacza nieujemne rozmiary w kolejności rosnącej od odpowiednich wymiarów (nazywanych też osiami) o numerach od 0 do R-1. Liczba wymiarów R to ranking. Na przykład tensor<2x3xf32> to tensor typu tensor o kształcie 2x3 i typie elementu f32. Ma 2 wymiary (czyli 2 osie) – 0 i 1. Ma pozycję 2.

Definiuje obsługę kształtów statycznych, w których rozmiary wymiarów są statycznie znane. W przyszłości zamierzamy też wprowadzić obsługę kształtów dynamicznych, w przypadku których rozmiary wymiarów są częściowo lub całkowicie nieznane (#8). Planujemy też poszerzyć zakres typów tensorów poza rozmiarami i typami elementów, np. o układy (#629) i rozproszenie (#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
Nazwa Typ Ograniczenia
storage_type typ liczby całkowitej (C1–C4), (C9)
storage_min stała liczba całkowita [C2], [C4], [C8]
storage_max stała liczba całkowita [C3], [C4], [C8]
expressed_type typ liczby zmiennoprzecinkowej (C1), [C5]
quantization_dimension opcjonalna stała liczba całkowita (C11-C13)
scales zmienna stałej liczby zmiennoprzecinkowej (C5-C7), (C10), (C11), (C13)
zero_points zmienna stałych liczb całkowitych (C8-C10)

Typy elementów kwantyzowanych reprezentują wartości całkowite typu pamięci z zakresu od storage_min do storage_max (włącznie), które odpowiadają wartościom zmiennoprzecinkowym typu wyrażonego. Dla danej liczby całkowitej i odpowiednią wartość zmiennoprzecinkową f można obliczyć jako f = (i - zero_point) * scale, gdzie scale i zero_point są nazywane parametrami kwantyzacyjnymi. Znaczniki storage_min i storage_max są w gramacie opcjonalne, ale mają wartości domyślne odpowiednio min_value(storage_type) i max_value(storage_type). Typy elementów skwantyzowanych podlegają tym ograniczeniom:

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

W tej chwili QuantizationScale jest stałą zmiennoprzecinkową, ale istnieje duże zainteresowanie skalami opartymi na liczbach całkowitych, reprezentowanych za pomocą mnożników i przesunięć. Planujemy zbadać tę funkcję w najbliższej przyszłości (#1404).

Trwa dyskusja na temat semantyki obiektu QuantizationZeroPoint, w tym jego typu, wartości oraz tego, czy w skwantyzowanym typie tensora może istnieć tylko 1 punkt zerowy lub potencjalnie wiele punktów zerowych. Na podstawie wyników tej dyskusji stwierdziliśmy, że specyfikacja około 0 punktów może w przyszłości ulec zmianie (#1405).

Kolejnym tematem jest semantyka QuantizationStorageMin i QuantizationStorageMax w celu określenia, czy na te wartości i na wartości skwantyzowanych tensorów należy wprowadzić jakieś ograniczenia (#1406).

Na koniec zajmiemy się reprezentowaniem nieznanych skal i punktów zerowych, podobnie jak w przypadku reprezentowania nieznanych rozmiarów wymiarów (#1407).

Typy tensorów skwantyzowanych reprezentują tensory z elementami skwantyzowanymi. Tensory te są dokładnie takie same jak tensory zwykłe, z tą różnicą, że ich elementy zawierają skwantyzowane typy elementów zamiast zwykłych typów.

W skwantyzowanych tensorach kwantyzacja może odbywać się na tensor, co oznacza, że dla całego tensora może być przypisana 1 scale i zero_point dla całego tensora lub na osi, co oznacza, że kwantyzowanie może obejmować wiele elementów scales i zero_points, po 1 parze na wycinek danego wymiaru quantization_dimension. Dokładniej rzecz biorąc, w tensorze t z kwantyzacją według osi znajdują się wycinki quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :] itd. Wszystkie elementy i używają scales[i] i zero_points[i] jako parametrów kwantyzacji.dim(t, quantization_dimension) Kwantyzowane typy tensorów mają te ograniczenia:

  • Kwantyzacja na intensywność:
    • Bez dodatkowych ograniczeń.
  • Kwantyzowanie według osi:
    • (C12) quantization_dimension < rank(self).
    • (C13) dim(self, quantization_dimension) = size(scales).
TokenType ::= 'token'

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

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

Typy tupli reprezentują krotki, czyli listy heterogeniczne. Kropki to starsza funkcja, która istnieje tylko na potrzeby zgodności z HLO. W HLO krotki są używane do reprezentowania zróżnicowanych danych wejściowych i wyjściowych. W stableHLO zmienne dane wejściowe i wyjści są natywnie obsługiwane, a jedyne użycie krotek w wersji StableHLO polega na wszechstronnym reprezentowaniu interfejsu HLO ABI, gdzie na przykład 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ą umożliwić usunięcie typów krotek ze StableHLO (#598).

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

Typy elementów reprezentują elementy typów tensorów. W przeciwieństwie do wielu języków programowania te typy nie należą do pierwszej klasy w języku StableHLO. Oznacza to, że programy StableHLO nie mogą bezpośrednio przedstawiać wartości tego typu (w związku z tym reprezentowanie wartości skalarnych typu T z 0-wymiarowymi wartościami tensorów typu tensor<T> jest idiomatyczne).

  • Typ wartości logicznej reprezentuje wartości logiczne true i false.
  • Typy liczb całkowitych mogą być podpisane (si) lub nieoznaczone (ui) i mieć jedną z obsługiwanych szerokości bitowych (4, 8, 16, 32 lub 64). Podpisane typy siN reprezentują wartości całkowite od -2^(N-1) do 2^(N-1)-1, a nieoznaczone typy uiN reprezentują wartości całkowite z zakresu od 0 do 2^N-1 włącznie.
  • Możliwe typy liczby zmiennoprzecinkowej:
  • Typy złożone to wartości złożone, które mają część rzeczywistą i część urojoną tego samego typu elementu. Obsługiwane typy złożone 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 danych wejściowych (lista typów po lewej stronie interfejsu ->) i typy danych wyjściowych (lista typów po prawej stronie interfejsu ->). W wielu językach programowania typy funkcji są klasy pierwszej, ale nie w stableHLO.

StringType ::= 'string'

Typ ciągu znaków 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 jest używany tylko do określania statycznych metadanych elementów programu.

Operacje

Operacje stabilnego HLO (nazywane też operacjami) to zamknięty zestaw operacji wysokiego poziomu w modelach systemów uczących się. Jak już wspomnieliśmy, składnia wersji StableHLO w dużym stopniu bazuje na MLIR, który niekoniecznie jest najbardziej ergonomiczną alternatywą, ale zapewne najlepiej pasuje do celu StableHLO, którym jest zapewnienie większej interoperacyjności między platformami ML a kompilatorami ML.

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

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

Obecnie programy StableHLO w środowisku dzikiej przyrody czasami zawierają operacje, które nie zostały opisane w tym dokumencie. W przyszłości planujemy uwzględnić te operacje w opcji StableHLO lub zabronić ich pojawiania się w programach StableHLO. Na razie możesz zapoznać się z listą tych operacji:

  • builtin.module, func.func, func.call i func.return (#425).
  • Operacje chlo (#602).
  • Kategoria „Not in HLO” w operacjach StableHLO – początkowo była częścią operacji StableHLO, ale później uznano, że nie pasuje do tej kategorii: broadcast, create_token, cross-replica-sum, dot, einsum, torch_index_select, unary_einsum (3).
  • Kategoria „Dynamism” operacji StableHLO – zostały one pobrane z MHLO, ale jeszcze ich nie specyfikowaliśmy: compute_reshape_shape, cstr_reshapable, dynamic_broadcast_in_dim, dynamic_conv, dynamic_gather, dynamic_iota, dynamic_pad, dynamic_reshape, real_dynamic_slice i set_dimension_size (#8).
  • Obliczanie kształtów, w tym operacje arith, shape i tensor (#8).
OpInputs        ::= OpInputValues OpInputFuncs OpInputAttrs
OpInputValues   ::= '(' [OpInputValue {',' OpInputValue}] ')'
OpInputValue    ::= ValueId
OpInputFuncs    ::= ['(' OpInputFunc {',' OpInputFunc} ')']
OpInputAttrs    ::= ['{' OpInputAttr {',' OpInputAttr} '}']
OpOutputs       ::= [OpOutput {',' OpOutput} '=']
OpOutput        ::= ValueId

Operacje zużywają dane wejściowe i generują dane wyjściowe. Dane wejściowe są podzielone na wartości wejściowe (obliczane podczas wykonywania), funkcje wejściowe (podawane statycznie, ponieważ w przypadku funkcji StableHLO nie są wartościami pierwszej klasy) i atrybuty wejściowe (również statycznie). Rodzaj danych wejściowych i wyjściowych zużywanych i wytwarzanych przez operację zależy od jej mnemotechniki. Na przykład operacja add wykorzystuje 2 wartości wejściowe i generuje jedną wartość wyjściową. W porównaniu ta operacja select_and_scatter wykorzystuje 3 wartości wejściowe, 2 funkcje wejściowe i 3 atrybuty wejściowe.

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

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

Składnia funkcji wejściowych obejmuje obecnie nieużywaną część (patrz sekcja produkcyjna Unused powyżej), która zapewnia zgodność z MLIR. W MLIR istnieje bardziej ogólna koncepcja „regionów”, które mogą obejmować wiele „bloków” operacji połączonych ze sobą za pomocą operacji skoków. Te bloki mają identyfikatory odpowiadające produkcji Unused, dzięki czemu można je od siebie odróżnić. StableHLO nie ma operacji skoków, więc odpowiadająca jej część składni MLIR nie jest używana (ale nadal tam jest).

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

Atrybuty wejściowe mają nazwę i wartość, które są jedną z obsługiwanych stałych. Stanowią one podstawowy sposób określania statycznych metadanych elementów programu. Na przykład operacja concatenate używa atrybutu dimension do określenia wymiaru, z którym są łączone jej wartości wejściowe. Podobnie operacja slice używa wielu atrybutów, takich jak start_indices i limit_indices, aby określić granice służące do wycięcia wartości wejściowej.

Obecnie programy StableHLO w środowisku dzikiej przyrody czasami zawierają atrybuty, które nie zostały opisane w tym dokumencie. W przyszłości planujemy uwzględnić te atrybuty w opcji StableHLO lub zabronić ich pojawiania się w programach StableHLO. Na razie przejrzyj listę 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 operacyjny składa się z typów wszystkich wartości wejściowych (listy typów po lewej stronie interfejsu ->) i typów wszystkich wartości wyjściowych (listy typów po prawej stronie ->). Ściśle mówiąc, typy danych wejściowych są nadmiarowe, a typy danych wyjściowych są niemal zawsze nadmiarowe (ponieważ w większości operacji StableHLO typy danych wyjściowych można określać na podstawie danych wejściowych). Podpis op jest jednak celowo częścią składni StableHLO, aby zapewnić zgodność z MLIR.

Poniżej znajduje się przykład działania, którego mnemotechnika to select_and_scatter. Wykorzystuje 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 zawiera tylko typy jej wartości wejściowych (nie obejmuje typów funkcji wejściowych ani atrybutów, które są dostępne 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 stabilny HLO mają literał i typ, które razem reprezentują wartość StableHLO. Ogólnie ten typ jest częścią stałej składni, chyba że jest jednoznaczny (np. stała wartość logiczna ma typ i1, a stała liczba całkowita może mieć wiele możliwych typów).

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

Stałe logiczne to 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 liczby całkowite przedstawiają wartości całkowite za pomocą ciągów znaków w notacji dziesiętnej lub szesnastkowej. Inne wartości podstawowe, np. binarne lub ósemkowe, nie są obsługiwane. Stałe liczby całkowite mają te 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 liczby zmiennoprzecinkowe reprezentują wartości zmiennoprzecinkowe za pomocą ciągów znaków w zapisie dziesiętnym lub naukowym. Dodatkowo można użyć notacji szesnastkowej, aby bezpośrednio określić bazowe bity w formacie zmiennoprzecinkowym odpowiedniego typu. Stałe liczby zmiennoprzecinkowe mają następujące ograniczenia:

  • (C1) Jeśli używany jest zapis nieszesnastkowy, is_wellformed(float_literal, float_type).
  • (C2) Jeśli używasz zapisu szesnastkowego, 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 przy użyciu list części rzeczywistej (przechodzi ona pierwsza) i części urojonej (następuje jako druga). Na przykład (1.0, 0.0) : complex<f32> oznacza 1.0 + 0.0i, a (0.0, 1.0) : complex<f32> – 0.0 + 1.0i. Kolejność, w jakiej te części będą przechowywane w pamięci, jest określana 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 Tensor reprezentują wartości tensorów za pomocą zagnieżdżonych list określonych w notacji NumPy. Na przykład dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> reprezentuje wartość tensora za pomocą tego mapowania 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 będą przechowywane w pamięci, jest określana przez implementację. Stałe Tensor 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) '>'

Stałe tensorowe reprezentują skwantyzowane wartości tensorów z użyciem tej samej notacji co stałe tensor, przy czym elementy są określone jako stałe typu przechowywania. Kwantyzowane stałe tensorowe mają te 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))

Litery ciągu znaków składają się z bajtów określonych przy użyciu znaków ASCII i sekwencji zmiany znaczenia. Te bajty są niezależne od kodowania, więc ich interpretacja jest definiowana w zależności od implementacji. Litery ciągu znaków mają typ string.

Operacje

abs

Semantyka

Wykonuje operacje absorpcji oparte na elementach na tensorze operand i generuje tensor result. W zależności od typu elementu wykona te działania:

  • W przypadku liczb całkowitych: moduł liczby całkowitej
  • Liczba zmiennoprzecinkowa: abs (IEEE-754).
  • W przypadku liczb zespolonych: moduł sprężystości zespolonej.
  • W przypadku typów skwantyzowanych: dequantize_op_quantize(abs, operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor liczby całkowitej, zmiennoprzecinkowej, zespolonej lub tensora skwantyzowanego według intensywności (C1-C2)

Wyniki

Nazwa Typ Ograniczenia
result tensor liczby całkowitej lub liczby zmiennoprzecinkowej albo tensor skwantyzowany według intensywności (C1-C2)

Ograniczenia

  • (C1) shape(result) = shape(operand).
  • (C2) 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 2 tensorów lhs i rhs na poziomie elementów, generując tensor result. W zależności od typu elementu wykona te działania:

  • W przypadku wartości logicznych: logiczne LUB.
  • W przypadku liczb całkowitych: dodawanie liczby całkowitej.
  • Liczba zmiennoprzecinkowa: addition (IEEE-754).
  • W przypadku liczb zespolonych: dodawanie zespolone.
  • W przypadku typów skwantyzowanych: dequantize_op_quantize(add, lhs, rhs, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor lub tensor kwantyzowany na intensywność (C1)
(I2) rhs tensor lub tensor kwantyzowany na intensywność (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C1)

Ograniczenia

  • (C1) baseline_type(lhs) = baseline_type(rhs) = baseline_type(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 wszelkimi operacjami zależnymi od result. Wykonanie tej operacji nie powoduje żadnych skutków, służy jedynie do ustalenia zależności danych z okresu od result do inputs.

Dane wejściowe

Etykieta Nazwa Typ
(I1) inputs zmienna token

Wyniki

Nazwa Typ
result token

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 tensora operand z każdego procesu wzdłuż all_gather_dim, tworząc 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ś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 ciągu każdego process_group:

  • operands@receiver = [operand@sender for sender in process_group] dla wszystkich receiver w grupie process_group.
  • result@process = concatenate(operands@process, all_gather_dim) dla wszystkich process w grupie process_group.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor lub tensor kwantyzowany na intensywność [C1] (C6)
(I2) all_gather_dim stała typu si64 [C1] (C6)
(I3) replica_groups Dwuwymiarowa stała tensorowa typu si64 (C2–C4)
(I4) channel_id stała typu si64 (C5)
(I5) use_global_device_ids stała typu i1 (C5)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C6)

Ograniczenia

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

Przykłady

// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [[1, 2], [3, 4]]
// %operand@(1, 0): [[5, 6], [7, 8]]
%result = "stablehlo.all_gather"(%operand) {
  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<2x4xi64>
// %result@(0, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result@(1, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]

Więcej przykładów

all_reduce

Semantyka

W ramach każdej grupy procesów w siatce procesów StableHLO zastosuj funkcję redukcji computation do wartości tensora operand z każdego procesu i wygeneruje 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ś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 ciągu każdego process_group:

  • result@process[result_index] = exec(schedule) w przypadku drzewa binarnego, schedule, gdzie:
    • exec(node) = computation(exec(node.left), exec(node.right))
    • exec(leaf) = leaf.value
  • schedule to zdefiniowane w implementacji drzewo binarne, którego wartość przemierzania w kolejności wynosi to_destination_type(operands@process_group...[result_index], type(func_inputs(computation)[0])).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor lub tensor kwantyzowany na intensywność [C5] [C6]
(I2) replica_groups zmienna jednowymiarowa stałych tensorów typu si64 (C1–C3)
(I3) channel_id stała typu si64 (C4)
(I4) use_global_device_ids stała typu i1 (C4)
(I5) computation funkcja (C5)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C6-C7)

Ograniczenia

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

Przykłady

// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [1, 2, 3, 4]
// %operand@(1, 0): [5, 6, 7, 8]
%result = "stablehlo.all_reduce"(%operand) ({
  ^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_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor<i64>) -> tensor<i64>
// %result@(0, 0): [6, 8, 10, 12]
// %result@(1, 0): [6, 8, 10, 12]

Więcej przykładów

all_to_all

Semantyka

W ramach każdej grupy procesów w siatce procesów StableHLO dzieli wartości tensora operand wzdłuż split_dimension na części, rozprasza podzielone części między procesy, łączy rozproszone części wzdłuż concat_dimension i generuje 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śli channel_id <= 0.
  • cross_partition(replica_groups), jeśli channel_id > 0.

Następnie w ciągu każdego process_group:

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

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor lub tensor kwantyzowany na intensywność (C1-C3), (C9)
(I2) split_dimension stała typu si64 [C1], [C2], [C9]
(I3) concat_dimension stała typu si64 (C3) (C9)
(I4) split_count stała typu si64 [C2], [C4], [C8], [C9]
(I5) replica_groups Dwuwymiarowa stała tensorowa typu si64 (C5-C8)
(I6) channel_id stała typu si64

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C9)

Ograniczenia

  • (C1) 0 <= split_dimension < rank(operand).
  • (C2) dim(operand, split_dimension) % split_count = 0.
  • (C3) 0 <= concat_dimension < rank(operand).
  • (C4) 0 < split_count.
  • (C5) is_unique(replica_groups).
  • (C6) size(replica_groups) jest zdefiniowany jako:
    • num_replicas, jeśli używana jest wartość cross_replica.
    • num_partitions, jeśli używana jest wartość cross_partition.
  • (C7) 0 <= replica_groups < size(replica_groups).
  • (C8) dim(replica_groups, 1) = split_count.
  • (C9) type(result) = type(operand) z wyjątkiem:
    • dim(result, split_dimension) = dim(operand, split_dimension) / split_count.
    • dim(result, concat_dimension) = dim(operand, concat_dimension) * split_count.

Przykłady

// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [[1, 2, 3, 4],
//                   [5, 6, 7, 8]]
// %operand@(1, 0): [[9, 10, 11, 12],
//                   [13, 14, 15, 16]]
%result = "stablehlo.all_to_all"(%operand) {
  split_dimension = 1 : i64,
  concat_dimension = 0 : i64,
  split_count = 2 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
} : (tensor<2x4xi64>) -> tensor<4x2xi64>
// %result@(0, 0): [[1, 2],
//                  [5, 6],
//                  [9, 10],
//                  [13, 14]]
// %result@(1, 0): [[3, 4],
//                  [7, 8],
//                  [11, 12],
//                  [15, 16]]

Więcej przykładów

i

Semantyka

Wykonuje parametr ORAZ z uwzględnieniem 2 tensorów lhs i rhs, generując tensor result. W zależności od typu elementu wykona te działania:

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

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor typu logicznego lub liczby całkowitej (C1)
(I2) rhs tensor typu logicznego lub liczby całkowitej (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu logicznego lub liczby całkowitej (C1)

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

atan2

Semantyka

Wykonuje operacje atan2 z wykorzystaniem elementów na tensorze lhs i rhs oraz generuje tensor result. W zależności od typu elementu wykona te działania:

  • Liczba zmiennoprzecinkowa: atan2 (IEEE-754).
  • W przypadku liczb zespolonych: atan2 zespolony.
  • W przypadku typów skwantyzowanych: dequantize_op_quantize(atan2, lhs, rhs, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)
(I2) rhs tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

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 batch_norm_training z propagacją wsteczną z grad_output i generuje intensywność grad_operand, grad_scale i grad_offset. Bardziej formalnie ta operacja może być przedstawiona jako rozkład na istniejące operacje 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 skwantyzowanych wykonuje 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

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C1-C3), (C5)
(I2) scale Jednowymiarowy tensor typu zmiennoprzecinkowego lub kwantyzowanego na intensywność [C2], [C4], [C5]
(I3) mean Jednowymiarowy tensor typu zmiennoprzecinkowego lub kwantyzowanego na intensywność (C2) (C4)
(I4) variance Jednowymiarowy tensor typu zmiennoprzecinkowego lub kwantyzowanego na intensywność (C2) (C4)
(I5) grad_output tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C2) (C3)
(I6) epsilon stała typu f32
(I7) feature_index stała typu si64 (C1), [C5]

Wyniki

Nazwa Typ Ograniczenia
grad_operand tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C2) (C3)
grad_scale Jednowymiarowy tensor typu zmiennoprzecinkowego lub kwantyzowanego na intensywność (C2) (C4)
grad_offset Jednowymiarowy tensor typu zmiennoprzecinkowego lub kwantyzowanego na intensywność (C2) (C4)

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) 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 generuje tensor result. Bardziej formalnie ta operacja może być przedstawiona jako rozkład istniejących operacji StableHLO przy użyciu 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 skwantyzowanych wykonuje 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

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C1–C7)
(I2) scale Jednowymiarowy tensor typu zmiennoprzecinkowego lub kwantyzowanego na intensywność (C2) (C3)
(I3) offset Jednowymiarowy tensor typu zmiennoprzecinkowego lub kwantyzowanego na intensywność (C2) (C4)
(I4) mean Jednowymiarowy tensor typu zmiennoprzecinkowego lub kwantyzowanego na intensywność (C5)
(I5) variance Jednowymiarowy tensor typu zmiennoprzecinkowego lub kwantyzowanego na intensywność (C2) (C6)
(I6) epsilon stała typu f32
(I7) feature_index stała typu si64 [C1] (C3-C6)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C2) (C7)

Ograniczenia

  • (C1) 0 <= feature_index < rank(operand).
  • (C2) operand, scale, offset, mean, variance i result mają tę samą wartość 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 z wyjątkiem wymiaru feature_index oraz normalizuje tensor operand, tworząc tensory output, batch_mean i batch_var. Bardziej formalnie ta operacja może być przedstawiona jako rozkład istniejących operacji StableHLO przy użyciu składni Pythona w ten sposób:

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 skwantyzowanych wykonuje 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

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C1)
(I2) scale Jednowymiarowy tensor zmiennoprzecinkowej lub kwantyzowany na intensywność (C2) (C3)
(I3) offset Jednowymiarowy tensor zmiennoprzecinkowej lub kwantyzowany na intensywność (C2) (C4)
(I4) epsilon stała typu f32 [C1] (C3-C6)
(I5) feature_index stała typu si64 [C1] (C3-C6)

Wyniki

Nazwa Typ Ograniczenia
output tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C7)
batch_mean Jednowymiarowy tensor zmiennoprzecinkowej lub kwantyzowany na intensywność (C2) (C5)
batch_var Jednowymiarowy tensor zmiennoprzecinkowej lub kwantyzowany na intensywność (C2) (C6)

Ograniczenia

  • (C1) 0 <= feature_index < rank(operand).
  • (C2) operand, scale, offset, batch_mean, batch_var i output mają tę samą wartość 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

Przeprowadza operację transmisji bitów na tensorze operand i generuje tensor result, w którym bity całego tensora operand są ponownie interpretowane za pomocą typu tensora result.

Bardziej oficjalnie. Według danych 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 zachowanie jest definiowane za pomocą implementacji, ponieważ dokładna reprezentacja tensorów jest zdefiniowana za pomocą implementacji, a dokładna reprezentacja typów elementów również jest definiowana przez implementację.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor lub tensor skwantyzowany (C1-C2)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor skwantyzowany (C1-C2)

Ograniczenia

  • (C1) Dla danych 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 we tensorze operand i generuje tensor result. Bardziej formalnie result[result_index] = operand[operand_index], gdzie wszystkie 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

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor lub tensor skwantyzowany (C1-C2), (C5-C6)
(I2) broadcast_dimensions Jednowymiarowa stała tensorowa typu si64 (C2–C6)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor skwantyzowany (C1), [C3], (C5-C6)

Ograniczenia

  • (C1) Wartość element_type(result) jest obliczana przez:
    • element_type(operand), jeśli !is_per_axis_quantized(operand).
    • element_type(operand) z wyjątkiem tych wartości: quantization_dimension(operand), scales(operand) i zero_points(operand) mogą się różnić od odpowiedzi quantization_dimension(result), scales(result) i zero_points(result) z innego powodu.
  • (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

zgłoszenie

Semantyka

Produkuje dane wyjściowe podczas wykonywania 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śli 0 <= index < size(branches).
  • W przeciwnym razie selected_branch = branches[-1].

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) index Tensor 0-wymiarowy typu si32
(I2) branches zmienna liczba funkcji (C1–C4)

Wyniki

Nazwa Typ Ograniczenia
results zmienna tensorów, tensorów kwantowych lub tokenów (C4)

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 generuje tensor result. W zależności od typu elementu wykona te działania:

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

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

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 tensor ceila operand i generuje tensor result. Implementuje operację roundToIntegralTowardPositive ze specyfikacji IEEE-754. W przypadku typów skwantyzowanych wykonuje dequantize_op_quantize(ceil, operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C1)

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 dla wsadu matryc.

Ogólnie rzecz biorąc, dla wszystkich i w tabeli index_space(result) result[i0, ..., iR-3, :, :] to rozkład Cholesky’ego a[i0, ..., iR-3, :, :] w postaci macierzy dolnego trójkąta (jeśli lower to true) lub górnej trójkąta (jeśli lower to false). Wartości wyjściowe w przeciwległym trójkącie, tj. odpowiednio ścisły górny trójkąt lub ścisły dolny trójkąt, są zdefiniowane w ramach implementacji.

Jeśli i ma tablicę wejściową, która nie jest macierą nieoznaczoną hermicką, działanie jest niezdefiniowane.

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

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) a tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1–C3)
(I2) lower 0-wymiarowa stała tensorowa typu i1

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

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 generuje 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 skwantyzowanych wykonuje dequantize_op_quantize(clamp, min, operand, max, type(result)).

Ustalanie kolejności liczb zespolonych może być zaskakujące, dlatego w przyszłości planujemy wycofanie w tej operacji obsługi liczb zespolonych (#560).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) min tensor lub tensor kwantyzowany na intensywność (C1) (C3)
(I2) operand tensor lub tensor kwantyzowany na intensywność (C1–C4)
(I3) max tensor lub tensor kwantyzowany na intensywność (C2) (C3)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C4)

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ść 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śli channel_id <= 0.
  • cross_partition(replica_groups), jeśli channel_id > 0.

Następnie result@process przyznaje:

  • operand@process_groups[i, 0], jeśli istnieje dyrektywa i, która oznacza, że proces jest w elemencie process_groups[i].
  • broadcast_in_dim(constant(0, element_type(result)), [], type(result)) w przeciwnym razie.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor (C3)
(I2) replica_groups zmienna jednowymiarowa stałych tensorów typu si64 (C1) (C2)
(I3) channel_id stała typu si64

Wyniki

Nazwa Typ Ograniczenia
result tensor (C3)

Ograniczenia

  • (C1) is_unique(replica_groups).
  • (C2) 0 <= replica_groups < N, gdzie N jest zdefiniowana jako:
    • num_replicas, jeśli używana jest wartość cross_replica.
    • num_partitions, jeśli używana jest wartość 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 ramach każdej grupy procesów w siatce procesów StableHLO wysyła wartość tensora operand z procesu źródłowego do procesu docelowego i generuje 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(source_target_pairs), jeśli channel_id <= 0.
  • cross_partition(source_target_pairs), jeśli channel_id > 0.

Następnie result@process przyznaje:

  • operand@process_groups[i, 0], jeśli istnieje 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

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor lub tensor kwantyzowany na intensywność (C5)
(I2) source_target_pairs Dwuwymiarowa stała tensorowa typu si64 (C1–C4)
(I3) channel_id stała typu si64

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C1)

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 zdefiniowana jako:
    • num_replicas, jeśli używana jest wartość cross_replica.
    • num_partitions, jeśli używana jest wartość 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 na podstawie elementów zgodnie z zasadami comparison_direction i compare_type oraz generuje tensor result.

Wartości comparison_direction i compare_type mają tę semantykę:

W przypadku elementów zawierających wartości logiczne i całkowite:

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

W przypadku elementów zmiennoprzecinkowych compare_type = FLOAT operacja implementuje te operacje IEEE-754:

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

W przypadku elementów zmiennoprzecinkowych compare_type = TOTALORDER operacja używa kombinacji operacji totalOrder i compareQuietEqual z IEEE-754. Wygląda na to, że ta funkcja nie jest używana, więc w przyszłości planujemy ją usunąć (#584).

W przypadku elementów złożonych następuje porównanie leksykograficzne par (real, imag) z użyciem podanych parametrów comparison_direction i compare_type. Ustalanie kolejności liczb zespolonych może być zaskakujące, dlatego w przyszłości planujemy wycofanie obsługi liczb zespolonych, gdy comparison_direction to GE, GT, LE lub LT (#560).

W przypadku typów skwantyzowanych skuteczność wynosi dequantize_compare(lhs, rhs, comparison_direction).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor lub tensor kwantyzowany na intensywność (C1–C3)
(I2) rhs tensor lub tensor kwantyzowany na intensywność (C1-C2)
(I3) comparison_direction wyliczenie: EQ, NE, GE, GT, LE i LT
(I4) compare_type wyliczenie: FLOAT, TOTALORDER, SIGNED i UNSIGNED (C3)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu logicznego (C2)

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

Przekształca wartości związane z elementami na wartość zespoloną z pary wartości rzeczywistych i urojonych (lhs i rhs), a potem generuje tensor result.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor typu f32 lub f64 (C1–C3)
(I2) rhs tensor typu f32 lub f64 (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zespolonego (C2) (C3)

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

concatenate

Semantyka

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

  1. id = d0 + ... + dk-1 + kd.
  2. d równa się dimension, a d0 (...) to d rozmiar wymiarów inputs.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) inputs zmienna liczba tensorów lub tensorów kwantyzowanych na intensywność (C1–C6)
(I2) dimension stała typu si64 [C2], [C4], [C6]

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C5-C6)

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 na podstawie stałej value.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) value stała (C1)

Wyniki

Nazwa Typ Ograniczenia
output tensor lub tensor skwantyzowany (C1)

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 z jednego typu elementu konwersję na inny za pomocą tensora operand i generuje tensor result.

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

W przypadku konwersji uwzględniających liczbę całkowitą na liczbę całkowitą, liczbę całkowitą na liczbę zmiennoprzecinkową lub zmiennoprzecinkową, o ile wartość źródłowa może być dokładnie reprezentowana w typie miejsca docelowego, wartość wyniku jest właśnie dokładną reprezentacją. W przeciwnym razie zachowanie jest do ustalenia (#180).

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

Konwersja polegająca na przekształceniu części ze złożonym i złożonym działa tak samo jak konwersje z liczby zmiennoprzecinkowej na zmiennoprzecinkowe w przypadku przekształcania części rzeczywistych i urojonych.

W przypadku konwersji complex-to-any-other-type i complex-to-any-other-type źródłowa wartość urojona jest ignorowana odpowiednio, a wartość urojona miejsca docelowego jest wyzerowana. Przekształcanie części rzeczywistej następuje po konwersjach z liczbą zmiennoprzecinkową.

Zasadniczo operacja ta może wyrażać dekwantyzację (konwersję tensorów kwantyzowanych na tensory zwykłe), kwantyzację (konwersję ze zwykłych tensorów na tensory skwantyzowane) i rekwantyzację (konwersję między kwantyzowanymi tensorami), ale obecnie mamy do tego specjalne operacje – uniform_dequantize w pierwszym i trzecim przypadku użycia, a w trzecim przypadku użycia: uniform_quantize. W przyszłości te 2 operacje mogą zostać scalone w usługę convert (#1576).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor (C1)

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 lhs i wycinkami rhs oraz daje result. Na diagramie poniżej widać, jak za pomocą konkretnego przykładu obliczamy elementy w atrybutach result z wartości lhs i rhs.

Bardziej formalnie rozważ takie przekształcenie danych wejściowych w kierunku funkcji lhs, aby móc wyrazić okno 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).

Przygotowując tę strukturę, wykorzystuje się następujące 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]). Wygląda na to, że ta funkcja nie jest używana, więc w przyszłości planujemy ją usunąć (#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 skwantyzowanych wykonuje 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)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor lub tensor kwantyzowany na intensywność (C1), (C10-C11), (C14) (C25), (C27-C30)
(I2) rhs tensor lub tensor skwantyzowany (C1), (C14-C16), (C25), (C27-C32)
(I3) window_strides Jednowymiarowa stała tensorowa typu si64 (C2-C3), (C25)
(I4) padding Dwuwymiarowa stała tensorowa typu si64 [C4] (C25)
(I5) lhs_dilation Jednowymiarowa stała tensorowa typu si64 (C5-C6), (C25)
(I6) rhs_dilation Jednowymiarowa stała tensorowa typu si64 (C7-C8), (C25)
(I7) window_reversal Jednowymiarowa stała tensorowa typu i1 (C9)
(I8) input_batch_dimension stała typu si64 (C10), (C13) i (C25)
(I9) input_feature_dimension stała typu si64 (C11), (C13-C14)
(I10) input_spatial_dimensions Jednowymiarowa stała tensorowa typu si64 (C12), (C13) i (C25)
(I11) kernel_input_feature_dimension stała typu si64 (C14) (C18)
(I12) kernel_output_feature_dimension stała typu si64 (C15-C16), (C18), (C25), (C32)
(I13) kernel_spatial_dimensions Jednowymiarowa stała tensorowa typu si64 (C17-C18), (C25)
(I14) output_batch_dimension stała typu si64 (C20) (C25)
(I15) output_feature_dimension stała typu si64 (C20), (C25) i (C33)
(I16) output_spatial_dimensions Jednowymiarowa stała tensorowa typu si64 (C19-C20), (C25)
(I17) feature_group_count stała typu si64 (C11), (C14), (C16), (C21) i (C23)
(I18) batch_group_count stała typu si64 (C10), (C15), (C22), (C23) i (C25)
(I19) precision_config zmienna liczba wyliczeń DEFAULT, HIGH i HIGHEST (C24)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor skwantyzowany (C25-C28), (C30-C31), (C33)

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) Dla danych 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) Dla danych 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) Dla danych 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) został 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.
    • W przeciwnym razie num_windows, 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 niekwantyzowanych:
    • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
  • Jeśli operacja używa tensorów skwantyzowanych:
    • (C28) is_quantized_tensor(lhs) and is_quantized_tensor(rhs) and is_quantized_tensor(result).
    • (C29) storage_type(lhs) = storage_type(rhs).
    • (C30) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
    • (C31) Jeśli is_per_tensor_quantized(rhs), to is_per_tensor_quantized(result).
    • (C32) Jeśli is_per_axis_quantized(rhs), to quantization_dimension(rhs) = kernel_output_feature_dimension.
    • (C33) Jeśli is_per_axis_quantized(result), to quantization_dimension(result) = output_feature_dimension.

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 = dense<4> : tensor<2xi64>,
  padding = dense<0> : tensor<2x2xi64>,
  lhs_dilation = dense<2> : tensor<2xi64>,
  rhs_dilation = dense<1> : tensor<2xi64>,
  window_reversal = dense<false> : tensor<2xi1>,
  // 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]>,
  feature_group_count = 1 : i64,
  batch_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi32>, tensor<3x3x1x1xi32>) -> tensor<1x2x2x1xi32>
// %result: [[
//            [[10], [26]],
//            [[46], [62]]
//          ]]

cosinus

Semantyka

Wykonuje operację cosinusową elementu na tensorze operand i generuje tensor result. W zależności od typu elementu wykona te działania:

  • Liczba zmiennoprzecinkowa: cos (IEEE-754).
  • W przypadku liczb zespolonych: cosinus zespolony.
  • W przypadku typów skwantyzowanych: dequantize_op_quantize(cosine, operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

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 zliczanie na początku liczby zerowych bitów na początku tensora operand i generuje tensor result.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu liczby całkowitej (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu liczby całkowitej (C1)

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ą w implementacji operację call_target_name, która pobiera inputs i called_computations oraz generuje results. has_side_effect, backend_config i api_version mogą być używane do przekazywania dodatkowych metadanych określonych w implementacji.

W tej chwili ta operacja zawiera dość nieuporządkowany zbiór metadanych, który odzwierciedla organiczną ewolucję jej odpowiednika w kompilatorze XLA. W przyszłości planujemy ujednolicić te metadane (#741).

Dane wejściowe

Etykieta Nazwa Typ
(I1) inputs zmienna liczba wartości
(I2) call_target_name stała typu string
(I3) has_side_effect stała typu i1
(I4) backend_config stała typu string
(I5) api_version stała typu si32
(I6) called_computations zmienna stałych typu string

Wyniki

Nazwa Typ
results zmienna liczba wartości

Przykłady

%results = "stablehlo.custom_call"(%input0) {
  call_target_name = "foo",
  has_side_effect = false,
  backend_config = "bar",
  api_version = 1 : i32,
  called_computations = [@foo]
} : (tensor<f64>) -> tensor<f64>

dzielenie

Semantyka

Wykonuje dzielenie tensorów dzielnika lhs i dzielnika rhs oraz generuje tensor result. W zależności od typu elementu wykona te działania:

  • W przypadku liczb całkowitych: dzielenie, które daje iloraz algebraiczny z odrzuconą częścią ułamkową.
  • Liczba zmiennoprzecinkowa: division (IEEE-754).
  • W przypadku liczb zespolonych: dzielenie zespolone.
  • W przypadku typów skwantyzowanych:
    • dequantize_op_quantize(divide, lhs, rhs, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor liczby całkowitej, zmiennoprzecinkowej lub zespolonej albo tensora skwantyzowanego według intensywności (C1)
(I2) rhs tensor liczby całkowitej, zmiennoprzecinkowej lub zespolonej albo tensora skwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu całkowitej, zmiennoprzecinkowej, zespolonej lub tensora skwantyzowanego według intensywności (C1)

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 wycinkami rhs, a następnie generuje tensor result.

Bardziej oficjalnie, 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 skwantyzowanych wykonuje 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)).

Określa tylko semantykę kwantyzacji na intensywność. Trwa kwantyzacja dla poszczególnych osi (#1574). W przyszłości możemy też rozważyć dodanie obsługi kwantyzacji hybrydowej (#1575).

precision_config określa kompromis między szybkością a dokładnością obliczeń w backendach akceleratora. Może to być jedna z tych sytuacji (na razie semantyka tych wartości wyliczeniowych jest nieokreślona, ale planujemy to zmienić w #755):

  • DEFAULT: Najszybsze obliczenie, ale najmniej dokładne przybliżenie liczby oryginalnej.
  • HIGH: wolniejsze obliczanie, ale dokładniejsze przybliżanie wartości oryginalnej.
  • HIGHEST: najwolniejsze obliczanie, ale z najwierniejszym przybliżeniem względem pierwotnej liczby.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor lub tensor kwantyzowany na intensywność (C5-C6), (C9-C10), (C12-C16)
(I2) rhs tensor lub tensor kwantyzowany na intensywność (C7-C10), (C12)
(I3) lhs_batching_dimensions Jednowymiarowa stała tensorowa typu si64 [C1], [C3], [C5], [C12]
(I4) rhs_batching_dimensions Jednowymiarowa stała tensorowa typu si64 [C1], [C4], [C7], [C9]
(I5) lhs_contracting_dimensions Jednowymiarowa stała tensorowa typu si64 [C2], [C3], [C6], [C10]
(I6) rhs_contracting_dimensions Jednowymiarowa stała tensorowa typu si64 [C2], [C4], [C8], [C10]
(I7) precision_config zmienna liczba wyliczeń DEFAULT, HIGH i HIGHEST (C11)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C12), (C14) i (C16)

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 niekwantyzowanych:
    • (C13) element_type(lhs) = element_type(rhs).
  • Jeśli operacja używa tensorów skwantyzowanych:
    • (C14) is_quantized(lhs) and is_quantized(rhs) and is_quantized(result).
    • (C15) storage_type(lhs) = storage_type(rhs).
    • (C16) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
    • (C17) zero_points(rhs) = 0.

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>]
} : (tensor<2x2x2xi64>, tensor<2x2x2xi64>) -> tensor<2x2x2xi64>
// %result: [
//           [[1, 2],
//            [3, 4]],
//           [[5, 6],
//            [7, 8]]
//          ]

Więcej przykładów

dynamic_slice

Semantyka

Wyodrębnia wycinek z obiektu operand za pomocą dynamicznie obliczanych indeksów początkowych i generuje tensor result. start_indices zawiera początkowe indeksy wycinka dla każdego wymiaru, który może zostać skorygowany, a slice_sizes zawiera rozmiary wycinka każdego wymiaru. Bardziej oficjalnie, result[result_index] = operand[operand_index], gdzie:

  • adjusted_start_indices = clamp(0, start_indices, shape(operand) - slice_sizes).
  • operand_index = adjusted_start_indices + result_index.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor lub tensor kwantyzowany na intensywność (C1), [C2], (C4)
(I2) start_indices zmienna tensorowa 0-wymiarowego typu liczby całkowitej (C2) (C3)
(I3) slice_sizes Jednowymiarowa stała tensorowa typu si64 [C2], [C4], [C5]

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C1), [C5]

Ograniczenia

  • (C1) element_type(operand) = element_type(result).
  • (C2) size(start_indices) = size(slice_sizes) = rank(operand).
  • (C3) same(type(start_indices...)).
  • (C4) 0 <= slice_sizes <= shape(operand).
  • (C5) shape(result) = slice_sizes.

Przykłady

// %operand: [
//            [0, 0, 1, 1],
//            [0, 0, 1, 1],
//            [0, 0, 0, 0],
//            [0, 0, 0, 0]
//           ]
// %start_indices0: -1
// %start_indices1: 3
%result = "stablehlo.dynamic_slice"(%operand, %start_indices0, %start_indices1) {
  slice_sizes = dense<[2, 2]> : tensor<2xi64>
} : (tensor<4x4xi32>, tensor<i64>, tensor<i64>) -> tensor<2x2xi32>
// %result: [
//           [1, 1],
//           [1, 1]
//          ]

Więcej przykładów

dynamic_update_slice

Semantyka

Generuje tensor result, który jest równy tensorowi operand, z tą różnicą, że wycinek zaczynający się od start_indices jest aktualizowany wartościami z parametru update. Dokładniej rzecz ujmując, result[result_index] to:

  • update[update_index], jeśli 0 <= update_index < shape(update), gdzie:
    • adjusted_start_indices = clamp(0, start_indices, shape(operand) - shape(update)).
    • update_index = result_index - adjusted_start_indices.
  • W przeciwnym razie operand[result_index].

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor lub tensor kwantyzowany na intensywność (C1-C4), (C6)
(I2) update tensor lub tensor kwantyzowany na intensywność [C2], [C3], [C6]
(I3) start_indices zmienna tensorowa 0-wymiarowego typu liczby całkowitej [C4] (C5)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C1)

Ograniczenia

  • (C1) type(operand) = type(result).
  • (C2) element_type(update) = element_type(operand).
  • (C3) rank(update) = rank(operand).
  • (C4) size(start_indices) = rank(operand).
  • (C5) same(type(start_indices...)).
  • (C6) 0 <= shape(update) <= shape(operand).

Przykłady

// %operand: [
//            [1, 1, 0, 0],
//            [1, 1, 0, 0],
//            [1, 1, 1, 1],
//            [1, 1, 1, 1]
//           ]
// %update: [
//           [1, 1],
//           [1, 1]
//          ]
// %start_indices0: -1
// %start_indices1: 3
%result = "stablehlo.dynamic_update_slice"(%operand, %update, %start_indices0, %start_indices1)
  : (tensor<4x4xi32>, tensor<2x2xi32>, tensor<i64>, tensor<i64>) -> tensor<4x4xi32>
// %result: [
//           [1, 1, 1, 1],
//           [1, 1, 1, 1],
//           [1, 1, 1, 1],
//           [1, 1, 1, 1]
//          ]

Więcej przykładów

wykładniczo

Semantyka

Wykonuje operację wykładniczą dotyczącą elementów na tensorze operand i generuje tensor result. W zależności od typu elementu wykona te działania:

  • Liczba zmiennoprzecinkowa: exp (IEEE-754).
  • W przypadku liczb zespolonych: funkcja wykładnicza zespolona.
  • W przypadku typów skwantyzowanych: dequantize_op_quantize(exponential, operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Ograniczenia

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

Przykłady

// %operand: [[0.0, 1.0], [2.0, 3.0]]
%result = "stablehlo.exponential"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
// %result: [[1.0, 2.7182818284590451], [7.3890560989306504, 20.085536923187668]]

Więcej przykładów

exponential_minus_one

Semantyka

Wykonuje operację wykładniczą minus 1 zależną od elementów na tensorze operand i generuje tensor result. W zależności od typu elementu wykona te działania:

  • Liczba zmiennoprzecinkowa: expm1 (IEEE-754).
  • W przypadku liczb zespolonych: wykładnik wykładniczy zespolony minus 1.
  • W przypadku typów skwantyzowanych: dequantize_op_quantize(exponential_minus_one, operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Ograniczenia

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

Przykłady

// %operand: [0.0, 1.0]
%result = "stablehlo.exponential_minus_one"(%operand) : (tensor<2xf64>) -> tensor<2xf64>
// %result: [0.0, 1.71828187]

Więcej przykładów

FFT

Semantyka

Wykonuje transformatę Fouriera do przodu i odwrotnego dla rzeczywistych i złożonych danych wejściowych/wyjściowych.

fft_type należy do jednej z tych wartości:

  • FFT: przekształcanie zespolonego na złożony.
  • IFFT: odwrotność funkcji zespolonej do złożonej.
  • RFFT: przekształcanie funkcji FFT od rzeczywistej do złożonej.
  • IRFFT: funkcja odwrotna funkcji FFT z funkcji rzeczywistej do złożonej (tj. funkcja przyjmuje wartości zespolone, zwraca wartość rzeczywistą).

Bardziej formalnie funkcja fft, która przyjmuje jednowymiarowe tensory typów złożonych jako dane wejściowe, generuje jednowymiarowe tensory tego samego typu co dane wyjściowe i oblicza dyskretną transformatę Fouriera:

W przypadku funkcji fft_type = FFT wartość result jest zdefiniowana jako końcowy wynik serii obliczeń L, gdzie L = size(fft_length). Na przykład w przypadku L = 3:

  • result1[i0, ..., :] = fft(operand[i0, ..., :]).
  • result2[i0, ..., :, iR-1] = fft(result1[i0, ..., :, iR-1]).
  • result[i0, ..., :, iR-2, iR-1] = fft(result2[i0, ..., :, iR-2, iR-1]).

Ponadto funkcja ifft ma podpis tego samego typu i oblicza odwrotność fft:

W przypadku fft_type = IFFT funkcja result jest zdefiniowana jako odwrotność obliczeń fft_type = FFT. Na przykład w przypadku L = 3:

  • result1[i0, ..., :, iR-2, iR-1] = ifft(operand[i0, ..., :, iR-2, iR-1]).
  • result2[i0, ..., :, iR-1] = ifft(result1[i0, ..., :, iR-1]).
  • result[i0, ..., :] = ifft(result2[i0, ..., :]).

Poza tym funkcja rfft, która pobiera tensory jednowymiarowe typów zmiennoprzecinkowych, generuje jednowymiarowe tensory złożone tego samego semantyki zmiennoprzecinkowej. Działa w ten sposób:

  • rfft(real_operand) = truncated_result, gdzie
  • complex_operand... = (real_operand..., 0.0).
  • complex_result = fft(complex_operand).
  • truncated_result = complex_result[:(rank(complex_result) / 2 + 1)].

(Gdy dyskretna transformata Fouriera jest obliczana dla rzeczywistych operandów, pierwsze elementy N/2 + 1 jednoznacznie określają resztę wyniku, więc wynik funkcji rfft jest obcięty, aby uniknąć przetwarzania zbędnych elementów).

W przypadku funkcji fft_type = RFFT wartość result jest zdefiniowana jako końcowy wynik serii obliczeń L, gdzie L = size(fft_length). Na przykład w przypadku L = 3:

  • result1[i0, ..., :] = rfft(operand[i0, ..., :]).
  • result2[i0, ..., :, iR-1] = fft(result1[i0, ..., :, iR-1]).
  • result[i0, ..., :, iR-2, iR-1] = fft(result2[i0, ..., :, iR-2, iR-1]).

Na koniec z uwzględnieniem funkcji irfft, która ma ten sam typ podpisu i oblicza odwrotność rfft:

W przypadku fft_type = IRFFT funkcja result jest zdefiniowana jako odwrotność obliczeń fft_type = RFFT. Na przykład w przypadku L = 3:

  • result1[i0, ..., :, iR-2, iR-1] = ifft(operand[i0, ..., :, iR-2, iR-1]).
  • result2[i0, ..., :, iR-1] = ifft(result1[i0, ..., :, iR-1]).
  • result[i0, ..., :] = irfft(result2[i0, ..., :]).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub zespolonego [C1], [C2], [C4], [C5]
(I2) fft_type wyliczenie: FFT, IFFT, RFFT i IRFFT (C2) (C5)
(I3) fft_length Jednowymiarowa stała tensorowa typu si64 (C1), [C3], (C4)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub zespolonego [C2], [C4], [C5]

Ograniczenia

  • (C1) size(fft_length) <= rank(operand).
  • (C2) Relacje między typami elementów operand i result są różne:
    • Jeśli fft_type = FFT, element_type(operand) i element_type(result) mają ten sam typ zespolony.
    • Jeśli fft_type = IFFT, element_type(operand) i element_type(result) mają ten sam typ zespolony.
    • Jeśli fft_type = RFFT to typ zmiennoprzecinkowy, element_type(operand) jest typem zmiennoprzecinkowym, a element_type(result) to typ złożony o tej samej semantyce liczby zmiennoprzecinkowej.
    • Jeśli fft_type = IRFFT, element_type(operand) jest typem złożonym, a element_type(result) jest typem zmiennoprzecinkowym o tej samej semantyce liczby zmiennoprzecinkowej.
  • (C3) 1 <= size(fft_length) <= 3.
  • (C4) Jeśli między elementami operand i result występuje tensor real typu zmiennoprzecinkowego, występuje wartość shape(real)[-size(fft_length):] = fft_length.
  • (C5) shape(result) = shape(operand) z wyjątkiem tych sytuacji:
    • Jeśli fft_type = RFFT, dim(result, -1) = dim(operand, -1) = 0 ? 0 : dim(operand, -1) / 2 + 1.
    • Jeśli fft_type = IRFFT, dim(operand, -1) = dim(result, -1) = 0 ? 0 : dim(result, -1) / 2 + 1.

Przykłady

// %operand: [(1.0, 0.0), (0.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
%result = "stablehlo.fft"(%operand) {
  fft_type = #stablehlo<fft_type FFT>,
  fft_length = dense<4> : tensor<1xi64>
} : (tensor<4xcomplex<f32>>) -> tensor<4xcomplex<f32>>
// %result: [(1.0, 0.0), (1.0, 0.0), (1.0, 0.0), (1.0, 0.0)]

piętro

Semantyka

Wykonuje rzutowanie na element operand tensora i generuje tensor result. Implementuje operację roundToIntegralTowardNegative ze specyfikacji IEEE-754. W przypadku typów skwantyzowanych wykonuje dequantize_op_quantize(floor, operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C1)

Ograniczenia

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

Przykłady

// %operand: [-0.8166, -0.2530, 0.2530, 0.8166, 2.0]
%result = "stablehlo.floor"(%operand) : (tensor<5xf32>) -> tensor<5xf32>
// %result: [-1.0, -1.0, 0.0, 0.0, 2.0]

Więcej przykładów

zbierać

Semantyka

Pobiera wycinki z tensora operand z przesunięć określonych w start_indices i tworzy tensor result.

Na diagramie poniżej widać, jak za pomocą konkretnego przykładu mapujemy elementy w tabeli result na elementy w komponencie operand. Diagram wybiera kilka przykładowych indeksów result i dokładnie wyjaśnia, którym indeksom operand odpowiadają.

Bardziej formalnie, result[result_index] = operand[operand_index], gdzie:

  • batch_dims = [d for d in axes(result) and d not in offset_dims].
  • batch_index = result_index[batch_dims...].
  • start_index został zdefiniowany jako:
    • start_indices[bi0, ..., :, ..., biN], gdzie bi to poszczególne elementy w batch_index, a : jest wstawiony w indeksie index_vector_dim, jeśli index_vector_dim < rank(start_indices).
    • W przeciwnym razie [start_indices[batch_index]].
  • Przez d_operand w: axes(operand),
    • full_start_index[d_operand] = clamp(start_index[d_start], 0, dim(operand, d_operand) - slice_sizes[d_operand]) jeśli d_operand = start_index_map[d_start].
    • W przeciwnym razie full_start_index[d_operand] = 0.
  • offset_index = result_index[offset_dims...].
  • full_offset_index = [oi0, ..., 0, ..., oiN], gdzie oi to poszczególne elementy w offset_index, a 0 jest wstawiony w indeksach od collapsed_slice_dims.
  • operand_index = full_start_index + full_offset_index.

Jeśli indices_are_sorted ma wartość true, implementacja może przyjąć, że metoda start_indices jest posortowana względem argumentu start_index_map. W przeciwnym razie działanie jest niezdefiniowane. Oficjalnie dla wszystkich i1 < i2 od indices(result) full_start_index(i1) <= full_start_index(i2).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor lub tensor kwantyzowany na intensywność (C1), (C7), (C10-C12), (C14)
(I2) start_indices tensor typu liczby całkowitej (C2), (C3) i (C13)
(I3) offset_dims Jednowymiarowa stała tensorowa typu si64 (C1), (C4-C5), (C13)
(I4) collapsed_slice_dims Jednowymiarowa stała tensorowa typu si64 (C1), (C6-C8), (C13)
(I5) start_index_map Jednowymiarowa stała tensorowa typu si64 [C3], [C9], [C10]
(I6) index_vector_dim stała typu si64 (C2), (C3) i (C13)
(I7) slice_sizes Jednowymiarowa stała tensorowa typu si64 (C8) (C11-C13)
(I8) indices_are_sorted stała typu i1

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność [C5] [C13-C14]

Ograniczenia

  • (C1) rank(operand) = size(offset_dims) + size(collapsed_slice_dims).
  • (C2) 0 <= index_vector_dim <= rank(start_indices).
  • (C3) size(start_index_map) = index_vector_dim < rank(start_indices) ? dim(start_indices, index_vector_dim) : 1.
  • (C4) is_unique(offset_dims) and is_sorted(offset_dims).
  • (C5) 0 <= offset_dims < rank(result).
  • (C6) is_unique(collapsed_slice_dims) and is_sorted(collapsed_slice_dims).
  • (C7) 0 <= collapsed_slice_dims < rank(operand).
  • (C8) slice_sizes[collapsed_slice_dims...] <= 1.
  • (C9) is_unique(start_index_map).
  • (C10) 0 <= start_index_map < rank(operand).
  • (C11) size(slice_sizes) = rank(operand).
  • (C12) 0 <= slice_sizes <= shape(operand).
  • (C13) shape(result) = combine(batch_dim_sizes, offset_dim_sizes), gdzie:
    • batch_dim_sizes = shape(start_indices) z tym wyjątkiem, że rozmiar wymiaru start_indices odpowiadający atrybutowi index_vector_dim nie jest uwzględniony.
    • offset_dim_sizes = shape(slice_sizes) z tym wyjątkiem, że rozmiary podane w polu slice_sizes odpowiadające wartościom collapsed_slice_dims nie są uwzględniane.
    • combine umieszcza batch_dim_sizes na osiach odpowiadających wartościom batch_dims i offset_dim_sizes na osiach odpowiadających osiom offset_dims.
  • (C14) element_type(operand) = element_type(result).

Przykłady

// %operand: [
//            [[1, 2], [3, 4], [5, 6], [7, 8]],
//            [[9, 10],[11, 12], [13, 14], [15, 16]],
//            [[17, 18], [19, 20], [21, 22], [23, 24]]
//           ]
// %start_indices: [
//                  [[0, 0], [1, 0], [2, 1]],
//                  [[0, 1], [1, 1], [0, 2]]
//                 ]
%result = "stablehlo.gather"(%operand, %start_indices) {
  dimension_numbers = #stablehlo.gather<
    offset_dims = [2, 3],
    collapsed_slice_dims = [0],
    start_index_map = [1, 0],
    index_vector_dim = 2>,
  slice_sizes = dense<[1, 2, 2]> : tensor<3xi64>,
  indices_are_sorted = false
} : (tensor<3x4x2xi32>, tensor<2x3x2xi64>) -> tensor<2x3x2x2xi32>
// %result: [
//            [
//              [[1, 2], [3, 4]],
//              [[3, 4], [5, 6]],
//              [[13, 14], [15, 16]]
//            ],
//            [
//              [[9, 10], [11, 12]],
//              [[11, 12], [13, 14]],
//              [[17, 18], [19, 20]]
//            ]
//          ]

Więcej przykładów

get_dimension_size

Semantyka

Tworzy rozmiar podanej wartości dimension elementu operand. Oficjalnie: result = dim(operand, dimension).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor (C1)
(I2) dimension stała typu si64 (C1)

Wyniki

Nazwa Typ
result Tensor 0-wymiarowy typu si32

Ograniczenia

  • (C1) 0 <= dimension < rank(operand).

Przykłady

// %operand: [[1, 2, 3], [4, 5, 6]]
%result = "stablehlo.get_dimension_size"(%operand) {
  dimension = 1 : i64
} : (tensor<2x3xi64>) -> tensor<i32>
// %result: 3

Więcej przykładów

get_tuple_element

Semantyka

Wyodrębnia element w pozycji index krotki operand i generuje result. Oficjalnie: result = operand[index].

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tuple (C1) (C2)
(I2) index stała typu si32 (C1) (C2)

Wyniki

Nazwa Typ Ograniczenia
result dowolny obsługiwany typ (C2)

Ograniczenia

  • (C1) 0 <= index < size(operand).
  • (C2) type(result) = tuple_element_types(operand)[index].

Przykłady

// %operand: ([1.0, 2.0], (3))
%result = "stablehlo.get_tuple_element"(%operand) {
  index = 0 : i32
} : (tuple<tensor<2xf32>, tuple<tensor<i32>>>) -> tensor<2xf32>
// %result: [1.0, 2.0]

Więcej przykładów

if

Semantyka

Produkuje dane wyjściowe podczas wykonywania dokładnie 1 funkcji z funkcji true_branch lub false_branch w zależności od wartości pred. Oficjalnie: result = pred ? true_branch() : false_branch().

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) pred Tensor 0-wymiarowy typu i1
(I2) true_branch funkcja (C1–C3)
(I3) false_branch funkcja (C1) (C2)

Wyniki

Nazwa Typ Ograniczenia
results zmienna tensorów, tensorów kwantowych lub tokenów (C3)

Ograniczenia

  • (C1) input_types(true_branch) = input_types(false_branch) = [].
  • (C2) output_types(true_branch) = output_types(false_branch).
  • (C3) type(results...) = output_types(true_branch).

Przykłady

// %result_true_branch: 10
// %result_false_branch: 11
// %pred: true
%result = "stablehlo.if"(%pred) ({
  "stablehlo.return"(%result_true_branch) : (tensor<i32>) -> ()
}, {
  "stablehlo.return"(%result_false_branch) : (tensor<i32>) -> ()
}) : (tensor<i1>) -> tensor<i32>
// %result: 10

Więcej przykładów

Obraz

Semantyka

Wyodrębnia część urojoną (pod względem elementu) z obiektu operand i generuje tensor result. Bardziej formalnie dla każdego elementu x: imag(x) = is_complex(x) ? imaginary_part(x) : constant(0, element_type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub zespolonego (C1) (C2)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego (C1) (C2)

Ograniczenia

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

Przykłady

// %operand: [(1.0, 2.0), (3.0, 4.0)]
%result = "stablehlo.imag"(%operand) : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
// %result: [2.0, 4.0]

Więcej przykładów

In-Feed

Semantyka

Odczytuje dane z reklamy In-Feed i generuje results.

Semantyka funkcji infeed_config jest zdefiniowana w implementacji.

results składa się z wartości ładunków, które pojawiają się jako pierwsze, i tokena znajdującego się na końcu. W przyszłości planujemy podzielić ładunek i token na 2 osobne dane wyjściowe, aby zwiększyć przejrzystość (#670).

Dane wejściowe

Etykieta Nazwa Typ
(I1) token token
(I2) infeed_config stała typu string

Wyniki

Nazwa Typ Ograniczenia
results zmienna tensorów, tensorów kwantowych lub tokenów (C1–C3)

Ograniczenia

  • (C1) 0 < size(results).
  • (C2) is_empty(result[:-1]) lub is_tensor(type(results[:-1])).
  • (C3) is_token(type(results[-1])).

Przykłady

// %token: !stablehlo.token
// infeed_queue[0]: [[1, 2], [3, 4]]
// infeed_queue[1]: [[5, 6], [7, 8]]
%results0:2 = "stablehlo.infeed"(%token) {
  infeed_config = ""
} : (!stablehlo.token) -> (tensor<2x2xi64>, !stablehlo.token)
// results0#0: [[1, 2], [3, 4]]
%results1:2 = "stablehlo.infeed"(%token) {
  infeed_config = ""
} : (!stablehlo.token) -> (tensor<2x2xi64>, !stablehlo.token)
// results1#0: [[5, 6], [7, 8]]

Więcej przykładów

Iota

Semantyka

Wypełnia tensor output wartościami w kolejności rosnącej od zera wzdłuż wymiaru iota_dimension. Bardziej formalnie.

output[result_index] = constant(is_quantized(output) ? quantize(result_index[iota_dimension], element_type(output)) : result_index[iota_dimension], element_type(output)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) iota_dimension si64 (C1)

Wyniki

Nazwa Typ Ograniczenia
output tensor liczby całkowitej, zmiennoprzecinkowej lub zespolonej albo tensora skwantyzowanego według intensywności (C1)

Ograniczenia

  • (C1) 0 <= iota_dimension < rank(output).

Przykłady

%output = "stablehlo.iota"() {
  iota_dimension = 0 : i64
} : () -> tensor<4x5xi32>
// %output: [
//           [0, 0, 0, 0, 0],
//           [1, 1, 1, 1, 1],
//           [2, 2, 2, 2, 2],
//           [3, 3, 3, 3, 3]
//          ]

%output = "stablehlo.iota"() {
  iota_dimension = 1 : i64
} : () -> tensor<4x5xi32>
// %output: [
//           [0, 1, 2, 3, 4],
//           [0, 1, 2, 3, 4],
//           [0, 1, 2, 3, 4],
//           [0, 1, 2, 3, 4]
//          ]

Więcej przykładów

is_finite

Semantyka

Sprawdza z uwzględnieniem elementów, czy wartość w funkcji x jest skończona (czyli nie jest +Inf, -Inf ani NaN) i generuje tensor y. Implementuje operację isFinite ze specyfikacji IEEE-754. W przypadku typów skwantyzowanych wynik to zawsze true.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) x tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C1)

Wyniki

Nazwa Typ Ograniczenia
y tensor typu logicznego (C1)

Ograniczenia

  • (C1) shape(x) = shape(y).

Przykłady

// Logical values: -Inf, +Inf, NaN, ...
// %x: [0xFFF0000000000000, 0x7FF0000000000000, 0x7FF8000000000000, -10.0, -0.0, 0.0, 10.0]
%y = "stablehlo.is_finite"(%x) : (tensor<7xf64) -> tensor<7xi1>
// %y: [false, false, false, true, true, true, true]

Więcej przykładów

log

Semantyka

Wykonuje operację logarytmu z wykorzystaniem elementów na tensorze operand i generuje tensor result. W zależności od typu elementu wykona te działania:

  • Liczba zmiennoprzecinkowa: log (IEEE-754).
  • W przypadku liczb zespolonych: logarytm zespolony.
  • W przypadku typów skwantyzowanych: dequantize_op_quantize(log, operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Ograniczenia

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

Przykłady

// %operand: [[1.0, 2.0], [3.0, 4.0]]
%result = "stablehlo.log"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
// %result: [[0.0, 0.69314718055994529], [1.0986122886681098, 1.3862943611198906]]

Więcej przykładów

log_plus_one

Semantyka

Wykonuje logarytm związany z elementami i 1 operację na tensorze operand oraz generuje tensor result. W zależności od typu elementu wykona te działania:

  • Liczba zmiennoprzecinkowa: logp1 (IEEE-754).
  • W przypadku liczb zespolonych: logarytm zespolony plus jedynka.
  • W przypadku typów skwantyzowanych: dequantize_op_quantize(log_plus_one, operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Ograniczenia

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

Przykłady

// %operand: [0.0, -0.999, 7.0, 6.38905621, 15.0]
%result = "stablehlo.log_plus_one"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// %result: [0.0, -6.90776825, 2.07944155, 2.0, 2.77258873]

Więcej przykładów

logistyka

Semantyka

Wykonuje operacje logistyczne z wykorzystaniem elementów na tensorze operand i generuje tensor result. W zależności od typu elementu wykona te działania:

  • Liczba zmiennoprzecinkowa: division(1, addition(1, exp(-x))) (IEEE-754).
  • W przypadku liczb zespolonych: logistyka zespolona.
  • W przypadku typów skwantyzowanych: dequantize_op_quantize(logistic, operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Ograniczenia

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

Przykłady

// %operand: [[0.0, 1.0], [2.0, 3.0]]
%result = "stablehlo.logistic"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
// %result: [[0.5, 0.73105858], [0.88079708, 0.95257413]]

Więcej przykładów

mapa

Semantyka

Stosuje funkcję mapy computation do parametru inputs wzdłuż dimensions i generuje tensor result.

Oficjalnie: result[result_index] = computation(inputs...[result_index]). Zwróć uwagę, że elementy dimensions nie są obecnie używane i prawdopodobnie zostaną usunięte w przyszłości (nr 487).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) inputs zmienna liczba tensorów lub tensorów kwantyzowanych na intensywność (C1–C4)
(I2) dimensions Jednowymiarowa stała tensorowa typu si64 (C3)
(I3) computation funkcja (C4)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność [C1] (C4)

Ograniczenia

  • (C1) shape(inputs...) = shape(result).
  • (C2) 0 < size(inputs) = N.
  • (C3) dimensions = range(rank(inputs[0])).
  • (C4) computation ma typ (tensor<E0>, ..., tensor<EN-1>) -> tensor<E'>, gdzie Ei = element_type(inputs[i]) i E' = element_type(result).

Przykłady

// %input0: [[0, 1], [2, 3]]
// %input1: [[4, 5], [6, 7]]
%result = "stablehlo.map"(%input0, %input1) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = stablehlo.multiply %arg0, %arg1 : tensor<i64>
    stablehlo.return %0 : tensor<i64>
}) {
  dimensions = dense<[0, 1]> : tensor<2xi64>
} : (tensor<2x2xi64>, tensor<2x2xi64>) -> tensor<2x2xi64>
// %result: [[0, 5], [12, 21]]

Więcej przykładów

maksimum

Semantyka

Wykonuje maksymalną operację dotyczącą elementów na tensorach lhs i rhs oraz generuje tensor result. W zależności od typu elementu wykona te działania:

  • W przypadku wartości logicznych: logiczne LUB.
  • W przypadku liczb całkowitych: maksymalna liczba całkowita.
  • Liczba zmiennoprzecinkowa: maximum (IEEE-754).
  • W przypadku liczb zespolonych: maksymalna liczba leksykograficzna dla pary (real, imaginary). Ustalanie kolejności liczb zespolonych może być zaskakujące, dlatego w przyszłości planujemy wycofanie w tej operacji obsługi liczb zespolonych (#560).
  • W przypadku typów skwantyzowanych:
    • dequantize_op_quantize(maximum, lhs, rhs, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor lub tensor kwantyzowany na intensywność (C1)
(I2) rhs tensor lub tensor kwantyzowany na intensywność (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C1)

Ograniczenia

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

Przykłady

// %lhs: [[1, 2], [7, 8]]
// %rhs: [[5, 6], [3, 4]]
%result = "stablehlo.maximum"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[5, 6], [7, 8]]

Więcej przykładów

minimum

Semantyka

Wykonuje operacje minimalnego na poziomie elementów na tensorach lhs i rhs oraz generuje tensor result. W zależności od typu elementu wykona te działania:

  • W przypadku wartości logicznych: logiczne ORAZ.
  • W przypadku liczb całkowitych: minimalna liczba całkowita.
  • Liczba zmiennoprzecinkowa: minimum (IEEE-754).
  • W przypadku liczb zespolonych: minimum leksykograficzne dla pary (real, imaginary). Ustalanie kolejności liczb zespolonych może być zaskakujące, dlatego w przyszłości planujemy wycofanie w tej operacji obsługi liczb zespolonych (#560).
  • W przypadku typów skwantyzowanych:
    • dequantize_op_quantize(minimum, lhs, rhs, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor lub tensor kwantyzowany na intensywność (C1)
(I2) rhs tensor lub tensor kwantyzowany na intensywność (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C1)

Ograniczenia

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

Przykłady

// %lhs: [[1, 2], [7, 8]]
// %rhs: [[5, 6], [3, 4]]
%result = "stablehlo.minimum"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[1, 2], [3, 4]]

Więcej przykładów

mnożenie

Semantyka

Wykonuje iloczyn elementów 2 tensorów lhs i rhs oraz generuje tensor result. W zależności od typu elementu wykona te działania:

  • W przypadku wartości logicznych: logiczne ORAZ.
  • W przypadku liczb całkowitych: mnożenie przez liczbę całkowitą.
  • Liczba zmiennoprzecinkowa: multiplication (IEEE-754).
  • W przypadku liczb zespolonych: mnożenie zespolone.
  • W przypadku typów skwantyzowanych:
    • dequantize_op_quantize(multiply, lhs, rhs, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor lub tensor kwantyzowany na intensywność (C1)
(I2) rhs tensor lub tensor kwantyzowany na intensywność (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C1)

Ograniczenia

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

Przykłady

// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.multiply"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[5, 12], [21, 32]]

Więcej przykładów

usuń

Semantyka

Wykonuje negację tensora operand na poziomie elementów i generuje tensor result. W zależności od typu elementu wykona te działania:

  • W przypadku liczb całkowitych: negacja liczb całkowitych.
  • W przypadku liczb całkowitych bez znaku: bitcast na liczbę całkowitą ze znakiem, negacja liczby całkowitej, bitcast z powrotem na nieoznaczoną liczbę całkowitą.
  • Liczba zmiennoprzecinkowa: negate (IEEE-754).
  • W przypadku liczb zespolonych: negacja zespolona.
  • W przypadku typów skwantyzowanych: dequantize_op_quantize(negate, operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu całkowitej, zmiennoprzecinkowej, zespolonej lub tensora skwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu całkowitej, zmiennoprzecinkowej, zespolonej lub tensora skwantyzowanego według intensywności (C1)

Ograniczenia

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

Przykłady

// Negation operation with integer Tensors
// %operand: [0, -2]
%result = "stablehlo.negate"(%operand) : (tensor<2xi32>) -> tensor<2xi32>
// %result: [0, 2]

// Negation operation with with complex tensors
// %operand: (2.5, 0.0)
%result = "stablehlo.negate"(%operand) : (tensor<1xcomplex<f32>>) -> tensor<1xcomplex<f32>>
// %result: [-2.5, -0.0]

Więcej przykładów

nie

Semantyka

Wykonuje wykonywanie na poziomie elementów NIE tensora operand i generuje tensor result. W zależności od typu elementu wykona te działania:

  • W przypadku wartości logicznych: logiczne NIE.
  • W przypadku liczb całkowitych: bitowe NIE.

Argumenty

Nazwa Typ Ograniczenia
operand tensor typu logicznego lub liczby całkowitej (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu logicznego lub liczby całkowitej (C1)

Ograniczenia

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

Przykłady

// Bitwise operation with with integer tensors
// %operand: [[1, 2], [3, 4]]
%result = "stablehlo.not"(%operand) : (tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[-2, -3], [-4, -5]]

// Bitwise operation with with boolean tensors
// %operand: [true, false]
%result = "stablehlo.not"(%operand) : (tensor<2xi1>) -> tensor<2xi1>
// %result: [false, true]

optimization_barrier

Semantyka

Zapewnia, że operacje, które generują operand, są wykonywane przed wszelkimi operacjami zależnymi od result, i zapobiega przenoszeniu operacji przez kompilator przez barierę. Poza tym operacja to tożsamość, tj. result = operand.

Argumenty

Nazwa Typ Ograniczenia
operand zmienna liczba tensorów, tensorów lub tokenów kwantyzowanych na poziomie intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result zmienna liczba tensorów, tensorów lub tokenów kwantyzowanych na poziomie intensywności (C1)

Ograniczenia

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

Przykłady

// %operand0: 0.0
// %operand1: 1.0
%result0, %result1 = "stablehlo.optimization_barrier"(%operand0, %operand1) : (tensor<f32>, tensor<f32>) -> (tensor<f32>, tensor<f32>)
// %result0: 0.0
// %result1: 1.0

Więcej przykładów

lub

Semantyka

Wykonuje parametr LUB 2 tensory lhs i rhs oraz generuje tensor result. W zależności od typu elementu wykona te działania:

  • W przypadku wartości logicznych: logiczne LUB.
  • W przypadku liczb całkowitych: bitowe LUB.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor typu całkowitej lub wartości logicznej (C1)
(I2) rhs tensor typu całkowitej lub wartości logicznej (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu całkowitej lub wartości logicznej (C1)

Ograniczenia

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

Przykłady

// Bitwise operation with with integer tensors
// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.or"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[5, 6], [7, 12]]

// Logical operation with with boolean tensors
// %lhs: [[false, false], [true, true]]
// %rhs: [[false, true], [false, true]]
%result = "stablehlo.or"(%lhs, %rhs) : (tensor<2x2xi1>, tensor<2x2xi1>) -> tensor<2x2xi1>
// %result: [[false, true], [true, true]]

Outfeed

Semantyka

Zapisuje inputs w pliku danych Out-Feed i generuje token result.

Semantyka funkcji outfeed_config jest zdefiniowana w implementacji.

Dane wejściowe

Etykieta Nazwa Typ
(I1) inputs zmienna tensorów lub tensorów skwantyzowanych
(I2) token token
(I3) outfeed_config stała typu string

Wyniki

Nazwa Typ
result token

Przykłady

%result = "stablehlo.outfeed"(%inputs0, %token) {
  outfeed_config = ""
} : (tensor<2x2x2xi64>, !stablehlo.token) -> !stablehlo.token

Więcej przykładów

pole do popisu

Semantyka

Rozwija element operand o dopełnienie wokół tensora oraz między jego elementami za pomocą podanej wartości padding_value.

edge_padding_low i edge_padding_high określają stopień dopełnienia dodanego odpowiednio na dole (obok indeksu 0) i w najwyższej (obok najwyższego indeksu) każdego wymiaru. Ilość dopełnienia może być ujemna, gdzie bezwzględna wartość dopełnienia ujemnego wskazuje liczbę elementów do usunięcia z określonego wymiaru.

interior_padding określa stopień dopełnienia dodanego między dowolnymi dwoma elementami w każdym wymiarze, który nie może być liczbą ujemną. Dopełnienie wewnętrzne ma miejsce przed dopełnieniem krawędzi, dzięki któremu dopełnienie ujemne usuwa elementy z operandu z wewnętrznym dopełnieniem.

Dokładniej rzecz ujmując, result[result_index] to:

  • operand[operand_index], jeśli result_index = edge_padding_low + operand_index * (interior_padding + 1).
  • W przeciwnym razie padding_value.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor lub tensor kwantyzowany na intensywność (C1), [C2], (C4)
(I2) padding_value tensor 0-wymiarowy lub tensor skwantyzowany na intensywność (C1)
(I3) edge_padding_low Jednowymiarowa stała tensorowa typu si64 [C1] (C4)
(I4) edge_padding_high Jednowymiarowa stała tensorowa typu si64 [C1] (C4)
(I5) interior_padding Jednowymiarowa stała tensorowa typu si64 (C2–C4)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C3–C6)

Ograniczenia

  • (C1) element_type(operand) = element_type(padding_value) = element_type(result).
  • (C2) size(edge_padding_low) = size(edge_padding_high) = size(interior_padding) = rank(operand).
  • (C3) 0 <= interior_padding.
  • (C4) shape(result) = shape(operand) + edge_padding_low + max(shape(operand) - 1, 0) * interior_padding + edge_padding_high.

Przykłady

// %operand: [
//            [1, 2, 3],
//            [4, 5, 6]
//           ]
// %padding_value: 0
%result = "stablehlo.pad"(%operand, %padding_value) {
  edge_padding_low = dense<[0, 1]> : tensor<2xi64>,
  edge_padding_high = dense<[2, 1]> : tensor<2xi64>,
  interior_padding = dense<[1, 2]> : tensor<2xi64>
} : (tensor<2x3xi32>, tensor<i32>) -> tensor<5x9xi32>
// %result: [
//           [0, 1, 0, 0, 2, 0, 0, 3, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0],
//           [0, 4, 0, 0, 5, 0, 0, 6, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0]
//          ]

Więcej przykładów

partition_id

Semantyka

Generuje partition_id bieżącego procesu.

Wyniki

Nazwa Typ
result Tensor 0-wymiarowy typu ui32

Przykłady

%result = "stablehlo.partition_id"() : () -> tensor<ui32>

Więcej przykładów

Popcnt

Semantyka

Oblicza liczbę bitów ustawioną we tensorze operand i generuje tensor result.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu liczby całkowitej (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu liczby całkowitej (C1)

Ograniczenia

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

Przykłady

// %operand: [0, 1, 2, 127]
%result = "stablehlo.popcnt"(%operand) : (tensor<4xi64>) -> tensor<4xi64>
// %result: [0, 1, 1, 7]

Więcej przykładów

moc

Semantyka

Wykonuje wykładnicze wykładnicze tensora lhs za pomocą tensora rhs i generuje tensor result. W zależności od typu elementu wykona te działania:

  • W przypadku liczb całkowitych: wykładnik liczb całkowity.
  • Liczba zmiennoprzecinkowa: pow (IEEE-754).
  • W przypadku liczb zespolonych: wykładnik zespolony.
  • W przypadku typów skwantyzowanych: dequantize_op_quantize(power, lhs, rhs, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor typu całkowitej, zmiennoprzecinkowej, zespolonej lub tensora skwantyzowanego według intensywności (C1)
(I2) rhs tensor typu całkowitej, zmiennoprzecinkowej, zespolonej lub tensora skwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu całkowitej, zmiennoprzecinkowej, zespolonej lub tensora skwantyzowanego według intensywności (C1)

Ograniczenia

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

Przykłady

// %lhs: [-2.0, -0.0, -36.0, 5.0, 3.0, 10000.0]
// %rhs: [2.0, 2.0, 1.1, 2.0, -1.0, 10.0]
%result = "stablehlo.power"(%lhs, %rhs) : (tensor<6xf64>, tensor<6xf64>) -> tensor<6xf64>
// %result: [4.0, 0.0, -nan, 25.0, 0.333333343, inf]

Więcej przykładów

prawdziwy

Semantyka

Wyodrębnia część rzeczywistą z elementu operand i generuje ten parametr result. Bardziej formalnie dla każdego elementu x: real(x) = is_complex(x) ? real_part(x) : x.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub zespolonego (C1) (C2)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego (C1) (C2)

Ograniczenia

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

Przykłady

// %operand: [(1.0, 2.0), (3.0, 4.0)]
%result = "stablehlo.real"(%operand) : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
// %result: [1.0, 3.0]

Więcej przykładów

recv

Semantyka

Odbiera dane z kanału z channel_id i generuje results.

Jeśli is_host_transfer ma wartość true, operacja przenosi dane z hosta. W przeciwnym razie dane zostaną przeniesione z innego urządzenia. To znaczy, że jest ono definiowane przez wdrożenie. Ta flaga powiela informacje podane w polu channel_type, więc w przyszłości planujemy zachować tylko jedną z nich (#666).

results składa się z wartości ładunków, które pojawiają się jako pierwsze, i tokena znajdującego się na końcu. W przyszłości planujemy podzielić ładunek i token na 2 osobne dane wyjściowe, aby zwiększyć przejrzystość (#670).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) token token (C4)
(I2) channel_id stała typu si64
(I3) channel_type wyliczenie DEVICE_TO_DEVICE i HOST_TO_DEVICE (C1)
(I4) is_host_transfer stała typu i1 (C1)

Wyniki

Nazwa Typ Ograniczenia
results zmienna tensorów, tensorów kwantowych lub tokenów (C2–C4)

Ograniczenia

  • (C1) channel_type jest zdefiniowany jako:
    • HOST_TO_DEVICE, jeśli is_host_transfer = true,
    • W przeciwnym razie DEVICE_TO_DEVICE.
  • (C2) 0 < size(results).
  • (C3) is_empty(result[:-1]) lub is_tensor(type(results[:-1])).
  • (C4) is_token(type(results[-1])).

Przykłady

%results0, %results1 = "stablehlo.recv"(%token) {
  channel_handle = #stablehlo.channel_handle<handle = 1, type = 3>,
  is_host_transfer = true
} : (!stablehlo.token) -> (tensor<2x2xi64>, !stablehlo.token)

Więcej przykładów

zmniejsz

Semantyka

Stosuje funkcję redukcji body do elementów inputs i init_values wzdłuż linii dimensions i generuje results tensorów.

Kolejność redukcji jest określona przez implementację, co oznacza, że body i init_values muszą tworzyć monoid, aby uzyskać taki sam efekt w przypadku wszystkich danych wejściowych we wszystkich implementacjach. Taka sytuacja nie wystarczy jednak w przypadku wielu popularnych redukcji. Na przykład dodanie liczby zmiennoprzecinkowej dla body i 0 w przypadku elementu init_values nie tworzy monoidu, ponieważ dodawanie liczby zmiennoprzecinkowej nie jest skojarzone.

Bardziej formalnie, results...[j0, ..., jR-1] = reduce(input_slices_converted), gdzie:

  • input_slices = inputs...[j0, ..., :, ..., jR-1], gdzie : wstawiany jest dimensions.
  • input_slices_converted = to_destination_type(input_slices..., type(func_inputs(body)[:len(func_inputs(body))//2])...).
  • init_values_converted = to_destination_type(init_values..., type(func_inputs(body)[len(func_inputs(body))//2:])...).
  • reduce(input_slices_converted) = exec(schedule) w przypadku drzewa binarnego, schedule, gdzie:
    • exec(node) = body(exec(node.left), exec(node.right)).
    • exec(leaf) = leaf.value.
  • schedule to zdefiniowane w implementacji pełne drzewo binarne, którego przemierzanie w kolejności obejmuje:
    • input_slices_converted...[index] dla wszystkich index w index_space(input_slices_converted) w rosnącej kolejności leksykograficznej index.
    • Przeplatany z zadeklarowaną w implementacji wartością init_values_converted w miejscach określonych przez implementację.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) inputs zmienna liczba tensorów lub tensorów kwantyzowanych na intensywność (C1–C4), (C6), (C7)
(I2) init_values zmienna liczba tensorów 0-wymiarowych lub tensorów skwantyzowanych na intensywność (C2) (C3)
(I3) dimensions Jednowymiarowa stała tensorowa typu si64 [C4], [C5], [C7]
(I4) body funkcja (C6)

Wyniki

Nazwa Typ Ograniczenia
results zmienna liczba tensorów lub tensorów kwantyzowanych na intensywność (C3), [C7], (C8)

Ograniczenia

  • (C1) same(shape(inputs...)).
  • (C2) element_type(inputs...) = element_type(init_values...).
  • (C3) 0 < size(inputs) = size(init_values) = size(results) = N.
  • (C4) 0 <= dimensions < rank(inputs[0]).
  • (C5) is_unique(dimensions).
  • (C6) body ma typ (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>), gdzie is_promotable(element_type(inputs[i]), Ei).
  • (C7) shape(results...) = shape(inputs...) z tym wyjątkiem, że rozmiary wymiarów inputs... odpowiadające dimensions nie są uwzględniane.
  • (C8) element_type(results[i]) = Ei dla wszystkich i w [0,N).

Przykłady

// %input = [[0, 1, 2, 3, 4, 5]]
// %init_value = 0
%result = "stablehlo.reduce"(%input, %init_value) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  dimensions = dense<1> : tensor<1xi64>
} : (tensor<1x6xi64>, tensor<i64>) -> tensor<1xi64>
// %result = [15]

Więcej przykładów

reduce_precision

Semantyka

Przeprowadza z uwzględnieniem elementów konwersję typu operand na inny typ zmiennoprzecinkowy, który korzysta z exponent_bits i mantissa_bits, oraz z powrotem do pierwotnego typu zmiennoprzecinkowego, generując tensor output.

Bardziej formalnie:

  • Bity mantissy pierwotnej wartości są zaokrąglane w taki sposób, aby zaokrąglały ją do najbliższej wartości, którą można reprezentować za pomocą mantissa_bits z wykorzystaniem semantyki roundToIntegralTiesToEven.
  • Następnie, jeśli wartość mantissa_bits jest mniejsza niż liczba bitów mantissy w pierwotnej wartości, bity mantissy są obcinane do wartości mantissa_bits.
  • Następnie, jeśli bity wykładnicze wyniku pośredniego nie mieszczą się w zakresie podanym przez funkcję exponent_bits, wynik pośredni wypływa z nieskończoności do nieskończoności, używając znaku oryginalnego, lub rozszerza się do zera przy użyciu znaku pierwotnego.
  • W przypadku typów skwantyzowanych wykonuje dequantize_op_quantize( lambda operand: reduce_precision(operand, exponent_bits, mantissa_bits), operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C1)
(I2) exponent_bits stała typu si32 (C2)
(I3) mantissa_bits stała typu si32 (C3)

Wyniki

Nazwa Typ Ograniczenia
output tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C1)

Ograniczenia

  • (C1) baseline_type(operand) = baseline_type(output).
  • (C2) 1 <= exponent_bits.
  • (C3) 0 <= mantissa_bits.

Przykłady

// Logical values: +Inf, NaN, +Denormal, 0.0, 65519.0, 65520.0
// %operand: [0x7FF0000000000000, 0x7FFFFFFFFFFFFFFF, 0x0000000000000001, 0.0, 65519.0, 65520.0]
%output = "stablehlo.reduce_precision"(%operand) {
  exponent_bits = 5 : i32,
  mantissa_bits = 10 : i32
} : (tensor<6xf64>) -> tensor<6xf64>
// Logical values: +Inf, NaN, 0.0, 0.0, 65504.0, +Inf
// %output: [0x7FF0000000000000, 0x7FFFFFFFFFFFFFFF, 0.0, 0.0, 65504.0, 0x7FF0000000000000]

Więcej przykładów

reduce_scatter

Semantyka

W ramach każdej grupy procesów w siatce procesów StableHLO przeprowadza zmniejszenie za pomocą funkcji computations wartości tensora operand z każdego procesu, dzieli wynik redukcji na części scatter_dimension na części i rozprasza te fragmenty między procesy w celu uzyskania 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ś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 ciągu każdego process_group:

  • reduced_value = all_reduce(operand, replica_groups, channel_id, use_global_device_ids, computation).
  • parts@sender = split(reduced_value@sender, dim(process_groups, 1), scatter_dimension).
  • result@receiver = parts@sender[receiver_index] dla wszystkich sender w process_group, gdzie receiver_index = process_group.index(receiver).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor lub tensor kwantyzowany na intensywność [C1], [C2], [C7], [C8]
(I2) scatter_dimension stała typu si64 (C1), [C2], (C8)
(I3) replica_groups Dwuwymiarowa stała tensorowa typu si64 (C3-C5)
(I4) channel_id stała typu si64 (C6)
(I5) use_global_device_ids stała typu i1 (C6)
(I6) computation funkcja (C7)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C8-C9)

Ograniczenia

  • (C1) dim(operand, scatter_dimension) % dim(process_groups, 1) = 0.
  • (C2) 0 <= scatter_dimension < rank(operand).
  • (C3) is_unique(replica_groups).
  • (C4) size(replica_groups) jest zdefiniowany jako:
    • num_replicas, jeśli używana jest wartość cross_replica.
    • num_replicas, jeśli używana jest wartość cross_replica_and_partition.
    • num_processes, jeśli używana jest wartość flattened_ids.
  • (C5) 0 <= replica_groups < size(replica_groups).
  • (C6) Jeśli use_global_device_ids = true – channel_id > 0.
  • (C7) computation ma typ (tensor<E>, tensor<E>) -> (tensor<E>), gdzie is_promotable(element_type(operand), E).
  • (C8) shape(result) = shape(operand) z wyjątkiem:
    • dim(result, scatter_dimension) = dim(operand, scatter_dimension) / dim(process_groups, 1).
  • (C9) element_type(result) = E.

Przykłady

// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [[1, 2, 3, 4],
//                   [5, 6, 7, 8]]
// %operand@(1, 0): [[9, 10, 11, 12],
//                   [13, 14, 15, 16]]
%result = "stablehlo.reduce_scatter"(%operand) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
  "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  scatter_dimension = 1 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor<2x4xi64>) -> tensor<2x2xi64>
//
// %result@(0, 0): [[10, 12],
//                  [18, 20]]
// %result@(1, 0): [[14, 16],
//                  [22, 24]]

Więcej przykładów

reduce_window

Semantyka

Stosuje funkcję redukcji body do okien inputs i init_values oraz daje results.

Na diagramie poniżej widać, jak za pomocą konkretnego przykładu obliczamy elementy w funkcji results..., korzystając z tabeli inputs....

Bardziej oficjalnie: results...[result_index] = reduce(windows, init_values, axes(inputs...), body) (patrz reduce), gdzie:

  • padded_inputs = pad(inputs..., init_values..., padding[:, 0], padding[:, 1], base_dilations - 1).
  • window_start = result_index * window_strides.
  • window_end = window_start + (window_dimensions - 1) * window_dilations + 1.
  • windows = slice(padded_inputs..., window_start, window_end, window_dilations).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) inputs zmienna liczba tensorów lub tensorów kwantyzowanych na intensywność (C1–C4), (C6), (C8), (C10), (C12), (C13), (C15)
(I2) init_values zmienna liczba tensorów 0-wymiarowych lub tensorów skwantyzowanych na intensywność [C1] (C13)
(I3) window_dimensions Jednowymiarowa stała tensorowa typu si64 [C4], [C5], [C15]
(I4) window_strides Jednowymiarowa stała tensorowa typu si64 (C6), [C7], (C15)
(I5) base_dilations Jednowymiarowa stała tensorowa typu si64 [C8], [C9], [C15]
(I6) window_dilations Jednowymiarowa stała tensorowa typu si64 (C10), (C11) i (C15)
(I7) padding Dwuwymiarowa stała tensorowa typu si64 (C12) (C15)
(I8) body funkcja (C13)

Wyniki

Nazwa Typ Ograniczenia
results zmienna liczba tensorów lub tensorów kwantyzowanych na intensywność [C1] (C14-C16)

Ograniczenia

  • (C1) 0 < size(inputs) = size(init_values) = size(results) = N.
  • (C2) same(shape(inputs...)).
  • (C3) element_type(inputs...) = element_type(init_values...).
  • (C4) size(window_dimensions) = rank(inputs[0]).
  • (C5) 0 < window_dimensions.
  • (C6) size(window_strides) = rank(inputs[0]).
  • (C7) 0 < window_strides.
  • (C8) size(base_dilations) = rank(inputs[0]).
  • (C9) 0 < base_dilations.
  • (C10) size(window_dilations) = rank(inputs[0]).
  • (C11) 0 < window_dilations.
  • (C12) shape(padding) = [rank(inputs[0]), 2].
  • (C13) body ma typ (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>), gdzie is_promotable(element_type(inputs[i]), Ei).
  • (C14) same(shape(results...)).
  • (C15) shape(results[0]) = num_windows, gdzie:
    • dilated_input_shape = shape(inputs[0]) = 0 ? 0 : (shape(inputs[0]) - 1) * base_dilations + 1.
    • padded_input_shape = padding[:, 0] + dilated_input_shape + padding[:, 1].
    • dilated_window_shape = (window_dimensions - 1) * window_dilations + 1.
    • is_empty_window = padded_input_shape = 0 || dilated_window_shape > padded_input_shape.
    • num_windows = is_empty_window ? 0 : floor((padded_input_shape - dilated_window_shape) / window_strides) + 1.
  • (C16) element_type(results[i]) = Ei dla wszystkich i w [0,N).

Przykłady

// %input = [[1, 2], [3, 4], [5, 6]]
// %init_value = 0
%result = "stablehlo.reduce_window"(%input, %init_value) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  window_dimensions = dense<[2, 1]> : tensor<2xi64>,
  window_strides = dense<[4, 1]> : tensor<2xi64>,
  base_dilations = dense<[2, 1]> : tensor<2xi64>,
  window_dilations = dense<[3, 1]> : tensor<2xi64>,
  padding = dense<[[2, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<3x2xi64>, tensor<i64>) -> tensor<2x2xi64>
// %result = [[0, 0], [3, 4]]

Więcej przykładów

reszta

Semantyka

Służy do wykonywania reszty dzielnika lhs i dzielnika rhs na podstawie elementu oraz generuje tensor result.

Bardziej formalnie znak wyniku jest brany z dzielnicy, a wartość bezwzględna wyniku jest zawsze mniejsza niż wartość bezwzględna dzielnika. Reszta jest obliczana jako lhs - d * rhs, przy czym d jest obliczany według wzoru:

  • W przypadku liczb całkowitych: stablehlo.divide(lhs, rhs).
  • W przypadku liczb zmiennoprzecinkowych division(lhs, rhs) z IEEE-754 z atrybutem zaokrąglania roundTowardZero.
  • W przypadku liczb zespolonych: TBD (#997).
  • W przypadku typów skwantyzowanych:
    • dequantize_op_quantize(remainder, lhs, rhs, type(result)).

W przypadku elementów zmiennoprzecinkowych ta operacja różni się od operacji remainder ze specyfikacji IEEE-754, gdzie d to wartość całkowita zbliżona do dokładnej wartości lhs/rhs z uwzględnieniem wartości parzystych.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor liczby całkowitej, zmiennoprzecinkowej lub zespolonej albo tensora skwantyzowanego według intensywności (C1)
(I2) rhs tensor liczby całkowitej, zmiennoprzecinkowej lub zespolonej albo tensora skwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor liczby całkowitej, zmiennoprzecinkowej lub zespolonej albo tensora skwantyzowanego według intensywności (C1)

Ograniczenia

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

Przykłady

// %lhs: [17, -17, 17, -17]
// %rhs: [3, 3, -3, -3]
%result = "stablehlo.remainder"(%lhs, %rhs) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi64>
// %result: [2, -2, 2, -2]

Więcej przykładów

replica_id

Semantyka

Generuje replica_id bieżącego procesu.

Wyniki

Nazwa Typ
result Tensor 0-wymiarowy typu ui32

Przykłady

%result = "stablehlo.replica_id"() : () -> tensor<ui32>

Więcej przykładów

zmieniać kształt

Semantyka

Zmienia kształt tensora operand na tensor result. Zasadniczo chodzi o zachowanie tej samej reprezentacji kanonicznej, ale potencjalnie zmianę kształtu, np. z tensor<2x3xf32> na tensor<3x2xf32> lub tensor<6xf32>.

Bardziej oficjalnie, result[result_index] = operand[operand_index], gdzie result_index i operand_index mają tę samą pozycję w leksykograficznym porządku elementów index_space(result) i index_space(operand).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor lub tensor skwantyzowany (C1–C3)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor skwantyzowany (C1–C3)

Ograniczenia

  • (C1) Wartość element_type(result) jest obliczana przez:
    • element_type(operand), jeśli !is_per_axis_quantized(operand).
    • element_type(operand) z tym wyjątkiem, że quantization_dimension(operand) i quantization_dimension(result) mogą się różnić. W przeciwnym razie te wartości mogą się różnić.
  • (C2) size(operand) = size(result).
  • (C3) Jeśli is_per_axis_quantized(operand):
    • reduce(dims(operand, [0, 1, ..., quantization_dimension(operand) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) = reduce(dims(result, [0, 1, ..., quantization_dimension(result) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y).
    • dim(operand, quantization_dimension(operand)) = dim(result, quantization_dimension(result)).
    • reduce(dims(operand, [quantization_dimension(operand) + 1, ..., rank(operand) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) = reduce(dims(result, [quantization_dimension(result) + 1, ..., rank(result) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y).

Przykłady

// %operand: [[1, 2, 3], [4, 5, 6]]]
%result = "stablehlo.reshape"(%operand) : (tensor<2x3xi32>) -> tensor<3x2xi32>
// %result: [[1, 2], [3, 4], [5, 6]]

Więcej przykładów

odwróć

Semantyka

Odwraca kolejność elementów w funkcji operand wzdłuż podanej wartości dimensions i generuje tensor result. Bardziej oficjalnie, result[result_index] = operand[operand_index], gdzie:

  • operand_index[d] = dim(result, d) - result_index[d] - 1 jeśli d w dimensions.
  • W przeciwnym razie operand_index[d] = result_index[d].

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor lub tensor kwantyzowany na intensywność (C1) (C3)
(I2) dimensions Jednowymiarowa stała tensorowa typu si64 (C2) (C3)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C1) (C3)

Ograniczenia

  • (C1) type(operand) = type(result).
  • (C2) is_unique(dimensions).
  • (C3) 0 <= dimensions < rank(result).

Przykłady

// %operand = [[1, 2], [3, 4], [5, 6]]
%result = "stablehlo.reverse"(%operand) {
  dimensions = dense<1> : tensor<1xi64>
} : (tensor<3x2xi32>) -> tensor<3x2xi32>
// %result: [[2, 1], [4, 3], [6, 5]]

Więcej przykładów

Rng

Semantyka

Generuje liczby losowe za pomocą algorytmu rng_distribution i zwłaszcza tensor result danego kształtu shape.

Jeśli jest ustawiona wartość rng_distribution = UNIFORM, liczby losowe są generowane zgodnie z równomiernym rozkładem w przedziale [a, b). Jeśli ustawiona jest wartość a >= b, działanie jest nieokreślone.

Jeśli ustawiona jest wartość rng_distribution = NORMAL, liczby losowe są generowane według rozkładu normalnego, dla którego średnia = a, a odchylenie standardowe = b. Jeśli ustawiona jest wartość b < 0, zachowanie jest nieokreślone.

Dokładny sposób generowania liczb losowych zależy od implementacji. Na przykład mogą, ale nie muszą, i mogą korzystać z stanu ukrytego, ale nie muszą.

W rozmowach z wieloma zainteresowanymi osobami ta operacja została tak samo skutecznie wycofana, więc w przyszłości planujemy ją usunąć (#597).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) a 0-wymiarowy tensor typu całkowitej, logicznej lub zmiennoprzecinkowej (C1) (C2)
(I2) b 0-wymiarowy tensor typu całkowitej, logicznej lub zmiennoprzecinkowej (C1) (C2)
(I3) shape Jednowymiarowa stała tensorowa typu si64 (C3)
(I4) rng_distribution wyliczenie UNIFORM i NORMAL (C2)

Wyniki

Nazwa Typ Ograniczenia
result tensor liczb całkowitych, logicznych lub zmiennoprzecinkowych (C1–C3)

Ograniczenia

  • (C1) element_type(a) = element_type(b) = element_type(result).
  • (C2) Jeśli rng_distribution = NORMAL, to is_float(a).
  • (C3) shape(result) = shape.

Przykłady

// %a = 0
// %b = 2
// %shape = [3, 3]
%result = "stablehlo.rng"(%a, %b, %shape) {
  rng_distribution = #stablehlo<rng_distribution UNIFORM>
} : (tensor<i32>, tensor<i32>, tensor<2xi64>) -> tensor<3x3xi32>
// %result: [
//           [1, 0, 1],
//           [1, 1, 1],
//           [0, 0, 0]
//          ]

rng_bit_generator

Semantyka

Zwraca element output wypełniony jednolitymi losowymi bitami i zaktualizowanym stanem wyjściowym output_state za pomocą algorytmu generatora liczb pseudolosowych rng_algorithm przy określonym stanie początkowym initial_state. Gwarantujemy, że dane wyjściowe będą deterministyczną funkcją initial_state, ale nie możemy zagwarantować, że będą deterministyczne między implementacjami.

rng_algorithm należy do jednej z tych wartości:

  • DEFAULT: algorytm zdefiniowany w implementacji.
  • THREE_FRY: zdefiniowany w implementacji wariant algorytmu Threefry.*
  • PHILOX: zdefiniowany w ramach implementacji wariant algorytmu Philox*.

* Patrz: Salmon i in. SC 2011. Równoległe liczby losowe: tak proste jak 1, 2, 3.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) rng_algorithm wyliczenie DEFAULT, THREE_FRY i PHILOX (C2)
(I2) initial_state Jednowymiarowy tensor typu ui64 (C1) (C2)

Wyniki

Nazwa Typ Ograniczenia
output_state Jednowymiarowy tensor typu ui64 (C1)
output tensor liczby całkowitej lub liczby zmiennoprzecinkowej

Ograniczenia

  • (C1) type(initial_state) = type(output_state).
  • (C2) size(initial_state) jest zdefiniowany jako:
    • zdefiniowana, jeśli rng_algorithm = DEFAULT.
    • 2, jeśli rng_algorithm = THREE_FRY.
    • 2 lub 3, jeśli rng_algorithm = PHILOX.

Przykłady

// %initial_state: [1, 2]
%output_state, %output = "stablehlo.rng_bit_generator"(%initial_state) {
  rng_algorithm = #stablehlo<rng_algorithm THREE_FRY>
} : (tensor<2xui64>) -> (tensor<2xui64>, tensor<2x2xui64>)
// %output_state: [1, 6]
// %output: [
//           [9236835810183407956, 16087790271692313299],
//           [18212823393184779219, 2658481902456610144]
//          ]

round_nearest_afz

Semantyka

Wykonuje zaokrąglanie według elementów w kierunku najbliższej liczby całkowitej, przerywając reprezentacje od zera na tensorze operand i generuje tensor result. Implementuje operację roundToIntegralTiesToAway ze specyfikacji IEEE-754. W przypadku typów skwantyzowanych wykonuje dequantize_op_quantize(round_nearest_afz, operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C1)

Ograniczenia

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

Przykłady

// %operand = [-2.5, 0.4, 0.5, 0.6, 2.5]
%result = "stablehlo.round_nearest_afz"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// %result: [-3.0, 0.0, 1.0, 1.0, 3.0]

Więcej przykładów

round_nearest_even

Semantyka

Wykonuje zaokrąglanie według elementów w kierunku najbliższej liczby całkowitej, przerywając wiązania w kierunku parzystej liczby całkowitej na tensorze operand i generuje tensor result. Implementuje operację roundToIntegralTiesToEven ze specyfikacji IEEE-754. W przypadku typów skwantyzowanych wykonuje dequantize_op_quantize(round_nearest_even, operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub tensor skwantyzowany na intensywność (C1)

Ograniczenia

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

Przykłady

// %operand = [-2.5, 0.4, 0.5, 0.6, 2.5]
%result = "stablehlo.round_nearest_even"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// %result: [-2.0, 0.0, 0.0, 1.0, 2.0]

Więcej przykładów

rsqrt

Semantyka

Wykonuje operację odwrotności pierwiastka kwadratowego na tensorze operand i generuje tensor result. W zależności od typu elementu wykona te działania:

  • Liczba zmiennoprzecinkowa: rSqrt (IEEE-754).
  • W przypadku liczb zespolonych: zespolony odwrotny pierwiastek kwadratowy.
  • W przypadku typów skwantyzowanych: dequantize_op_quantize(rsqrt, operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Ograniczenia

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

Przykłady

// %operand: [[1.0, 4.0], [9.0, 25.0]]
%result = "stablehlo.rsqrt"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[1.0, 0.5], [0.33333343, 0.2]]

Więcej przykładów

scatter

Semantyka

Generuje tensor results, który jest równy tensorom inputs, z tą różnicą, że kilka wycinków określonych przez scatter_indices zostało zaktualizowanych przy użyciu wartości updates za pomocą parametru update_computation.

Na diagramie poniżej widać, jak za pomocą konkretnego przykładu mapujemy elementy w tabeli updates... na elementy w komponencie results.... Diagram wybiera kilka przykładowych indeksów updates... i dokładnie wyjaśnia, którym indeksom results... odpowiadają.

Oficjalnie dla wszystkich update_index w index_space(updates[0]):

  • update_scatter_dims = [d for d in axes(updates[0]) and d not in update_window_dims].
  • update_scatter_index = update_index[update_scatter_dims...].
  • start_index został zdefiniowany jako:
    • scatter_indices[si0, ..., :, ..., siN], gdzie si to poszczególne elementy w update_scatter_index, a : jest wstawiony w indeksie index_vector_dim, jeśli index_vector_dim < rank(scatter_indices).
    • W przeciwnym razie [scatter_indices[update_scatter_index]].
  • Przez d_input w: axes(inputs[0]),
    • full_start_index[d_input] = start_index[d_start], jeśli d_input = scatter_dims_to_operand_dims[d_start].
    • W przeciwnym razie full_start_index[d_input] = 0.
  • update_window_index = update_index[update_window_dims...].
  • full_window_index = [wi0, ..., 0, ..., wiN], gdzie wi to poszczególne elementy w update_window_index, a 0 jest wstawiony w indeksach od inserted_window_dims.
  • result_index = full_start_index + full_window_index.

W związku z tym results = exec(schedule, inputs), gdzie:

  • schedule to zdefiniowana w implementacji permutacja funkcji index_space(updates[0]).
  • exec([update_index, ...], results) = exec([...], updated_results), gdzie:
    • Jeśli result_index mieści się w zakresie shape(results...)
    • updates_converted = to_destination_type( updates...[update_index], type(func_inputs(update_computation) [len(func_inputs(update_computation))//2:])... )
    • updated_values = update_computation(results...[result_index], updates_converted)
    • updated_results to kopia klucza results z wartością results...[result_index] ustawioną na updated_values....
    • W innym przypadku
    • updated_results = results.
  • exec([], results) = results.

Jeśli indices_are_sorted ma wartość true, implementacja może zakładać, że scatter_indices jest posortowany według właściwości scatter_dims_to_operand_dims. W przeciwnym razie zachowanie jest niezdefiniowane. Oficjalnie dla wszystkich i1 < i2 z indices(result), full_start_index(i1) <= full_start_index(i2).

Jeśli unique_indices ma wartość true, implementacja może przyjąć, że wszystkie rozproszone indeksy result_index są unikalne. Jeśli unique_indices ma wartość true, ale rozkładane indeksy nie są unikalne, zachowanie jest nieokreślone.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) inputs zmienna liczba tensorów lub tensorów kwantyzowanych na intensywność (C1), (C2), (C4-C6), (C10), (C13), (C15-C16)
(I2) scatter_indices tensor typu liczby całkowitej (C4), [C11], (C14)
(I3) updates zmienna liczba tensorów lub tensorów kwantyzowanych na intensywność (C3-C6), (C8)
(I4) update_window_dims Jednowymiarowa stała tensorowa typu si64 [C2], [C4], [C7], [C8]
(I5) inserted_window_dims Jednowymiarowa stała tensorowa typu si64 [C2], [C4], [C9], [C10]
(I6) scatter_dims_to_operand_dims Jednowymiarowa stała tensorowa typu si64 (C11-C13)
(I7) index_vector_dim stała typu si64 (C4), [C11], (C14)
(I8) indices_are_sorted stała typu i1
(I9) unique_indices stała typu i1
(I10) update_computation funkcja (C15)

Wyniki

Nazwa Typ Ograniczenia
results zmienna liczba tensorów lub tensorów kwantyzowanych na intensywność (C15-C17)

Ograniczenia

  • (C1) same(shape(inputs...)).
  • (C2) rank(inputs[0]) = size(update_window_dims) + size(inserted_window_dims).
  • (C3) same(shape(updates...)).
  • (C4) shape(updates[0]) = combine(update_scatter_dim_sizes, update_window_dim_sizes), gdzie:
    • update_scatter_dim_sizes = shape(scatter_indices) z tym wyjątkiem, że rozmiar wymiaru scatter_indices odpowiadający wartości index_vector_dim nie jest uwzględniony.
    • update_window_dim_sizes <= shape(inputs[0]) z tym wyjątkiem, że rozmiary podane w polu inputs[0] odpowiadające wartościom inserted_window_dims nie są uwzględniane.
    • combine umieszcza update_scatter_dim_sizes na osiach odpowiadających wartościom update_scatter_dims i update_window_dim_sizes na osiach odpowiadających wartościom update_window_dims.
  • (C5) 0 < size(inputs) = size(updates) = N.
  • (C6) element_type(updates...) = element_type(inputs...).
  • (C7) is_unique(update_window_dims) and is_sorted(update_window_dims).
  • (C8) 0 <= update_window_dims < rank(updates[0]).
  • (C9) is_unique(inserted_window_dims) and is_sorted(update_window_dims).
  • (C10) 0 <= inserted_window_dims < rank(inputs[0]).
  • (C11) size(scatter_dims_to_operand_dims) = index_vector_dim < rank(scatter_indices) ? dim(scatter_indices, index_vector_dim) : 1.
  • (C12) is_unique(scatter_dims_to_operand_dims).
  • (C13) 0 <= scatter_dims_to_operand_dims < rank(inputs[0]).
  • (C14) 0 <= index_vector_dim <= rank(scatter_indices).
  • (C15) update_computation ma typ (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>), gdzie is_promotable(element_type(inputs[i]), Ei).
  • (C16) shape(inputs...) = shape(results...).
  • (C17) element_type(results[i]) = Ei dla wszystkich i w [0,N).

Przykłady

// %input: [
//          [[1, 2], [3, 4], [5, 6], [7, 8]],
//          [[9, 10], [11, 12], [13, 14], [15, 16]],
//          [[17, 18], [19, 20], [21, 22], [23, 24]]
//         ]
// %scatter_indices: [[[0, 2], [1, 0], [2, 1]], [[0, 1], [1, 0], [0, 9]]]
// %update: [
//           [[[1, 1], [1, 1]], [[1, 1], [1, 1]], [[1, 1], [1, 1]]],
//           [[[1, 1], [1, 1]], [[1, 1], [1, 1]], [[1, 1], [1, 1]]]
//          ]
%result = "stablehlo.scatter"(%input, %scatter_indices, %update) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  scatter_dimension_numbers = #stablehlo.scatter<
    update_window_dims = [2, 3],
    inserted_window_dims = [0],
    scatter_dims_to_operand_dims = [1, 0],
    index_vector_dim = 2>,
  indices_are_sorted = false,
  unique_indices = false
} : (tensor<3x4x2xi64>, tensor<2x3x2xi64>, tensor<2x3x2x2xi64>) -> tensor<3x4x2xi64>
// %result: [
//           [[1, 2], [5, 6], [7, 8], [7, 8]],
//           [[10, 11], [12, 13], [14, 15], [16, 17]],
//           [[18, 19], [20, 21], [21, 22], [23, 24]]
//          ]

Więcej przykładów

wybierz

Semantyka

Powoduje wygenerowanie tensora result, w którym każdy element jest wybierany z tensora on_true lub on_false na podstawie wartości odpowiadającego mu elementu pred. Bardziej formalnie: result[result_index] = pred_element ? on_true[result_index] : on_false[result_index], gdzie pred_element = rank(pred) = 0 ? pred[] : pred[result_index]. W przypadku typów skwantyzowanych wykonuje dequantize_select_quantize(pred, on_true, on_false, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) pred tensor typu i1 (C1)
(I2) on_true tensor lub tensor kwantyzowany na intensywność (C1-C2)
(I3) on_false tensor lub tensor kwantyzowany na intensywność (C2)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C2)

Ograniczenia

  • (C1) rank(pred) = 0 or shape(pred) = shape(on_true).
  • (C2) baseline_type(on_true) = baseline_type(on_false) = baseline_type(result).

Przykłady

// %pred: [[false, true], [true, false]]
// %on_true: [[1, 2], [3, 4]]
// %on_false: [[5, 6], [7, 8]]
%result = "stablehlo.select"(%pred, %on_true, %on_false) : (tensor<2x2xi1>, tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[5, 2], [3, 8]]

Więcej przykładów

select_and_scatter

Semantyka

Rozkłada wartości z tensora source za pomocą scatter na podstawie wyniku reduce_window tensora input przy użyciu select i generuje tensor result.

Na diagramie poniżej widać, jak za pomocą konkretnego przykładu obliczamy elementy w atrybutach result z wartości operand i source.

Bardziej formalnie:

  • selected_values = reduce_window_without_init(...) tymi danymi:

    • `inputs = [operand].
    • window_dimensions, window_strides i padding, które są używane w niezmienionej postaci.
    • base_dilations = windows_dilations = 1.
    • body został zdefiniowany jako:
    def body(arg0: tensor<E>, arg1: tensor<E>) -> tensor<E>:
  return select(arg0, arg1) ? arg0 : arg1;

    gdzie E = element_type(operand) i reduce_window_without_init działają dokładnie tak samo jak reduce_window z tą różnicą, że schedule bazowego elementu reduce (patrz reduce) nie zawiera wartości init. Obecnie nie jest określone, co się stanie, jeśli odpowiednie okno nie ma wartości (#731).

  • result[result_index] = reduce([source_values], [init_value], [0], scatter), gdzie:

    • source_values = [source[source_index] for source_index in source_indices].
    • selected_index(source_index) = operand_index, jeśli selected_values[source_index] zawiera element operand z operand_index.
    • source_indices = [source_index for source_index in indices(source) if selected_index(source_index) = result_index].

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor lub tensor kwantyzowany na intensywność (C1–C4), (C6), (C8-C11)
(I2) source tensor lub tensor kwantyzowany na intensywność (C1) (C2)
(I3) init_value tensor 0-wymiarowy lub tensor skwantyzowany na intensywność (C3)
(I4) window_dimensions Jednowymiarowa stała tensorowa typu si64 [C2], [C4], [C5]
(I5) window_strides Jednowymiarowa stała tensorowa typu si64 [C2], [C6], [C7]
(I6) padding Dwuwymiarowa stała tensorowa typu si64 (C2) (C8)
(I7) select funkcja (C9)
(I8) scatter funkcja (C10)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C11-C12)

Ograniczenia

  • (C1) element_type(operand) = element_type(source).
  • (C2) shape(source) = num_windows, gdzie:
    • padded_operand_shape = padding[:, 0] + shape(operand) + padding[:, 1].
    • is_empty_window = padded_operand_shape = 0 || window_dimensions > padded_operand_shape.
    • num_windows = is_empty_window ? 0 : floor((padded_operand_shape - window_dimensions) / window_strides) + 1.
  • (C3) element_type(init_value) = element_type(operand).
  • (C4) size(window_dimensions) = rank(operand).
  • (C5) 0 < window_dimensions.
  • (C6) size(window_strides) = rank(operand).
  • (C7) 0 < window_strides.
  • (C8) shape(padding) = [rank(operand), 2].
  • (C9) select ma typ (tensor<E>, tensor<E>) -> tensor<i1>, gdzie E = element_type(operand).
  • (C10) scatter ma typ (tensor<E>, tensor<E>) -> tensor<E>, gdzie is_promotable(element_type(operand), E).
  • (C11) shape(operand) = shape(result).
  • (C12) element_type(result) = E.

Przykłady

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

Więcej przykładów

wyślij

Semantyka

Wysyła inputs do kanału channel_id i generuje token result.

Jeśli is_host_transfer ma wartość true, operacja przenosi dane do hosta. W przeciwnym razie dane zostaną przeniesione na inne urządzenie. To znaczy, że jest ono definiowane przez wdrożenie. Ta flaga powiela informacje podane w polu channel_type, więc w przyszłości planujemy zachować tylko jedną z nich (#666).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) inputs zmienna tensorów lub tensorów skwantyzowanych
(I2) token token
(I3) channel_id stała typu si64
(I4) channel_type wyliczenie DEVICE_TO_DEVICE i DEVICE_TO_HOST (C1)
(I5) is_host_transfer stała typu i1 (C1)

Wyniki

Nazwa Typ
result token

Ograniczenia

  • (C1) channel_type jest zdefiniowany jako:
    • DEVICE_TO_HOST, jeśli is_host_transfer = true,
    • W przeciwnym razie DEVICE_TO_DEVICE.

Przykłady

%result = "stablehlo.send"(%operand, %token) {
  channel_handle = #stablehlo.channel_handle<handle = 1, type = 2>,
  is_host_transfer = true
} : (tensor<2x2xi64>, !stablehlo.token) -> !stablehlo.token

Więcej przykładów

shift_left

Semantyka

Wykonuje operację przesunięcia w lewo na tensorze lhs o określoną liczbę bitów (rhs) i generuje tensor result.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor typu liczby całkowitej (C1)
(I2) rhs tensor typu liczby całkowitej (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu liczby całkowitej (C1)

Ograniczenia

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

Przykłady

// %lhs: [-1, 0, 1]
// %rhs: [1, 2, 3]
%result = "stablehlo.shift_left"(%lhs, %rhs): (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
// %result: [-2, 0, 8]

Więcej przykładów

shift_right_arithmetic

Semantyka

Wykonuje operacje arytmetyczne przesunięcia w prawo na tensorze lhs o określoną liczbę bitów (rhs) i generuje tensor result.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor typu liczby całkowitej (C1)
(I2) rhs tensor typu liczby całkowitej (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu liczby całkowitej (C1)

Ograniczenia

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

Przykłady

// %lhs: [-1, 0, 8]
// %rhs: [1, 2, 3]
%result = "stablehlo.shift_right_arithmetic"(%lhs, %rhs): (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
// %result: [-1, 0, 1]

Więcej przykładów

shift_right_logical

Semantyka

Wykonuje operację logicznego przesunięcia w prawo dla tensora lhs o określoną liczbę bitów (rhs) i generuje tensor result.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor typu liczby całkowitej (C1)
(I2) rhs tensor typu liczby całkowitej (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu liczby całkowitej (C1)

Ograniczenia

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

Przykłady

// %lhs: [-1, 0, 8]
// %rhs: [1, 2, 3]
%result = "stablehlo.shift_right_logical"(%lhs, %rhs): (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
// %result: [9223372036854775807, 0, 1]

Więcej przykładów

podpisywanie

Semantyka

Zwraca znak elementu operand i generuje tensor result. Bardziej formalnie semantyka każdego elementu x można wyrazić za pomocą składni Pythona w ten sposób:

def sign(x):
  if is_integer(x):
    if compare(x, 0, LT, SIGNED): return -1
    if compare(x, 0, EQ, SIGNED): return 0
    return 1
  elif is_float(x):
    if is_nan(x): return NaN
    if compare(x, -0.0, EQ, FLOAT): return -0.0
    if compare(x, +0.0, EQ, FLOAT): return +0.0
    if compare(x, 0.0, LT, FLOAT): return -1.0
    return 1.0
  elif is_complex(x):
    if is_nan(real(x)) or is_nan(imag(x)): return (NaN, NaN)
    if compare(x, (0.0, 0.0), EQ, FLOAT): return (0.0, 0.0)
    return divide(x, convert(abs(x), type(x)))

W przypadku typów skwantyzowanych wykonuje dequantize_op_quantize(sign, operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor liczby całkowitej, zmiennoprzecinkowej, zespolonej lub tensora skwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor liczby całkowitej, zmiennoprzecinkowej, zespolonej lub tensora skwantyzowanego według intensywności (C1)

Ograniczenia

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

Przykłady

// Logical values: +NaN, -1.0, -0.0, +0.0, 1.0
// operand: [0x7FFFFFFFFFFFFFFF, -1.0, -0.0, 0.0, 1.0]
%result = "stablehlo.sign"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// Logical values: +NaN, -1.0, -0.0, +0.0, 1.0
// %result: [0x7FFFFFFFFFFFFFFF, -1.0, -0.0, 0.0, 1.0]

Więcej przykładów

sinus

Semantyka

Wykonuje operację sinusową elementu na tensorze operand i generuje tensor result. W zależności od typu elementu wykona te działania:

  • Liczba zmiennoprzecinkowa: sin (IEEE-754).
  • W przypadku liczb zespolonych: sinus zespolony.
  • W przypadku typów skwantyzowanych: dequantize_op_quantize(sine, operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

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.sine"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[0.0, 1.0], [0.0, -1.0]]

Więcej przykładów

wycinek

Semantyka

Wyodrębnia wycinek z danych operand za pomocą statycznie obliczonych indeksów początkowych i generuje tensor result. start_indices zawiera początkowe indeksy wycinka dla każdego wymiaru, limit_indices – indeksy końcowe (wyłącznie) wycinka dla każdego wymiaru, a strides – kroki dla każdego wymiaru.

Bardziej formalnie: result[result_index] = operand[operand_index], gdzie operand_index = start_indices + result_index * strides.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor lub tensor kwantyzowany na intensywność (C1-C3), (C5)
(I2) start_indices Jednowymiarowa stała tensorowa typu si64 [C2], [C3], [C5]
(I3) limit_indices Jednowymiarowa stała tensorowa typu si64 [C2], [C3], [C5]
(I4) strides Jednowymiarowa stała tensorowa typu si64 (C2) (C4)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor kwantyzowany na intensywność (C1), [C5]

Ograniczenia

  • (C1) element_type(operand) = element_type(result).
  • (C2) size(start_indices) = size(limit_indices) = size(strides) = rank(operand).
  • (C3) 0 <= start_indices <= limit_indices <= shape(operand).
  • (C4) 0 < strides.
  • (C5) shape(result) = ceil((limit_indices - start_indices) / strides).

Przykłady

// %operand: [
//            [0, 0, 0, 0],
//            [0, 0, 1, 1],
//            [0, 0, 1, 1]
//           ]
%result = "stablehlo.slice"(%operand) {
  start_indices = dense<[1, 2]> : tensor<2xi64>,
  limit_indices = dense<[3, 4]> : tensor<2xi64>,
  strides = dense<1> : tensor<2xi64>
} : (tensor<3x4xi64>) -> tensor<2x2xi64>
// % result: [
//            [1, 1],
//            [1, 1]
//           ]

Więcej przykładów

sortuj

Semantyka

Sortuje jednowymiarowe wycinki kolumny inputs wzdłuż wymiaru dimension według kolumny comparator i generuje results.

W odróżnieniu od podobnych danych wejściowych w innych operacjach dimension dopuszcza wartości ujemne, o semantyce opisanych poniżej. W przyszłości może to być niedozwolone ze względu na spójność (#1377).

Jeśli is_stable ma wartość prawda, sortowanie jest stabilne, czyli względna kolejność elementów uważana za równe przez komparator. W przypadku jednej wartości wejściowej 2 elementy e1 i e2 są uznawane przez komparator tylko wtedy, gdy comparator(e1, e2) = comparator(e2, e1) = false. Zapoznaj się z formalizacją poniżej, aby dowiedzieć się, jak to uogólnia się na wiele danych wejściowych.

Oficjalnie dla wszystkich result_index w index_space(results[0]):

  • adjusted_dimension = dimension >= 0 ? dimension : rank(inputs[0]) + dimension.
  • result_slice = [ri0, ..., :, ..., riR-1], gdzie riN to poszczególne elementy w obrębie result_index, a : jest wstawiony w miejscu adjusted_dimension.
  • inputs_together = (inputs[0]..., ..., inputs[N-1]...).
  • results_together[result_slice] = sort(inputs_together[result_slice], comparator_together).
  • gdzie sort sortuje jednowymiarowy wycinek w kolejności niemalejącej, spodziewając się, że funkcja comparator_together zwraca wartość true, jeśli argument po lewej stronie jest mniejszy niż drugi argument po prawej stronie.

  • def comparator_together(lhs_together, rhs_together):
  args = []
  for (lhs_el, rhs_el) in zip(lhs_together, rhs_together):
    args.append(lhs_el)
    args.append(rhs_el)
  return comparator(*args)

  • (results[0]..., ..., results[N-1]...) = results_together.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) inputs zmienna liczba tensorów lub tensorów kwantyzowanych na intensywność (C1–C5)
(I2) dimension stała typu si64 (C4)
(I3) is_stable stała typu i1
(I4) comparator funkcja (C5)

Wyniki

Nazwa Typ Ograniczenia
results zmienna liczba tensorów lub tensorów kwantyzowanych na intensywność (C2) (C3)

Ograniczenia

  • (C1) 0 < size(inputs).
  • (C2) type(inputs...) = type(results...).
  • (C3) same(shape(inputs...) + shape(results...)).
  • (C4) -R <= dimension < R, gdzie R = rank(inputs[0]).
  • (C5) comparator ma typ (tensor<E1>, tensor<E1>, ..., tensor<EN-1>, tensor<EN-1>) -> tensor<i1>, gdzie Ei = element_type(inputs[i]).

Przykłady

// %input0 = [[1, 2, 3], [3, 2, 1]]
// %input1 = [[3, 2, 1], [1, 2, 3]]
%result0, %result1 = "stablehlo.sort"(%input0, %input1) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>, %arg2: tensor<i64>, %arg3: tensor<i64>):
    %predicate = "stablehlo.compare"(%arg0, %arg1) {
      comparison_direction = #stablehlo<comparison_direction GT>
    } : (tensor<i64>, tensor<i64>) -> tensor<i1>
    "stablehlo.return"(%predicate) : (tensor<i1>) -> ()
}) {
  dimension = 0 : i64,
  is_stable = true
} : (tensor<2x3xi64>, tensor<2x3xi64>) -> (tensor<2x3xi64>, tensor<2x3xi64>)
// %result0 = [[3, 2, 3], [1, 2, 1]]
// %result1 = [[1, 2, 1], [3, 2, 3]]

Więcej przykładów

sqrt

Semantyka

Wykonuje operację pierwiastka kwadratowego z elementu na tensorze operand i generuje tensor result. W zależności od typu elementu wykona te działania:

  • Liczba zmiennoprzecinkowa: squareRoot (IEEE-754).
  • W przypadku liczb zespolonych: pierwiastek zespolony.
  • W przypadku typów skwantyzowanych: dequantize_op_quantize(sqrt, operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Ograniczenia

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

Przykłady

// %operand: [[0.0, 1.0], [4.0, 9.0]]
%result = "stablehlo.sqrt"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[0.0, 1.0], [2.0, 3.0]]

Więcej przykładów

odejmowanie

Semantyka

Wykonuje odejmowanie od elementów 2 tensorów lhs i rhs i generuje tensor result. W zależności od typu elementu wykona te działania:

  • W przypadku liczb całkowitych: odejmowanie liczb całkowitych.
  • Liczba zmiennoprzecinkowa: subtraction (IEEE-754).
  • W przypadku liczb zespolonych: odejmowanie zespolone.
  • W przypadku typów skwantyzowanych:
    • dequantize_op_quantize(subtract, lhs, rhs, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor typu całkowitej, zmiennoprzecinkowej, zespolonej lub tensora skwantyzowanego według intensywności (C1)
(I2) rhs tensor typu całkowitej, zmiennoprzecinkowej, zespolonej lub tensora skwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu całkowitej, zmiennoprzecinkowej, zespolonej lub tensora skwantyzowanego według intensywności (C1)

Ograniczenia

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

Przykłady

// %lhs: [[6, 8], [10, 12]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.subtract"(%lhs, %rhs) : (tensor<2x2xf32>, tensor<2x2xf32>) -> (tensor<2x2xf32>)
// %result: [[1, 2], [3, 4]]

Więcej przykładów

Tanh

Semantyka

Wykonuje operację tangensa hiperbolicznego na poziomie pierwiastków na tensorze operand i generuje tensor result. W zależności od typu elementu wykona te działania:

  • Liczba zmiennoprzecinkowa: tanh (IEEE-754).
  • Dla liczb zespolonych: tangens hiperboliczny.
  • W przypadku typów skwantyzowanych:
    • dequantize_op_quantize(tanh, operand, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Ograniczenia

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

Przykłady

// %operand: [-1.0, 0.0, 1.0]
%result = "stablehlo.tanh"(%operand) : (tensor<3xf32>) -> tensor<3xf32>
// %result: [-0.76159416, 0.0, 0.76159416]

Więcej przykładów

transponować

Semantyka

Permituje wymiary tensora operand za pomocą funkcji permutation i tworzy tensor result. Oficjalnie: result[result_index] = operand[operand_index], gdzie result_index[d] = operand_index[permutation[d]].

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor lub tensor skwantyzowany (C1–C4)
(I2) permutation Jednowymiarowa stała tensorowa typu si64 (C2–C4)

Wyniki

Nazwa Typ Ograniczenia
result tensor lub tensor skwantyzowany [C1] (C3-C4)

Ograniczenia

  • (C1) Wartość element_type(result) jest obliczana przez:
    • element_type(operand), jeśli !is_per_axis_quantized(operand).
    • element_type(operand) z tym wyjątkiem, że quantization_dimension(operand) i quantization_dimension(result) mogą się różnić. W przeciwnym razie te wartości mogą się różnić.
  • (C2) permutation to permutacja wartości range(rank(operand)).
  • (C3) shape(result) = dim(operand, permutation...).
  • (C4) Jeśli is_per_axis_quantized(result), to quantization_dimension(operand) = permutation(quantization_dimension(result)).

Przykłady

// %operand: [
//            [[1,2], [3,4], [5,6]],
//            [[7,8], [9,10], [11,12]]
//           ]
%result = "stablehlo.transpose"(%operand) {
  permutation = dense<[2, 1, 0]> : tensor<3xi64>
} : (tensor<2x3x2xi32>) -> tensor<2x3x2xi32>
// %result: [
//           [[1,7], [3,9], [5,11]],
//           [[2,8], [4,10], [6,12]]
//          ]

Więcej przykładów

triangular_solve

Semantyka

Rozwiązywanie wsadów układów równań liniowych o matrycach współczynnika dolnego lub górnego trójkąta.

Bardziej formalnie, biorąc pod uwagę a i b, result[i0, ..., iR-3, :, :] jest rozwiązaniem op(a[i0, ..., iR-3, :, :]) * x = b[i0, ..., iR-3, :, :], gdy left_side to true lub x * op(a[i0, ..., iR-3, :, :]) = b[i0, ..., iR-3, :, :], gdy left_side ma wartość false, analizując zmienną x, gdzie op(a) jest określana przez transpose_a, która może być jedną z tych wartości:

  • NO_TRANSPOSE: wykonaj operację, używając parametru a w niezmienionej postaci.
  • TRANSPOSE: wykonaj operację transponowania a.
  • ADJOINT: wykonaj operację na transponowaniu sprzężonego elementu a.

Dane wejściowe są odczytywane tylko z dolnego trójkąta a, jeśli lower to true, lub z górnego trójkąta a, w przeciwnym razie. Dane wyjściowe są zwracane w tym samym trójkącie, a wartości w drugim trójkącie są definiowane przez implementację.

Jeśli zasada unit_diagonal ma wartość prawda, implementacja może zakładać, że ukośne elementy elementu a mają wartość 1. W przeciwnym razie zachowanie jest nieokreślone.

W przypadku typów skwantyzowanych wykonuje dequantize_op_quantize(lambda x, y: triangular_solve(x, y, left_side, lower, unit_diagonal, transpose_a), a, b, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) a tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1–C3)
(I2) b tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1–C4)
(I3) left_side stała typu i1 (C3)
(I4) lower stała typu i1
(I5) unit_diagonal stała typu i1
(I6) transpose_a wyliczenie NO_TRANSPOSE, TRANSPOSE i ADJOINT

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego lub zespolonego albo tensora kwantyzowanego według intensywności (C1)

Ograniczenia

  • (C1) baseline_element_type(a) = baseline_element_type(b).
  • (C2) 2 <= rank(a) = rank(b) = R.
  • (C3) Relacja między shape(a) a shape(b) jest zdefiniowana w następujący sposób:
    • shape(a)[:-3] = shape(b)[:-3].
    • dim(a, -2) = dim(a, -1) = dim(b, left_side ? -2 : -1).
  • (C4) baseline_type(b) = baseline_type(result).

Przykłady

// %a = [
//       [1.0, 0.0, 0.0],
//       [2.0, 4.0, 0.0],
//       [3.0, 5.0, 6.0]
//      ]
// %b = [
//       [2.0, 0.0, 0.0],
//       [4.0, 8.0, 0.0],
//       [6.0, 10.0, 12.0]
//      ]
%result = "stablehlo.triangular_solve"(%a, %b) {
  left_side = true,
  lower = true,
  unit_diagonal = false,
  transpose_a = #stablehlo<transpose NO_TRANSPOSE>
} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
// %result: [
//           [2.0, 0.0, 0.0],
//           [0.0, 2.0, 0.0],
//           [0.0, 0.0, 2.0]
//          ]

tuple

Semantyka

Generuje krotkę result z wartości val.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) val zmienna liczba wartości (C1)

Wyniki

Nazwa Typ Ograniczenia
result tuple (C1)

Ograniczenia

  • (C1) result ma typ tuple<E0, ..., EN-1>, gdzie Ei = type(val[i]).

Przykłady

// %val0: [1.0, 2.0]
// %val1: (3)
%result = "stablehlo.tuple"(%val0, %val1) : (tensor<2xf32>, tuple<tensor<i32>>) -> tuple<tensor<2xf32>, tuple<tensor<i32>>>
// %result: ([1.0, 2.0], (3))

Więcej przykładów

uniform_dequantize

Semantyka

Wykonuje konwersję skwantyzowanego tensora operand na tensora zmiennoprzecinkowego result zgodnie z parametrami kwantyzacji określonymi przez typ operand.

Oficjalnie: result = dequantize(operand).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor skwantyzowany (C1) (C2)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu zmiennoprzecinkowego (C1) (C2)

Ograniczenia

  • (C1) shape(operand) = shape(result).
  • (C2) element_type(result) = expressed_type(operand).

Przykłady

// %operand: [10, 10]
%result = "stablehlo.uniform_dequantize"(%operand) : (tensor<2x!quant.uniform<i8:f32:0, {0.1:-30,0.5:-20}>>) -> tensor<2xf32>
// %result: [4.0, 15.0]

uniform_quantize

Semantyka

Konwertuje na poziomie elementów tensora zmiennoprzecinkowego lub tensora skwantyzowanego operand na tensor skwantyzowany result zgodnie z parametrami kwantyzacyjnymi określonymi przez typ result.

Bardziej formalnie.

  • Jeśli is_float(operand):
    • result = quantize(operand, type(result)).
  • Jeśli is_quantized(operand):
    • float_result = dequantize(operand).
    • result = quantize(float_result, type(result)).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand tensor typu zmiennoprzecinkowego lub kwantyzowanego (C1) (C2)

Wyniki

Nazwa Typ Ograniczenia
result tensor skwantyzowany (C1) (C2)

Ograniczenia

  • (C1) shape(operand) = shape(result).
  • (C2) expressed_type(result) = is_float(operand) ? element_type(operand) : expressed_type(operand).

Przykłady

// %operand: [4.0, 15.0]
%result = "stablehlo.uniform_quantize"(%operand) : (tensor<2xf32>) -> tensor<2x!quant.uniform<i8:f32:0, {0.1:-30,0.5:-20}>>
// %result: [10, 10]

// %operand: [10, 10]
%result = "stablehlo.uniform_quantize"(%operand) : (tensor<2x!quant.uniform<i8:f32:0, {0.1:-30,0.5:-20}>>) -> tensor<2x!quant.uniform<i8:f32:0, {0.1:-20,0.2:-30}>>
// %result: [20, 45]

podczas

Semantyka

Tworzy dane wyjściowe podczas wykonywania funkcji body co najmniej 0 razy, gdy funkcja cond zwraca wartość true. Bardziej formalnie semantykę można wyrazić w składni Pythona w ten sposób:

internal_state = operand
while cond(*internal_state):
  internal_state = body(*internal_state)
results = internal_state

Zachowanie pętli nieskończonej jest do ustalenia (#383).

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) operand zmienna tensorów, tensorów kwantowych lub tokenów (C1–C3)
(I2) cond funkcja (C1)
(I3) body funkcja (C2)

Wyniki

Nazwa Typ Ograniczenia
results zmienna tensorów, tensorów kwantowych lub tokenów (C3)

Ograniczenia

  • (C1) cond ma typ (T0, ..., TN-1) -> tensor<i1>, gdzie Ti = type(operand[i]).
  • (C2) body ma typ (T0, ..., TN-1) -> (T0, ..., TN-1), gdzie Ti = type(operand[i]).
  • (C3) type(results...) = type(operand...).

Przykłady

// %init_i: 1
// %init_sum: 0
// %one: 1
// %ten: 10
%results0, %results1 = "stablehlo.while"(%init_i, %init_sum) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %cond = "stablehlo.compare"(%arg0, %ten) {
      comparison_direction = #stablehlo<comparison_direction LT>
    } : (tensor<i64>, tensor<i64>) -> tensor<i1>
    stablehlo.return %cond : tensor<i1>
  }, {
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %new_sum = stablehlo.add %arg1, %one : tensor<i64>
    %new_i = stablehlo.add %arg0, %one : tensor<i64>
    stablehlo.return %new_i, %new_sum : tensor<i64>, tensor<i64>
}) : (tensor<i64>, tensor<i64>) -> (tensor<i64>, tensor<i64>)
// %results0: 10
// %results1: 10

Więcej przykładów

XOR

Semantyka

Wykonuje operacje XOR związane z elementami z dwoma tensorami lhs i rhs oraz generuje tensor result. W zależności od typu elementu wykona te działania:

  • W przypadku wartości logicznych: logiczne XOR.
  • W przypadku liczb całkowitych: bitowe XOR.

Dane wejściowe

Etykieta Nazwa Typ Ograniczenia
(I1) lhs tensor typu logicznego lub liczby całkowitej (C1)
(I2) rhs tensor typu logicznego lub liczby całkowitej (C1)

Wyniki

Nazwa Typ Ograniczenia
result tensor typu logicznego lub liczby całkowitej (C1)

Ograniczenia

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

Przykłady

// Bitwise operation with with integer tensors
// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.xor"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[4, 4], [4, 12]]

// Logical operation with with boolean tensors
// %lhs: [[false, false], [true, true]]
// %rhs: [[false, true], [false, true]]
%result = "stablehlo.xor"(%lhs, %rhs) : (tensor<2x2xi1>, tensor<2x2xi1>) -> tensor<2x2xi1>
// %result: [[false, true], [true, false]]

Realizacja

Wykonanie sekwencyjne

Program StableHLO jest wykonywany przez podanie wartości wejściowych do funkcji main i obliczenie wartości wyjściowych. Wartości wyjściowe funkcji oblicza się, wykonując wykres operacji na poziomie głównym dla danej operacji return.

Zlecenie wykonania jest zdefiniowane w ramach implementacji, o ile jest zgodne z Dataflow, czyli jeśli operacje są wykonywane przed ich użyciem. W StableHLO wszystkie działania uboczne wykorzystują 1 token i generują 1 token (wiele tokenów można multipleksować w jeden token za pomocą after_all), więc kolejność wykonywania efektów ubocznych też jest zgodna z przepływem danych. Możliwe zamówienia wykonania w tym przykładowym programie to %0 → %1 → %2 → %3 → %4 → return lub %3 → %0 →%1 → %2 → %4 → return.

Bardziej formalnie proces StabilnyHLO łączy: 1) program StableHLO, 2) stany operacji (jeszcze niewykonano, już wykonano) i 3) wartości pośrednie, na których pracuje proces. Proces zaczyna się od wartości wejściowych do funkcji main, przechodzi przez wykres operacji aktualizowania stanów operacji i wartości pośrednich, a na koniec wartości wyjściowych. Dalsza formalizacja tego procesu zostanie opracowana (#484).

Wykonywanie równoległe

Programy StableHLO można uruchamiać równolegle w formie siatki 2D procesu num_replicas za pomocą num_partitions. Oba programy są typu ui32.

W siatce procesów StableHLO w tym samym czasie wykonywanych jest num_replicas * num_partitions procesów StableHLO. Każdy proces ma unikalny process_id = (replica_id, partition_id), gdzie replica_id w replica_ids = range(num_replicas) i partition_id w partition_ids = range(num_partitions) mają wartość ui32.

Rozmiar siatki procesów jest znany statycznie dla każdego programu (w przyszłości planujemy ustawić go jako część programów StableHLO#650), a położenie w siatce procesów jest znane statycznie dla każdego procesu. Każdy proces ma dostęp do swojej pozycji w siatce procesów za pomocą operacji replica_id i partition_id.

Programy w siatce procesów mogą być takie same (w stylu „Pojedynczy program, wiele danych”), różne (w stylu „Wiele programów, wiele danych”) lub między nimi. W przyszłości planujemy wprowadzić obsługę innych idiomów dotyczących definiowania równoległych programów StableHLO, w tym GSPMD (#619).

W siatce procesów procesy są w większości niezależne od siebie – mają oddzielne stany operacji, osobne wartości wejściowe/pośrednie/wyjściowe, a większość operacji jest realizowana osobno pomiędzy procesami, z wyjątkiem niewielkiej liczby operacji zbiorczych opisanych poniżej.

Biorąc pod uwagę, że większość operacji korzysta tylko z wartości z tego samego procesu, odwoływanie się do tych wartości za pomocą nazw jest zazwyczaj jednoznaczne. Jednak opisanie semantyki działań zbiorowych jest niewystarczające, ponieważ sprawia, że notacja name@process_id odnosi się do wartości name w obrębie konkretnego procesu. (z tego punktu widzenia niezakwalifikowane name to skrótowiec name@(replica_id(), partition_id())).

Kolejność wykonywania w różnych procesach jest zdefiniowana w ramach implementacji, z wyjątkiem synchronizacji spowodowanej komunikacją między punktami oraz operacjami zbiorowymi opisanymi poniżej.

Komunikacja między punktami

Procesy StableHLO mogą komunikować się ze sobą za pomocą kanałów StableHLO. Kanał jest reprezentowany przez dodatni identyfikator typu si64. Różne działania umożliwiają wysyłanie wartości do kanałów i odbieranie ich z kanałów.

Dalsza formalizacja, np. skąd pochodzą identyfikatory kanałów, sposób informowania o nich przez programy i rodzaj synchronizacji, to do ustalenia (#484).

Komunikacja strumieniowa

Każdy proces StableHLO ma dostęp do 2 interfejsów strumieniowania:

  • In-Feed, z której można odczytać.
  • Out-Feed, w których można zapisywać.

W odróżnieniu od kanałów, które służą do komunikacji między procesami, a tym samym mają procesy po obu stronach, w przypadku reklam In-Feed i out-Feed są określone inne końcowe implementacje.

Dalsza formalizacja, np. wpływ strumieniowania komunikacji na zamówienie i rodzaj synchronizacji, jest do ustalenia (#484).

Operacje zbiorcze

W StableHLO jest 6 zbiorczych operacji: all_gather, all_reduce, all_to_all, collective_broadcast, collective_permute i reduce_scatter. Wszystkie te operacje rozdzielają procesy w siatce procesów StableHLO na grupy procesów StableHLO i wykonują wspólne obliczenia w każdej grupie procesów niezależnie od innych grup procesów.

W każdej grupie procesów wspólne operacje mogą wprowadzać barierę synchronizacji. Dalsza formalizacja, np. doprecyzowanie, kiedy dokładnie następuje synchronizacja, w jaki sposób procesy docierają do tej bariery oraz co się dzieje w przypadku braku tej bariery, jest ustalona (#484).

Jeśli grupa procesów obejmuje komunikację między partycją, czyli w grupie procesów istnieją procesy, których identyfikatory partycji są różne, wykonanie operacji zbiorczej wymaga kanału, a operacja zbiorowa musi zapewniać dodatnią wartość channel_id typu si64. Komunikacja między replikami nie wymaga kanałów.

Obliczenia wykonywane przez zbiorcze działania dotyczą tylko poszczególnych operacji i zostały opisane powyżej w sekcjach dotyczących poszczególnych operacji. Jednak strategie, według których siatka procesów jest dzielona na grupy procesów, są wspólne dla tych operacji i zostały opisane w tej sekcji. Ogólnie rzecz biorąc, StableHLO obsługuje 4 strategie.

cross_replica

W ramach każdej grupy procesów zachodzi tylko komunikacja między replikami. Ta strategia wykorzystuje replica_groups – listę list identyfikatorów replik – i oblicza iloczyn kartezjański replica_groups według partition_ids. replica_groups musi mieć unikalne elementy i zawierać wszystkie replica_ids. Bardziej formalnie, używając składni Python:

def cross_replica(replica_groups: List[List[ReplicaId]]) -> List[List[ProcessId]]:
  for replica_group in replica_groups:
    for partition_id in partition_ids:
      process_group = []
      for replica_id in replica_group:
        process_group.append((replica_id, partition_id))
      yield process_group

Na przykład w przypadku replica_groups = [[0, 1], [2, 3]] i num_partitions = 2 cross_replica wygeneruje [[(0, 0), (1, 0)], [(0, 1), (1, 1)], [(2, 0), (3, 0)], [(2, 1), (3, 1)]].

cross_partition

W ramach każdej grupy procesów zachodzi tylko komunikacja między partycjami. Ta strategia wykorzystuje partition_groups – listę identyfikatorów partycji – i oblicza iloczyn kartezjański partition_groups według replica_ids. Element partition_groups musi zawierać unikalne elementy i zakrywać cały partition_ids. Bardziej formalnie, z użyciem składni Pythona:

def cross_partition(partition_groups: List[List[PartitionId]]) -> List[List[ProcessId]]:
  for partition_group in partition_groups:
    for replica_id in replica_ids:
      process_group = []
      for partition_id in partition_group:
        process_group.append((replica_id, partition_id))
      yield process_group

Na przykład w przypadku partition_groups = [[0, 1]] i num_replicas = 4 cross_partition wygeneruje [[(0, 0), (0, 1)], [(1, 0), (1, 1)], [(2, 0), (2, 1)], [(3, 0), (3, 1)]].

cross_replica_and_partition

W każdej grupie procesów może zachodzić zarówno komunikacja między replikami, jak i komunikacją między partycją. Ta strategia wykorzystuje replica_groups – listę list identyfikatorów replik – i oblicza iloczyn kartezjańskie każdego elementu replica_group za pomocą partition_ids. Element replica_groups musi zawierać unikalne elementy i zastosować wszystkie elementy replica_ids. Bardziej formalnie, z użyciem składni Pythona:

def cross_replica_and_partition(replica_groups: List[List[ReplicaId]]) -> List[List[ProcessId]]:
  for replica_group in replica_groups:
    process_group = []
    for partition_id in partition_ids:
      for replica_id in replica_group:
        process_group.append((replica_id, partition_id))
    yield process_group

Na przykład w przypadku replica_groups = [[0, 1], [2, 3]] i num_partitions = 2 cross_replica_and_partition wygeneruje [[(0, 0), (1, 0), (0, 1), (1, 1)], [(2, 0), (3, 0), (2, 1), (3, 1)]].

flattened_ids

Ta strategia pobiera flattened_id_groups – listę list „spłaszczonych” identyfikatorów procesów w postaci replica_id * num_partitions + partition_id – i przekształca je w identyfikatory procesów. Element flattened_id_groups musi zawierać unikalne elementy i zakrywać cały process_ids. Bardziej formalnie, z użyciem składni Pythona:

def flattened_ids(flattened_id_groups: List[List[ui32]]) -> List[List[ProcessId]]:
  for flattened_id_group in flattened_id_groups:
    process_group = []
    for flattened_id in flattened_id_group:
      replica_id = flattened_id // num_partitions
      partition_id = flattened_id % num_partitions
      process_group.append((replica_id, partition_id))
    yield process_group

Na przykład w przypadku flattened_id_groups = [[0, 1, 2, 3], [4, 5, 6, 7]], num_replicas = 4 i num_partitions = 2 flattened_ids wygeneruje [[(0, 0), (0, 1), (1, 0), (1, 1)], [(2, 0), (2, 1), (3, 0), (3, 1)]].

Dokładność

Obecnie StableHLO nie gwarantuje dokładności liczbowej, ale w przyszłości może się to zmienić (#1156).

Błędy

Programy StableHLO są sprawdzane przez obszerny zestaw ograniczeń dla poszczególnych operacji, co eliminuje wiele klas błędów przed czasem działania. Nadal jednak możliwe są warunki błędu, na przykład w wyniku przekroczenia liczby całkowitej, uzyskania dostępu poza granicami itp. O ile wyraźnie nie zostaną wywołane, wszystkie te błędy powodują zachowanie zdefiniowane w implementacji, ale w przyszłości może to ulec zmianie (#1157).

Wyjątkiem od tej reguły są wyjątki zmiennoprzecinkowe w programach StableHLO. Operacje, które stanowią wyjątki zdefiniowane przez standard IEEE-754 (nieprawidłowa operacja, dzielenie przez zero, przepełnienie, niedopełnienie lub niedokładne wyjątki), dają wyniki domyślne (zgodnie z definicją w standardzie) i kontynuują wykonywanie bez podnoszenia odpowiedniej flagi stanu, podobnie jak w przypadku obsługi wyjątku raiseNoFlag ze standardu. Wyjątki w przypadku niestandardowych operacji (np. złożonych funkcji arytmetycznych i pewnych funkcji transcendentalnych) są definiowane przez implementację.

Zapis

Do opisania składni w tym dokumencie wykorzystano zmodyfikowaną składnię ISO (ISO/IEC 14977:1996, Wikipedia) z 2 zmianami: 1) reguły są definiowane za pomocą ::=, a nie =;

2) konkatenacja jest wyrażana przez dopasowanie, a nie ,.

Do opisywania semantyki (np. w sekcjach „Typy”, „Stałe” i „Operacje”) używamy formuł opartych na składni Pythona poszerzonej o obsługę zwięzłego wyrażania operacji na tablicach w opisany poniżej sposób. Sprawdza się to w przypadku małych fragmentów kodu, ale w rzadkich przypadkach, gdy potrzebne są większe fragmenty kodu, stosujemy składnię vanilla Python, która zawsze jest wprowadzona w sposób jawny.

Wzory

Na podstawie przykładu ze specyfikacji dot_general przeanalizujmy, jak działają formuły. Jedno z ograniczeń tej operacji wygląda tak: dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...).

Nazwy używane w tej formule pochodzą z 2 źródeł: 1) funkcji globalnych, tj. dim, 2) definicji członkowskich odpowiedniego elementu programu, czyli danych wejściowych lhs, lhs_batching_dimensions, rhs i rhs_batching_dimensions zdefiniowanych w sekcji „Dane wejściowe” w dot_general.

Jak już wspomnieliśmy, składnia tej formuły jest oparta na Pythonie z kilkoma rozszerzeniami zorientowanymi na zwięzłość. Aby formuła była zrozumiała, przekonwertujmy ją na składnię waniliową Pythona.

A) W tych formułach do reprezentowania równości używamy =, więc pierwszym krokiem do uzyskania składni Pythona jest zastąpienie elementu = wartością == w następujący sposób: dim(lhs, lhs_batching_dimensions...) == dim(rhs, rhs_batching_dimensions...).

B) Poza tym te formuły obsługują elipsy (...), które zmieniają wyrażenia skalarne w wyrażenia tensorowe. W skrócie f(xs...) oznacza mniej więcej „dla każdego skalarnego x tensora xs oblicz skalarną f(x), a następnie zwróć wszystkie te wyniki skalarne jako wynik tensora”. W składni waniliowej Pythona nasza przykładowa formuła zmienia się w: [dim(lhs, dim1) for dim1 in lhs_batching_dimensions] == [dim(rhs, dim2) for dim2 in rhs_batching_dimensions].

Dzięki elipsom można często uniknąć pracy na poziomie poszczególnych skalarów. W niektórych trudnych przypadkach można jednak użyć składni półinformacyjnej niższego poziomu, jak we wzorze start_indices[bi0, ..., :, ..., biN] ze specyfikacji gather. W ramach usługi zwięzłości nie zapewniamy dokładnego formalizmu przy tłumaczeniu takiej składni na vanilla Python. Mamy nadzieję, że jest ona intuicyjna w każdym przypadku oddzielnie. Daj nam znać, jeśli niektóre formuły wyglądają na nieprzejrzyste, postaramy się je poprawić.

Zauważysz też, że formuły korzystają z elips do rozwijania różnych rodzajów list, m.in. tensorów, list tensorów (które np. mogą wynikać z różnej liczby tensorów) itp.W tym przypadku nie stosujemy dokładnego formalności (np. listy nie są nawet częścią systemu StableHLO w sposób intuicyjny).

C) Ostatnim ciekawym narzędziem informacyjnym, które stosujemy, jest pośrednia transmisja treści. Mimo że ustawienie StableHLO nie obsługuje transmitowania niejawnego, formuły to umożliwiają, także, aby zapewnić zwięzłość. W skrócie, jeśli używany jest skalar w kontekście, gdy spodziewany jest tensor, jest on transmitowany do oczekiwanego kształtu.

Aby kontynuować przykład z dot_general, oto kolejne ograniczenie: 0 <= lhs_batching_dimensions < rank(lhs). Zgodnie z definicją dot_general lhs_batching_dimensions jest tensorem, ale 0 i rank(lhs) są skalarami. Po zastosowaniu transmisji niejawnych formuła będzie miała postać [0, ..., 0] <= lhs_batching_dimensions < [rank(lhs), ..., rank(lhs)].

Jeśli ta formuła zostanie zastosowana do konkretnej operacji dot_general, będzie oceniać jako tensor wartości logicznych. Gdy formuły są używane jako ograniczenia, te ograniczenia są stosowane, gdy formuła zwraca wartość true lub tensora, który ma tylko elementy true.

Nazwy

Zakres leksyki w formułach obejmuje: 1) funkcje globalne, 2) definicje członków,

3) lokalne definicje. Poniżej znajdziesz listę funkcji globalnych. Lista definicji elementów zależy od elementu programu, do którego ma zastosowanie notacja:

  • W przypadku operacji definicje członków obejmują nazwy wprowadzone w sekcjach „Dane wejściowe” i „Dane wyjściowe”.
  • We wszystkich innych przypadkach definicje członków zawierają elementy strukturalne elementu programu, których nazwy pochodzą od odpowiednich elementów nieterminowych EBNF. W większości przypadków nazwy tych części strukturalnych uzyskuje się przez przekształcanie nazw elementów nieterminalowych na typu „snake” (np. IntegerLiteral =>integer_literal), ale czasami nazwy zostają skrócone (np. QuantizationStorageType => storage_type), co oznacza, że są one wprowadzane w sposób wyraźnie podobny do nazw w sekcjach „Dane wejściowe” lub „Dane wyjściowe”.
  • Poza tym definicje członków zawsze zawierają ciąg self, który pozwala odwoływać się do odpowiedniego elementu programu.

Wartości

Podczas oceniania formuł działają one z następującymi typami wartości: 1) Value (wartości rzeczywiste, np. dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>; zawsze znają swój typ), 2) Placeholder (wartości przyszłe, np. lhs, rhs lub result; ich wartości nie są jeszcze znane; są znane tylko ich wartości), 3) Type (typy określone w sekcji „Funkcje”), 4) Function (funkcje globalne).

W zależności od kontekstu nazwy mogą odnosić się do różnych wartości. Dokładniej rzecz ujmując, sekcja „Semantyki” dotycząca operacji (i ich odpowiedniki dla innych elementów programu) definiuje logikę środowiska wykonawczego, więc wszystkie dane wejściowe są dostępne jako Value. Z kolei sekcja „Ograniczenia” dotycząca operacji (i ich odpowiedników) definiuje logikę „czasu kompilacji”, czyli coś, co jest zwykle wykonywane przed środowiskiem wykonawczym. Dlatego jako Value dostępne są tylko stałe dane wejściowe, a inne – tylko jako Placeholder.

Nazwy W sekcji „Semantyki” W sekcji „Ograniczenia”
Funkcje globalne Function Function
Stałe dane wejściowe Value Value
Stałe dane wejściowe Value Placeholder
Wyniki Value Placeholder
Lokalne definicje Zależy od definicji Zależy od definicji

Przyjrzyjmy się przykładowej operacji transpose:

%result = "stablehlo.transpose"(%operand) {
  permutation = dense<[2, 1, 0]> : tensor<3xi64>
} : (tensor<2x3x2xi32>) -> tensor<2x3x2xi32>

W przypadku tej operacji permutation jest stałą, więc jest dostępna jako Value zarówno w semantyce, jak i w ograniczeniach. W przeciwieństwie do elementów operand i result są dostępne jako element Value w zasadzie semantyki, ale tylko jako Placeholder w ograniczeniach.

Funkcje

Konstrukcja typów

Nie ma funkcji, za pomocą których można tworzyć typy. Zamiast tego używamy składni typów bezpośrednio, ponieważ jest ona zwykle bardziej zwięzła. Przykład: (tensor<E>, tensor<E>) -> (tensor<E>) zamiast function_type( [tensor_type([], E), tensor_type([], E)], [tensor_type([], E)]).

Funkcje w typach

  • Funkcja element_type jest zdefiniowana dla typów tensorów i typów i zwróconych tensorów skwantyzowanych odpowiednio jako części TensorElementType lub QuantizedTensorElementType odpowiedniego elementu TensorType lub QuantizedTensorType.
def element_type(x: Value | Placeholder | Type):
 if type(x) == TensorType:
    return tensor_element_type(x)
  if type(x) == QuantizedTensorType:
    return quantized_tensor_element_type(x)
  if type(x) is not Type:
    return element_type(type(x))

  • is_per_axis_quantized(x: Value | Placeholder | Type) -> Value to skrót do is_quantized(x) and quantization_dimension(x) is not None.

  • is_per_tensor_quantized(x: Value | Placeholder | Type) -> Value to skrót nazwy is_quantized(x) and quantization_dimension(x) is None.

  • is_promotable(x: Type, y: Type) -> bool sprawdza, czy typ x może być awansowany na typ y. Gdy x i y mają wartość QuantizedTensorElementType, promocja jest stosowana tylko do storage_type. Ta konkretna wersja promocji jest obecnie używana w kontekście obliczania redukcji wartości (więcej informacji znajdziesz w RFC).

def is_promotable(x: Type, y: Type) -> Value:
  is_same_type = (is_bool(x) and is_bool(y)) or
    (is_integer(x) and is_integer(y)) or (is_float(x) and is_float(y)) or
    (is_complex(x) and is_complex(y)) or
    (is_quantized(x) and is_quantized(y) and expressed_type(x) = expressed_type(y))

  if is_same_type == False:
    return False

  if is_integer(x) or is_float(x):
    return bitwidth(x) <= bitwidth(y)

  if is_complex(x):
    return bitwidth(element_type(x)) <= bitwidth(element_type(y))

  if is_quantized(x):
    return bitwidth(storage_type(x)) <= bitwidth(storage_type(y))

  return false

  • is_quantized(x: Value | Placeholder | Type) -> Value to skrót do is_quantized_tensor_element_type(x).

  • is_type_name(x: Value | Placeholder | Type) -> Value. Dostępne dla wszystkich typów. Na przykład funkcja is_float(x) zwraca wartość true, jeśli x to FloatType. Jeśli x jest wartością lub symbolem zastępczym, ta funkcja jest skrótem do is_type_name(type(x)).

  • max_value(x: Type) -> Value zwraca maksymalną wartość TensorElementType. Jeśli x nie jest wartością TensorElementType, zwraca None.

  • min_value(x: Type) -> Value zwraca minimalną możliwą wartość TensorElementType. Jeśli x nie jest wartością TensorElementType, zwraca None.

  • member_name(x: Value | Placeholder | Type) -> Any. Wartość dostępna w przypadku wszystkich definicji typu member_name każdego typu. Na przykład tensor_element_type(x) zwraca część TensorElementType odpowiedniej TensorType. Jeśli x jest wartością lub symbolem zastępczym, ta funkcja jest skrótem do member_name(type(x)). Jeśli x nie jest typem, który ma odpowiedni element, wartość lub obiekt zastępczy takiego typu, zwraca None.

Konstrukcja wartości

  • operation_name(*xs: Value | Type) -> Value. Dostępne dla wszystkich operacji. Na przykład add(lhs, rhs) przyjmuje 2 wartości tensorów lhs i rhs i zwraca dane wyjściowe oceny operacji add z tymi danymi. W przypadku niektórych operacji, np. broadcast_in_dim, dane wyjściowe to „przeciążenie”, czyli potrzebne do oceny operacji. W tym przypadku funkcja przyjmuje te typy jako argumenty.

Funkcja na wartościach

  • Dostępne są wszystkie operatory i funkcje Pythona. Na przykład adnotacje subskrypcji i wycinania z Pythona są dostępne do indeksowania za pomocą tensorów, kwantyzowanych tensorów i krotek.

  • Funkcja to_destination_type(x: Value, destination_type: Type) -> Value jest określona na parametrach i zwraca przekonwertowaną wartość x na podstawie danych type(x) i destination_type w ten sposób:

def to_destination_type(x: Value, destination_type: Type) -> Value:
  if type(x) == destination_type:
    return x

  if is_quantized(destination_type):
    if is_quantized(type(x)):
      return quantize(x, destination_type)
    assert is_float(type(x))
    return quantize(x, destination_type)

  if is_quantized(type(x)):
    assert destination_type = expressed_type(type(x))
    return dequantize(type(x))

  return convert(x, destination_type)

Obecnie trwają prace nad scalaniem operacji convert, uniform_quantize i uniform_dequantize (#1576). Po scaleniu nie jest potrzebna powyższa funkcja i możemy zamiast niej użyć nazwy operacji dla funkcji convert.

  • is_nan(x: Value) -> Value jest zdefiniowany w tensorach i zwraca true, jeśli wszystkie elementy funkcji x mają wartość NaN lub false. Jeśli x nie jest tensorem, zwraca None.

  • Element is_sorted(x: Value) -> Value jest zdefiniowany w tensorach i zwraca true, jeśli elementy właściwości x są posortowane rosnąco według kolejności leksykograficznej ich indeksów lub w przeciwnym razie zwraca wartość false. Jeśli x nie jest tensorem, zwraca None.

  • is_unique(x: Value) -> Value jest zdefiniowany w tensorach i zwraca true, jeśli x nie ma zduplikowanych elementów. W przeciwnym razie funkcja false nie jest określona. Jeśli x nie jest tensorem, zwraca None.

  • member_name(x: Value) -> Any jest zdefiniowany dla wszystkich definicji elementów member_name wszystkich wartości. Na przykład real_part(x) zwraca część RealPart odpowiedniej ComplexConstant. Jeśli x nie jest wartością, która ma odpowiedniego użytkownika, zwraca None.

  • Funkcja same(x: Value) -> Value jest definiowana w tensorach i zwraca true, jeśli wszystkie elementy funkcji x są sobie równe, a w przeciwnym razie zwraca wartość false. Jeśli tensor nie ma elementów, jest to liczone jako „wszystkie sobie równe”, czyli funkcja zwraca true. Jeśli x nie jest tensorem, zwraca None.

  • Funkcja split(x: Value, num_results: Value, axis: Value) -> Value jest definiowana na tensorach i zwraca wycinki (num_results) wartości x wzdłuż osi axis. Jeśli x nie jest tensorem ani dim(x, axis) % num_results != 0, zwraca None.

Obliczanie kształtów

  • axes(x: Value | Placeholder | Type) -> Value to skrót do range(rank(x)).

  • dim(x: Value | Placeholder | Type, axis: Value) -> Value to skrót do shape(x)[axis].

  • dims(x: Value | Placeholder | Type, axes: List) -> List to skrót do list(map(lambda axis: dim(x, axis), axes)).

  • index_space(x: Value | Placeholder | Type) -> Value jest zdefiniowany dla tensorów i zwraca indeksy size(x) dla odpowiedniego elementu TensorType posortowane w kolejności leksykograficznej, tj. [0, ..., 0], [0, ..., 1], ..., shape(x) - 1. Jeśli x nie jest typem tensora, skwantyzowanym typem tensora, wartością lub zmienną jednego z tych typów, zwraca None.

  • rank(x: Value | Placeholder | Type) -> Value to skrót do size(shape(x)).

  • shape(x: Value | Placeholder | Type) -> Value zdefiniowano w sekcji „Funkcje w typach” za pomocą member_name.

  • size(x: Value | Placeholder | Type) -> Value to skrót do reduce(lambda x, y: x * y, shape(x)).

Obliczenia kwantyzacji

  • def baseline_element_type(x: Value | Placeholder | Type) -> Type to skrót nazwy element_type(baseline_type(x)).

  • Funkcja baseline_type jest definiowana dla typów tensorów i tensorów skwantyzowanych i przekształca je w „element bazowy”, tj. typ o tym samym kształcie, ale z parametrami kwantyzacji typu elementu resetowanego do wartości domyślnych. Jest to przydatna sztuczka umożliwiająca jednolite porównywanie typów tensorów i kwantyzowanych typów tensorów, co jest potrzebne dość często. W przypadku typów skwantyzowanych umożliwia to porównywanie typów z ignorowaniem parametrów kwantyzacji, czyli shape, storage_type, expressed_type, storage_min, storage_max i quantization_dimension (w przypadku typu skwantyzowanego na osi), ale wartości scales i zero points mogą się różnić.

def baseline_type(x: Value | Placeholder | Type) -> Type:
  if type(x) == TensorType:
    return x
  if type(x) == QuantizedTensorType:
    element_type = quantized_tensor_element_type(x)
    baseline_element_type = QuantizedTensorElementType(
      storage_type = storage_type(element_type),
      storage_min = storage_min(element_type),
      storage_max = storage_max(element_type),
      expressed_type = expressed_type(element_type),
      quantization_dimension = quantization_dimension(element_type),
      scales = [constant(1.0, expressed_type(element_type))] * dim(x, quantization_dimension(element_type)),
      zero_points = [constant(0, storage_type(element_type))] * dim(x, quantization_dimension(element_type)))
    return QuantizedTensorType(shape(x), baseline_element_type)
  if type(x) is not Type:
    return baseline_element_type(type(x))
  • Funkcja dequantize jest definiowana w typach tensorów skwantyzowanych i przekształca je w tensorów zmiennoprzecinkowych. Odbywa się to poprzez konwertowanie elementów skwantyzowanych, które reprezentują wartości całkowite rodzaju pamięci masowej, na odpowiednie wartości zmiennoprzecinkowe określonego typu za pomocą punktu zerowego i skali powiązanego z typem elementu skwantyzowanego.
def compute_zero_points(quantized_type, result_type):
  if is_per_tensor_quantized(quantized_type):
    return broadcast_in_dim(constant(zero_point(quantized_type), storage_type(quantized_type)), [], result_type)
  if is_per_axis_quantized(quantized_type):
    for i in index_space(result_type):
      d = quantization_dimension(quantized_type)
      zero_points[i] = zero_points(quantized_type)[i[d]]
    return zero_points

def compute_scales(quantized_type, result_type):
  if is_per_tensor_quantized(quantized_type):
    return broadcast_in_dim(constant(scale(quantized_type), expressed_type(quantized_type)), [],
            type(result_type))
  if is_per_axis_quantized(quantized_type):
    for i in index_space(result_type):
      d = quantization_dimension(quantized_type)
      scales[i] = scales(quantized_type)[i[d]]
    return scales

def dequantize(x: Value) -> Value:
  assert is_quantized(x)
  x_storage = bitcast_convert(x, storage_type(x))
  x_storage_sub = x_storage - compute_zero_points(type(x), type(x_storage))
  x_expressed_sub = convert(x_storage_sub, expressed_type(x))
  return x_expressed_sub * compute_scales(type(x), type(x_expressed_sub))
  • Funkcja quantize jest definiowana dla tensorów zmiennoprzecinkowych i przekształca je w typy tensorów skwantyzowanych. Odbywa się to przez przekształcanie wartości zmiennoprzecinkowych określonego typu na odpowiednie wartości całkowite typu pamięci masowej z wykorzystaniem punktu zerowego i skali związanej z typem elementu skwantyzowanego.
def quantize(x: Value, type: Type) -> Value:
  assert is_float(x) and is_quantized(type)
  x_expressed_rounded = round_nearest_even(x / compute_scales(type, type(x)))
  x_storage_rounded = convert(x_expressed_rounded, storage_type(type))
  x_storage_add = x_storage_rounded + compute_zero_points(type, type(x_storage_rounded))
  x_storage = clamp(storage_min(type), x_storage_add, storage_max(type))
  return bitcast_convert(x_storage, type)
  • Element dequantize_op_quantize służy do określania obliczeń dotyczących elementów na skwantyzowanych tensorach. Dekwantyfikuje, czyli zamienia skwantyzowane elementy w ich typy wyrażone, następnie wykonuje operację, a następnie dokonuje ilości danych, czyli zamienia wyniki z powrotem w ich typy przechowywania. Obecnie ta funkcja działa tylko przy kwantyzacji według intensywności. Trwa kwantyzacja dla poszczególnych osi (#1574).
def dequantize_op_quantize(op, *inputs_and_output_type):
  inputs = inputs_and_output_type[:-1]
  output_type = inputs_and_output_type[-1]

  float_inputs = map(dequantize, inputs)
  float_result = op(*float_inputs)
  return quantize(float_result, output_type)

def dequantize_batch_norm_grad_or_training_quantize(op, *inputs_and_output_types):
  inputs = inputs_and_output_type[:-3]
  float_inputs = map(dequantize, inputs)
  float_results = op(*float_inputs)
  return map(quantize, float_results, inputs_and_output_type[-3:])

def dequantize_compare(lhs, rhs, comparison_direction):
  float_lhs = dequantize(lhs)
  float_rhs = dequantize(rhs)
  return compare(float_lhs, float_rhs, comparison_direction, FLOAT)

def dequantize_select_quantize(pred, on_true, on_false, output_type):
  float_on_true = dequantize(on_true)
  float_on_false = dequantize(on_false)
  float_result = select(pred, float_on_true, float_on_false)
  return quantize(float_result, output_type)

Obliczenia siatki

  • cross_partition(replica_groups: Value) -> Value. Zobacz sekcję „cross_replica” powyżej.

  • cross_replica(replica_groups: Value) -> Value. Zobacz sekcję „cross_replica” powyżej.

  • cross_replica_and_partition(replica_groups: Value) -> Value. Przeczytaj sekcję „cross_replica_and_partition” powyżej.

  • flattened_ids(replica_groups: Value) -> Value. Przeczytaj sekcję „flattened_ids” powyżej.