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)
, tosize(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)
.
- (C12)
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
ifalse
. - Typy liczb całkowitych mogą być podpisane (
si
) lub nieoznaczone (ui
) i mieć jedną z obsługiwanych szerokości bitowych (4
,8
,16
,32
lub64
). Podpisane typysiN
reprezentują wartości całkowite od-2^(N-1)
do2^(N-1)-1
, a nieoznaczone typyuiN
reprezentują wartości całkowite z zakresu od0
do2^N-1
włącznie. - Możliwe typy liczby zmiennoprzecinkowej:
- Typy
f8E4M3FN
if8E5M2
odpowiadające odpowiednio kodowaniuE4M3
iE5M2
w formacie FP8 opisanym w artykule Formaty FP8 na potrzeby deep learning. - Typy
f8E4M3FNUZ
if8E5M2FNUZ
odpowiadające kodomE4M3
iE5M2
w formatach FP8 opisanych w artykule 8-bitowe formaty liczbowe dla głębokich sieci neuronowych. - Typ
f8E4M3B11FNUZ
odpowiadający kodowaniuE4M3
formatów FP8 opisanych w artykule Hybrydowe 8-bitowe trenowanie zmiennoprzecinkowe (HFP8) i wnioskowanie dla głębokich sieci neuronowych. - Typ
bf16
odpowiadający formatowibfloat16
opisanemu w artykule BFloat16: tajemnica wysokiej wydajności w Cloud TPU. f16
,f32
if64
odpowiadające formatombinary16
(„połowa precyzji”),binary32
(„pojedyncza precyzja”) ibinary64
(„podwójna precyzja”) opisanymi w standardzie IEEE 754.
- Typy
- 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ą typuf32
) icomplex<f64>
(obie części są typuf64
).
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
ifunc.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
iset_dimension_size
(#8). - Obliczanie kształtów, w tym operacje
arith
,shape
itensor
(#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śliis_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]
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]]
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
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ślichannel_id <= 0 and use_global_device_ids = false
.cross_replica_and_partition(replica_groups)
jeślichannel_id > 0 and use_global_device_ids = false
.flattened_ids(replica_groups)
jeślichannel_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 wszystkichreceiver
w grupieprocess_group
.result@process = concatenate(operands@process, all_gather_dim)
dla wszystkichprocess
w grupieprocess_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]]
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ślichannel_id <= 0 and use_global_device_ids = false
.cross_replica_and_partition(replica_groups)
jeślichannel_id > 0 and use_global_device_ids = false
.flattened_ids(replica_groups)
jeślichannel_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 wynosito_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>)
, gdzieis_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]
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ślichannel_id <= 0
.cross_partition(replica_groups)
, jeślichannel_id > 0
.
Następnie w ciągu każdego process_group
:
split_parts@sender = split(operand@sender, split_count, split_dimension)
za wszystkiesender
w grupieprocess_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]]
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]
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
igrad_offset
mają te same wartościbaseline_element_type
. - (C3)
operand
,grad_output
igrad_operand
mają ten sam kształt. - (C4)
scale
,mean
,variance
,grad_scale
igrad_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
iresult
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
ioutput
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)
iR = 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 wszystkie0 <= 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 wszystkie0 <= i < R
.dim(operand, R - 1) * num_bits(E) = num_bits(E')
.
- Jeśli
- (C2) Jeśli
is_complex(operand) or is_complex(result)
, tois_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
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ślidim(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)
izero_points(operand)
mogą się różnić od odpowiedziquantization_dimension(result)
,scales(result)
izero_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
waxes(operand)
:dim(operand, d) = 1
lubdim(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
, toscales(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]
// ]
// ]
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śli0 <= 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]
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]
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]
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]
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ślichannel_id <= 0
.cross_partition(replica_groups)
, jeślichannel_id > 0
.
Następnie result@process
przyznaje:
operand@process_groups[i, 0]
, jeśli istnieje dyrektywai
, która oznacza, że proces jest w elemencieprocess_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
, gdzieN
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ślichannel_id <= 0
.cross_partition(source_target_pairs)
, jeślichannel_id > 0
.
Następnie result@process
przyznaje:
operand@process_groups[i, 0]
, jeśli istniejei
taki jakprocess_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
, gdzieN
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]]
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śliis_signed_integer(element_type(lhs))
.UNSIGNED
, jeśliis_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs))
.FLOAT
lubTOTALORDER
, jeśliis_float(element_type(lhs))
.FLOAT
, jeśliis_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]
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 typcomplex<E>
, gdzieE = 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)]
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:
id = d0 + ... + dk-1 + kd
.d
równa siędimension
, ad0
(...) tod
rozmiar wymiarówinputs
.
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óczdim(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]]
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]]
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)]
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]
, gdziej[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śliresult_dim = output_batch_dimension
.dim(rhs, kernel_output_feature_dimension)
, jeśliresult_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)
.
- (C27)
- 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)
, tois_per_tensor_quantized(result)
. - (C32) Jeśli
is_per_axis_quantized(rhs)
, toquantization_dimension(rhs) = kernel_output_feature_dimension
. - (C33) Jeśli
is_per_axis_quantized(result)
, toquantization_dimension(result) = output_feature_dimension
.
- (C28)
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]]
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]]
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]
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
gdziesize(result_batching_index) = size(lhs_batching_dimensions)
,size(result_lhs_index) = size(lhs_result_dimensions)
isize(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)
.
- (C13)
- 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
.
- (C14)
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]]
// ]
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]
// ]
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śli0 <= 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]
// ]
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]]
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]
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
, gdziecomplex_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
iresult
są różne:- Jeśli
fft_type = FFT
,element_type(operand)
ielement_type(result)
mają ten sam typ zespolony. - Jeśli
fft_type = IFFT
,element_type(operand)
ielement_type(result)
mają ten sam typ zespolony. - Jeśli
fft_type = RFFT
to typ zmiennoprzecinkowy,element_type(operand)
jest typem zmiennoprzecinkowym, aelement_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, aelement_type(result)
jest typem zmiennoprzecinkowym o tej samej semantyce liczby zmiennoprzecinkowej.
- Jeśli
- (C3)
1 <= size(fft_length) <= 3
. - (C4) Jeśli między elementami
operand
iresult
występuje tensorreal
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
.
- Jeśli
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]
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]
, gdziebi
to poszczególne elementy wbatch_index
, a:
jest wstawiony w indeksieindex_vector_dim
, jeśliindex_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ślid_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]
, gdzieoi
to poszczególne elementy woffset_index
, a0
jest wstawiony w indeksach odcollapsed_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 wymiarustart_indices
odpowiadający atrybutowiindex_vector_dim
nie jest uwzględniony.offset_dim_sizes = shape(slice_sizes)
z tym wyjątkiem, że rozmiary podane w poluslice_sizes
odpowiadające wartościomcollapsed_slice_dims
nie są uwzględniane.combine
umieszczabatch_dim_sizes
na osiach odpowiadających wartościombatch_dims
ioffset_dim_sizes
na osiach odpowiadających osiomoffset_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]]
// ]
// ]
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
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]
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
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śliis_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]
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])
lubis_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]]
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]
// ]
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]
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]]
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]
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]]
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'>
, gdzieEi = element_type(inputs[i])
iE' = 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]]
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]]
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]]
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]]
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]
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
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
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śliresult_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]
// ]
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>
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]
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]
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śliis_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]
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śliis_host_transfer = true
,- W przeciwnym razie
DEVICE_TO_DEVICE
.
- (C2)
0 < size(results)
. - (C3)
is_empty(result[:-1])
lubis_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)
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 jestdimensions
.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 wszystkichindex
windex_space(input_slices_converted)
w rosnącej kolejności leksykograficznejindex
.- 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>)
, gdzieis_promotable(element_type(inputs[i]), Ei)
. - (C7)
shape(results...) = shape(inputs...)
z tym wyjątkiem, że rozmiary wymiarówinputs...
odpowiadającedimensions
nie są uwzględniane. - (C8)
element_type(results[i]) = Ei
dla wszystkichi
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]
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 semantykiroundToIntegralTiesToEven
. - Następnie, jeśli wartość
mantissa_bits
jest mniejsza niż liczba bitów mantissy w pierwotnej wartości, bity mantissy są obcinane do wartościmantissa_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]
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ślichannel_id <= 0 and use_global_device_ids = false
.cross_replica_and_partition(replica_groups)
jeślichannel_id > 0 and use_global_device_ids = false
.flattened_ids(replica_groups)
jeślichannel_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 wszystkichsender
wprocess_group
, gdziereceiver_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>)
, gdzieis_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]]
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>)
, gdzieis_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 wszystkichi
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]]
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ąglaniaroundTowardZero
. - 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]
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>
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, żequantization_dimension(operand)
iquantization_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]]
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ślid
wdimensions
.- 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]]
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
, tois_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ślirng_algorithm = THREE_FRY
.2
lub3
, jeślirng_algorithm = PHILOX
.
- zdefiniowana, jeśli
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]
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]
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]]
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]
, gdziesi
to poszczególne elementy wupdate_scatter_index
, a:
jest wstawiony w indeksieindex_vector_dim
, jeśliindex_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ślid_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]
, gdziewi
to poszczególne elementy wupdate_window_index
, a0
jest wstawiony w indeksach odinserted_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 funkcjiindex_space(updates[0])
.exec([update_index, ...], results) = exec([...], updated_results)
, gdzie:- Jeśli
result_index
mieści się w zakresieshape(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 kluczaresults
z wartościąresults...[result_index]
ustawioną naupdated_values...
.- W innym przypadku
updated_results = results
.
- Jeśli
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 wymiaruscatter_indices
odpowiadający wartościindex_vector_dim
nie jest uwzględniony.update_window_dim_sizes <= shape(inputs[0])
z tym wyjątkiem, że rozmiary podane w poluinputs[0]
odpowiadające wartościominserted_window_dims
nie są uwzględniane.combine
umieszczaupdate_scatter_dim_sizes
na osiach odpowiadających wartościomupdate_scatter_dims
iupdate_window_dim_sizes
na osiach odpowiadających wartościomupdate_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>)
, gdzieis_promotable(element_type(inputs[i]), Ei)
. - (C16)
shape(inputs...) = shape(results...)
. - (C17)
element_type(results[i]) = Ei
dla wszystkichi
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]]
// ]
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]]
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
ipadding
, 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)
ireduce_window_without_init
działają dokładnie tak samo jakreduce_window
z tą różnicą, żeschedule
bazowego elementureduce
(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śliselected_values[source_index]
zawiera elementoperand
zoperand_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>
, gdzieE = element_type(operand)
. - (C10)
scatter
ma typ(tensor<E>, tensor<E>) -> tensor<E>
, gdzieis_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]]
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śliis_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
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]
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]
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]
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]
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]]
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]
// ]
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]
, gdzieriN
to poszczególne elementy w obrębieresult_index
, a:
jest wstawiony w miejscuadjusted_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 funkcjacomparator_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
, gdzieR = rank(inputs[0])
. - (C5)
comparator
ma typ(tensor<E1>, tensor<E1>, ..., tensor<EN-1>, tensor<EN-1>) -> tensor<i1>
, gdzieEi = 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]]
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]]
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]]
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]
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, żequantization_dimension(operand)
iquantization_dimension(result)
mogą się różnić. W przeciwnym razie te wartości mogą się różnić.
- (C2)
permutation
to permutacja wartościrange(rank(operand))
. - (C3)
shape(result) = dim(operand, permutation...)
. - (C4) Jeśli
is_per_axis_quantized(result)
, toquantization_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]]
// ]
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 parametrua
w niezmienionej postaci.TRANSPOSE
: wykonaj operację transponowaniaa
.ADJOINT
: wykonaj operację na transponowaniu sprzężonego elementua
.
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)
ashape(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 typtuple<E0, ..., EN-1>
, gdzieEi = 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))
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>
, gdzieTi = type(operand[i])
. - (C2)
body
ma typ(T0, ..., TN-1) -> (T0, ..., TN-1)
, gdzieTi = 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
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ęściTensorElementType
lubQuantizedTensorElementType
odpowiedniego elementuTensorType
lubQuantizedTensorType
.
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 dois_quantized(x) and quantization_dimension(x) is not None
.is_per_tensor_quantized(x: Value | Placeholder | Type) -> Value
to skrót nazwyis_quantized(x) and quantization_dimension(x) is None
.is_promotable(x: Type, y: Type) -> bool
sprawdza, czy typx
może być awansowany na typy
. Gdyx
iy
mają wartośćQuantizedTensorElementType
, promocja jest stosowana tylko dostorage_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 dois_quantized_tensor_element_type(x)
.is_type_name(x: Value | Placeholder | Type) -> Value
. Dostępne dla wszystkich typów. Na przykład funkcjais_float(x)
zwraca wartośćtrue
, jeślix
toFloatType
. Jeślix
jest wartością lub symbolem zastępczym, ta funkcja jest skrótem dois_type_name(type(x))
.max_value(x: Type) -> Value
zwraca maksymalną wartośćTensorElementType
. Jeślix
nie jest wartościąTensorElementType
, zwracaNone
.min_value(x: Type) -> Value
zwraca minimalną możliwą wartośćTensorElementType
. Jeślix
nie jest wartościąTensorElementType
, zwracaNone
.member_name(x: Value | Placeholder | Type) -> Any
. Wartość dostępna w przypadku wszystkich definicji typumember_name
każdego typu. Na przykładtensor_element_type(x)
zwraca częśćTensorElementType
odpowiedniejTensorType
. Jeślix
jest wartością lub symbolem zastępczym, ta funkcja jest skrótem domember_name(type(x))
. Jeślix
nie jest typem, który ma odpowiedni element, wartość lub obiekt zastępczy takiego typu, zwracaNone
.
Konstrukcja wartości
operation_name(*xs: Value | Type) -> Value
. Dostępne dla wszystkich operacji. Na przykładadd(lhs, rhs)
przyjmuje 2 wartości tensorówlhs
irhs
i zwraca dane wyjściowe oceny operacjiadd
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 danychtype(x)
idestination_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 zwracatrue
, jeśli wszystkie elementy funkcjix
mają wartośćNaN
lubfalse
. Jeślix
nie jest tensorem, zwracaNone
.Element
is_sorted(x: Value) -> Value
jest zdefiniowany w tensorach i zwracatrue
, jeśli elementy właściwościx
są posortowane rosnąco według kolejności leksykograficznej ich indeksów lub w przeciwnym razie zwraca wartośćfalse
. Jeślix
nie jest tensorem, zwracaNone
.is_unique(x: Value) -> Value
jest zdefiniowany w tensorach i zwracatrue
, jeślix
nie ma zduplikowanych elementów. W przeciwnym razie funkcjafalse
nie jest określona. Jeślix
nie jest tensorem, zwracaNone
.member_name(x: Value) -> Any
jest zdefiniowany dla wszystkich definicji elementówmember_name
wszystkich wartości. Na przykładreal_part(x)
zwraca częśćRealPart
odpowiedniejComplexConstant
. Jeślix
nie jest wartością, która ma odpowiedniego użytkownika, zwracaNone
.Funkcja
same(x: Value) -> Value
jest definiowana w tensorach i zwracatrue
, jeśli wszystkie elementy funkcjix
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 zwracatrue
. Jeślix
nie jest tensorem, zwracaNone
.Funkcja
split(x: Value, num_results: Value, axis: Value) -> Value
jest definiowana na tensorach i zwraca wycinki (num_results
) wartościx
wzdłuż osiaxis
. Jeślix
nie jest tensorem anidim(x, axis) % num_results != 0
, zwracaNone
.
Obliczanie kształtów
axes(x: Value | Placeholder | Type) -> Value
to skrót dorange(rank(x))
.dim(x: Value | Placeholder | Type, axis: Value) -> Value
to skrót doshape(x)[axis]
.dims(x: Value | Placeholder | Type, axes: List) -> List
to skrót dolist(map(lambda axis: dim(x, axis), axes))
.index_space(x: Value | Placeholder | Type) -> Value
jest zdefiniowany dla tensorów i zwraca indeksysize(x)
dla odpowiedniego elementuTensorType
posortowane w kolejności leksykograficznej, tj.[0, ..., 0]
,[0, ..., 1]
, ...,shape(x) - 1
. Jeślix
nie jest typem tensora, skwantyzowanym typem tensora, wartością lub zmienną jednego z tych typów, zwracaNone
.rank(x: Value | Placeholder | Type) -> Value
to skrót dosize(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 doreduce(lambda x, y: x * y, shape(x))
.
Obliczenia kwantyzacji
def baseline_element_type(x: Value | Placeholder | Type) -> Type
to skrót nazwyelement_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, czylishape
,storage_type
,expressed_type
,storage_min
,storage_max
iquantization_dimension
(w przypadku typu skwantyzowanego na osi), ale wartościscales
izero 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.