Specifica StableHLO

StableHLO è un insieme di operazioni per le operazioni di alto livello (HLO) nei modelli di machine learning (ML). StableHLO funziona come livello di portabilità tra diversi framework ML e compilatori ML: i framework ML che producono programmi StableHLO sono compatibili con i compilatori ML che utilizzano programmi StableHLO.

Il nostro obiettivo è semplificare e accelerare lo sviluppo di ML creando una maggiore interoperabilità tra vari framework di ML (come TensorFlow, JAX e PyTorch) e compilatori di ML (come XLA e IREE). A questo scopo, il documento fornisce una specifica per il linguaggio di programmazione StableHLO.

Questa specifica contiene tre sezioni principali. Innanzitutto, la sezione Programmi descrive la struttura dei programmi StableHLO, composti da funzioni StableHLO, a loro volta costituite da operazioni StableHLO. All'interno di questa struttura, la sezione Op specifica la semantica delle singole operazioni. La sezione Esecuzione fornisce la semantica per tutte queste operazioni che vengono eseguite insieme all'interno di un programma. Infine, la sezione Notazione illustra la notazione utilizzata nella specifica.

Per visualizzare le specifiche di una release precedente di StableHLO, apri il repository nella release con tag di tuo interesse. Ad esempio, la specifica StableHLO v0.19.0. Per visualizzare le modifiche apportate a ogni aggiornamento minore di StableHLO, consulta il log della versione in VhloDialect.td.

Programmi

Program ::= {Func}

I programmi StableHLO sono costituiti da un numero arbitrario di funzioni StableHLO. Di seguito è riportato un programma di esempio con una funzione @main che ha 3 input (%image, %weights e %bias) e 1 output. Il corpo della funzione ha 6 operazioni.

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

Funzioni

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

Le funzioni StableHLO (chiamate anche funzioni con nome) hanno un identificatore, input/output e un corpo. In futuro, prevediamo di introdurre metadati aggiuntivi per le funzioni per ottenere una migliore compatibilità con HLO (#425, #626, #740, #744).

Identificatori

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

Gli identificatori StableHLO sono simili a quelli di molti linguaggi di programmazione, con due peculiarità: 1) tutti gli identificatori hanno sigilli che distinguono i diversi tipi di identificatori, 2) gli identificatori di valore possono essere completamente numerici per semplificare la generazione di programmi StableHLO.

Tipi

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

I tipi StableHLO sono classificati in tipi di valore (chiamati anche tipi di prima classe) che rappresentano i valori StableHLO e tipi non di valore che descrivono altri elementi del programma. I tipi di HLO stabili sono simili a quelli in molti linguaggi di programmazione, ma la loro peculiarità principale è rappresentata dalla natura specifica del dominio di StableHLO, che genera risultati insoliti (ad esempio, i tipi scalari non sono tipi di valore).

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

I tipi di tensori rappresentano i tensori, ovvero gli array multidimensionali. Hanno una forma e un tipo di elemento, dove una forma rappresenta dimensioni non negative o sconosciute nell'ordine crescente delle dimensioni corrispondenti (chiamate anche assi) numerate da 0 a R-1. Il numero di dimensioni R è chiamato rango. Ad esempio, tensor<2x3xf32> è un tipo di tensore con forma 2x3 e tipo di elemento f32. Ha due dimensioni (o, in altre parole, due assi), la dimensione 0 e la dimensione 1, le cui dimensioni sono 2 e 3. Il suo ranking è 2.

Le forme possono essere parzialmente o completamente sconosciute (dinamiche), ad esempio tensor<?x2xf64> è parzialmente sconosciuta e tensor<?x?xf64> è completamente sconosciuta. Le dimensioni dinamiche sono rappresentate utilizzando un ?. Le forme non possono essere non classificate.

In futuro, prevediamo di estendere i tipi di tensori oltre le dimensioni e i tipi di elementi, ad esempio per includere i layout (#629) e la sparsità (#1078).

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

I tipi di elementi quantizzati rappresentano i valori interi di un tipo di archiviazione nell'intervallo compreso tra storage_min e storage_max (inclusi) che corrispondono ai valori in virgola mobile di un tipo espresso. Per un determinato valore intero i, il corrispondente valore in virgola mobile f può essere calcolato come f = (i - zero_point) * scale, dove scale e zero_point sono chiamati parametri di quantizzazione. storage_min e storage_max sono facoltativi nella grammatica, ma hanno rispettivamente i valori predefiniti min_value(storage_type) e max_value(storage_type). I tipi di elementi quantizzati hanno i seguenti vincoli:

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

Al momento, QuantizationScale è una costante a virgola mobile, ma c'è un forte interesse per le scale basate su numeri interi, rappresentate con moltiplicatori e spostamenti. Abbiamo in programma di esplorare questa opzione nel prossimo futuro (#1404).

È in corso una discussione sulla semantica di QuantizationZeroPoint, incluso il tipo, i valori e se possono esserci uno o più punti zero in un tipo di tensore quantizzato. In base ai risultati di questa discussione, la specifica relativa ai punti zero potrebbe cambiare in futuro (#1405).

Un'altra discussione in corso riguarda la semantica di QuantizationStorageMin e QuantizationStorageMax per determinare se debbano essere imposti vincoli su questi valori e sui valori dei tensori quantizzati (#1406).

Infine, prevediamo di esplorare la rappresentazione di scale e punti zero sconosciuti, in modo simile a come prevediamo di esplorare la rappresentazione delle dimensioni di dimensioni sconosciute (#1407).

I tipi di tensori quantizzati rappresentano i tensori con elementi quantizzati. Questi tensori sono esattamente come i tensori normali, ad eccezione del fatto che i loro elementi hanno tipi di elementi quantizzati anziché tipi di elementi normali.

Nei tensori quantizzati, la quantizzazione può essere per tensore, ovvero avere un scale e zero_point per l'intero tensore o essere per asse, il che significa avere più scales e zero_points, una coppia per sezione di una particolare dimensione quantization_dimension. In modo più formale, in un tensore t con quantizzazione per asse, ci sono dim(t, quantization_dimension) sezioni del quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :], ecc. Tutti gli elementi del iesimo slice utilizzano scales[i] e zero_points[i] come i loro parametri di quantizzazione. I tipi di tensori quantizzati hanno i seguenti vincoli:

• Per la quantizzazione per tensore:
  • Nessun vincolo aggiuntivo.
• Per la quantizzazione per asse:
  • (C12) quantization_dimension < rank(self).
  • (C13) dim(self, quantization_dimension) = size(scales).

TokenType ::= 'token'

I tipi di token rappresentano i token, ovvero valori opachi prodotti e consumati da alcune operazioni. I token vengono utilizzati per imporre l'ordine di esecuzione sulle operazioni, come descritto nella sezione Esecuzione.

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

I tipi di tuple rappresentano tuple, ovvero elenchi eterogenei. Le tuple sono una funzionalità legacy che esiste solo per la compatibilità con l'HLO. In HLO, le tuple vengono utilizzate per rappresentare input e output con parametri variabili. In StableHLO, gli input e gli output variabili sono supportati in modo nativo e l'unico utilizzo delle tuple in StableHLO è rappresentare in modo completo l'ABI HLO, dove ad esempio T, tuple<T> e tuple<tuple<T>> possono essere sostanzialmente diversi a seconda di una determinata implementazione. In futuro, prevediamo di apportare modifiche all'ABI HLO che potrebbero consentirci di rimuovere i tipi di tuple da StableHLO (#598).

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

I tipi di elementi rappresentano gli elementi dei tipi di tensori. A differenza di molti linguaggi di programmazione, questi tipi non sono di prima classe in StableHLO. Ciò significa che i programmi SttableHLO non possono rappresentare direttamente i valori di questi tipi. Di conseguenza, i valori scalari di tipo T sono idiomatici con valori tensoriali 0-dimensionali di tipo tensor<T>.

• Il tipo booleano rappresenta i valori booleani true e false.
• I tipi di numeri interi possono essere con segno (si) o senza segno (ui) e avere una delle larghezze di bit supportate (2, 4, 8, 16, 32 o 64). I tipi siN con segno rappresentano valori interi da -2^(N-1) a 2^(N-1)-1 inclusi, mentre i tipi uiN senza segno rappresentano valori interi da 0 a 2^N-1 inclusi.
• I tipi a virgola mobile possono essere uno dei seguenti:
  • f8E3M4, f8E4M3 e f8E5M2 numeri con virgola mobile a 8 bit conformi alle convenzioni IEEE-754.
  • Tipi f8E4M3FN e f8E5M2 corrispondenti rispettivamente alle codifiche E4M3 e E5M2 del formato FP8 descritto in Formati FP8 per il deep learning.
  • Tipi f8E4M3FNUZ e f8E5M2FNUZ corrispondenti alle codifiche E4M3 e E5M2 dei formati FP8 descritti in Formati numerici a 8 bit per reti neurali profonde.
  • Tipo f8E4M3B11FNUZ corrispondente alla codifica E4M3 dei formati FP8 descritti in Addestramento in virgola mobile ibrido a 8 bit (HFP8) e inferenza per reti neurali profonde.
  • Tipo bf16 corrispondente al formato bfloat16 descritto in BFloat16: il segreto delle alte prestazioni su Cloud TPU.
  • Tipi f16, f32 e f64 corrispondenti ai formati rispettivamente binary16 ("metà precisione"), binary32 ("precisione singola") e binary64 ("precisione doppia") descritti nello standard IEEE 754.
  • Il tipo tf32 corrisponde al formatoTensorFloat32 e ha un supporto limitato in StableHLO.
  • Tipi MX (microscalabilità) f4E2M1FN, f6E2M3FN, f6E3M2FN e f8E8M0FNU descritti nella Specifica dei formati di microscalabilità di OCP.
• I tipi complessi rappresentano valori complessi che hanno una parte reale e una parte immaginaria dello stesso tipo di elemento. I tipi complessi supportati sono complex<f32> (entrambe le parti sono di tipo f32) e complex<f64> (entrambe le parti sono di tipo f64).

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

I tipi di funzioni rappresentano sia le funzioni con nome che quelle anonime. Hanno tipi di input (l'elenco di tipi a sinistra di ->) e tipi di output (l'elenco di tipi a destra di ->). In molti linguaggi di programmazione, i tipi di funzione sono di prima classe, ma non in StableHLO.

StringType ::= 'string'

Il tipo di stringa rappresenta sequenze di byte. A differenza di molti linguaggi di programmazione, il tipo di stringa non è di prima classe in StableHLO e viene utilizzato solo per specificare i metadati statici per gli elementi del programma.

Operazioni

Le operazioni StableHLO (chiamate anche ops) rappresentano un insieme chiuso di operazioni di alto livello nei modelli di machine learning. Come discusso sopra, la sintassi di StableHLO è fortemente ispirata a MLIR, che non è necessariamente l'alternativa più ergonomica, ma è probabilmente la più adatta per lo scopo di StableHLO di creare una maggiore interoperabilità tra i framework ML e i compilatori ML.

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

Le operazioni SttableHLO (chiamate anche ops) hanno un nome, input/output e una firma. Il nome è costituito dal prefisso stablehlo. e da un mnemonico che identifica in modo univoco una delle operazioni supportate. Di seguito è riportato un elenco completo di tutte le operazioni supportate.

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

Le operazioni consumano input e producono output. Gli input sono classificati in valori di input (calcolati durante l'esecuzione), funzioni di input (fornite in modo statico, perché in StableHLO le funzioni non sono valori di prima classe) e attributi di input (forniti anche in modo statico). Il tipo di input e output consumati e prodotti da un'operazione dipende dal relativo mnemonico. Ad esempio, l'operatore add utilizza 2 valori di input e produce 1 valore di output. In confronto, l'operazione select_and_scatter richiede 3 valori di input, 2 funzioni di input e 3 attributi di input.

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

Le funzioni di input (chiamate anche funzioni anonime) sono molto simili alle funzioni con nome, tranne per il fatto che: 1) non hanno un identificatore (da qui il nome "anonimo"), 2) non dichiarano i tipi di output (i tipi di output sono dedotti dall'operazione return all'interno della funzione).

La sintassi per le funzioni di input include una parte attualmente non utilizzata (vedi la produzione Unused sopra) che è presente per la compatibilità con MLIR. In MLIR, esiste un concetto più generale di "regioni" che possono avere più "blocchi" di operazioni collegate tra loro tramite operazioni di salto. Questi blocchi hanno ID che corrispondono all'ambiente di produzione Unused, in modo che possano essere distinti tra loro. StableHLO non ha operazioni di salto, quindi la parte corrispondente della sintassi MLIR non viene utilizzata (ma è ancora presente).

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

Gli attributi di input hanno un nome e un valore che è una delle costanti supportate. Sono il modo principale per specificare metadati statici per gli elementi del programma. Ad esempio, l'operatore concatenate utilizza l'attributo dimension per specificare la dimensione lungo la quale vengono concatenati i relativi valori di input. Allo stesso modo, l'operazione slice utilizza più attributi, come start_indices e limit_indices, per specificare i limiti utilizzati per suddividere il valore di input.

Al momento, i programmi StableHLO in uso a volte contengono attributi che non sono descritti in questo documento. In futuro, prevediamo di integrare questi attributi nell'insieme di operazioni StableHLO o di vietarne la visualizzazione nei programmi StableHLO. Nel frattempo, ecco l'elenco di questi attributi:

• layout (#629).
• mhlo.frontend_attributes (#628).
• mhlo.sharding (#619).
• output_operand_aliases (#740).
• Metadati sulla località (#594).

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

La firma dell'operazione è costituita dai tipi di tutti i valori di input (l'elenco dei tipi sul lato sinistro di ->) e dai tipi di tutti i valori di output (l'elenco dei tipi sul lato destro di ->). Rigorosamente parlando, i tipi di input sono ridondanti e anche i tipi di output sono quasi sempre ridondanti (perché per la maggior parte delle operazioni StableHLO, i tipi di output possono essere dedotti dagli input). Tuttavia, la firma op fa deliberatamente parte della sintassi di StableHLO per la compatibilità con MLIR.

Di seguito è riportato un esempio di operazione il cui mnemonico è select_and_scatter. Utilizza 3 valori di input (%operand, %source e %init_value), 2 funzioni di input e 3 attributi di input (window_dimensions, window_strides e padding). Nota come la firma dell'operazione includa solo i tipi dei relativi valori di input (ma non i tipi di funzioni e attributi di input forniti in linea).

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

Costanti

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

Le costanti StableHLO hanno un valore letterale e un tipo che insieme rappresentano un valore StableHLO. In genere, il tipo fa parte della sintassi della costante, tranne quando è inequivocabile (ad es. una costante booleana ha inequivocabile il tipo i1, mentre una costante intera può avere più tipi possibili).

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

Le costanti booleane rappresentano i valori booleani true e false. Le costanti booleane hanno tipo i1.

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

Le costanti intere rappresentano valori interi tramite stringhe che utilizzano la notazione decimale o esadecimale. Altre basi, ad esempio binarie o ottali, non sono supportate. Le costanti intere hanno i seguenti vincoli:

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

Le costanti con virgola mobile rappresentano valori con virgola mobile tramite stringhe che utilizzano la notazione decimale o scientifica. Inoltre, la notazione esadecimale può essere utilizzata per specificare direttamente i bit sottostanti nel formato a virgola mobile del tipo corrispondente. Le costanti in virgola mobile presentano i seguenti vincoli:

• (C1) Se viene utilizzata la notazione non esadecimale, is_wellformed(float_literal, float_type).
• (C2) Se viene utilizzata la notazione esadecimale, size(hexadecimal_digits) = num_bits(float_type) / 4.

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

Le costanti complesse rappresentano valori complessi utilizzando elenchi di una parte reale (viene visualizzata per prima) e una parte immaginaria (viene visualizzata per seconda). Ad esempio, (1.0, 0.0) : complex<f32> rappresenta 1.0 + 0.0i e (0.0, 1.0) : complex<f32> rappresenta 0.0 + 1.0i. L'ordine in cui queste parti vengono archiviate in memoria è definito dall'implementazione. Le costanti complesse hanno i seguenti vincoli:

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

Le costanti tensore rappresentano i valori dei tensori utilizzando elenchi nidificati specificati tramite la notazione NumPy. Ad esempio, dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> rappresenta un valore tensoriale con la seguente mappatura dagli indici agli elementi:{0, 0} => 1, {0, 1} => 2, {0, 2} => 3, {1, 0} => 4, {1, 1} => 5,{1, 2} => 6. L'ordine in cui questi elementi vengono archiviati in memoria è definito dall'implementazione. Le costanti tensore hanno i seguenti vincoli:

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

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

Le costanti di tensore quantizzati rappresentano i valori dei tensori quantizzati utilizzando la stessa notazione delle costanti di tensore, con elementi specificati come costanti del loro tipo di archiviazione. Le costanti tensore quantizzate hanno i seguenti vincoli:

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

Le lettere di stringa sono costituite da byte specificati utilizzando caratteri ASCII e sequenze di escape. Sono indipendenti dalla codifica, pertanto l'interpretazione di questi byte è definita dall'implementazione. I valori letterali stringa sono di tipo string.

Operazioni

abs

Semantica

Esegue l'operazione di valore assoluto elemento per elemento sul tensore operand e produce un result. A seconda del tipo di elemento, esegue le seguenti operazioni:

• Per numeri interi con segno: modulo intero.
• Per i valori float: abs da IEEE-754.
• Per i numeri complessi: modulo complesso.
• Per i tipi quantizzati: dequantize_op_quantize(abs, operand, type(result)).

Input

Output

Vincoli

• (C1) shape(result) = shape(operand).
• (C2) baseline_element_type(result) è definito come:
  • complex_element_type(element_type(operand)) se is_complex(operand).
  • baseline_element_type(operand) in caso contrario.

Esempi

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

Altri esempi

add

Semantica

Esegue l'addizione elemento per elemento di due tensori lhs e rhs e produce un result. A seconda del tipo di elemento, esegue le seguenti operazioni:

• Per i valori booleani: OR logico.
• Per gli interi: addizione di numeri interi.
• Per i valori float: addition da IEEE-754.
• Per i numeri complessi: addizione complessa.
• Per i tipi quantizzati: dequantize_op_quantize(add, lhs, rhs, type(result)).

Input

Output

Vincoli

• Se l'operazione utilizza tensori non quantizzati:
  • (C1) type(lhs) = type(rhs) = type(result).
• Se l'operazione utilizza tensori quantizzati:
  • (C2) is_quantized(lhs) and is_quantized(rhs) and is_quantized(result).
  • (C3) storage_type(lhs) = storage_type(rhs) = storage_type(result).
  • (C4) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C5) (is_per_axis_quantized(lhs) or is_per_axis_quantized(rhs)) = is_per_axis_quantized(result).
  • (C6) Se is_per_axis_quantized(lhs), allora quantization_dimension(lhs) = quantization_dimension(result).
  • (C7) Se is_per_axis_quantized(rhs), allora quantization_dimension(rhs) = quantization_dimension(result).

Esempi

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

Altri esempi

after_all

Semantica

Garantisce che le operazioni che producono inputs vengano eseguite prima di qualsiasi operazione che dipende da result. L'esecuzione di questa operazione non fa nulla, esiste solo per stabilire le dipendenze dei dati da result a inputs.

Input

Output

Esempi

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

Altri esempi

all_gather

Semantica

All'interno di ogni gruppo di processi nella griglia di processi StableHLO, concatena i valori dei tensori operands di ogni processo lungo all_gather_dim e produce i tensori results.

L'operazione suddivide la griglia del processo StableHLO in process_groups, che è definita come segue:

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

Successivamente, all'interno di ogni process_group:

• operands...@receiver = [operand@sender for sender in process_group] per tutti receiver in process_group.
• results...@process = concatenate(operands...@process, all_gather_dim) per tutti process in process_group.

Input

Output

Vincoli

• (C1) 0 <= all_gather_dim < rank(operands...).
• (C2) is_unique(replica_groups).
• (C3) size(replica_groups) è definito come:
  • num_replicas se si utilizza cross_replica.
  • num_replicas se viene utilizzato cross_replica_and_partition.
  • num_processes se viene utilizzato flattened_ids.
• (C4) 0 <= replica_groups < size(replica_groups).
• (C5) Se use_global_device_ids = true, allora channel_id > 0.
• (C6) type(results...) = type(operands...) eccetto:
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1).

Esempi

// num_replicas: 2
// num_partitions: 1
// %operand0@(0, 0): [[1, 2], [3, 4]]
// %operand0@(1, 0): [[5, 6], [7, 8]]
// %operand1@(0, 0): [[11, 12], [13, 14]]
// %operand1@(1, 0): [[15, 16], [17, 18]]
%result:2 = "stablehlo.all_gather"(%operand0, %operand1) {
  all_gather_dim = 1 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  // channel_id = 0
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
  // use_global_device_ids = false
} : (tensor<2x2xi64>, tensor<2x2xi64>) -> (tensor<2x4xi64>, tensor<2x4xi64>)
// %result0@(0, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result0@(1, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result1@(0, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]
// %result1@(1, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]

 Altri esempi

all_reduce

Semantica

All'interno di ogni gruppo di processi nella griglia di processi StableHLO, applica una funzione di riduzione computation ai valori dei tensori operands di ogni processo e produce tensori results.

L'operazione suddivide la griglia del processo StableHLO in process_groups, che è definita come segue:

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

In seguito, all'interno di ogni process_group:

• results...@process[result_index] = exec(schedule) per un albero binario schedule dove:
  • exec(node) = computation(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value.
• schedule è un albero binario definito dall'implementazione il cui attraversamento dell'ordine è to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0])).

Input

Output

Vincoli

• (C1) is_unique(replica_groups).
• (C2) Per size(replica_groups) si intende:
  • num_replicas se viene utilizzato cross_replica.
  • num_replicas se viene utilizzato cross_replica_and_partition.
  • num_processes se viene utilizzato flattened_ids.
• (C3) 0 <= replica_groups < size(replica_groups).
• (C4) Se use_global_device_ids = true, allora channel_id > 0.
• (C5) computation è di tipo (tensor<E>, tensor<E>) -> (tensor<E>) dove is_promotable(element_type(operand), E).
• (C6) shape(results...) = shape(operands...).
• (C7) element_type(results...) = E.

Esempi

// num_replicas: 2
// num_partitions: 1
// %operand0@(0, 0): [1, 2, 3, 4]
// %operand0@(1, 0): [5, 6, 7, 8]
// %operand1@(0, 0): [9, 10, 11, 12]
// %operand1@(1, 0): [13, 14, 15, 16]
%result:2 = "stablehlo.all_reduce"(%operand0, %operand0) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  // channel_id = 0
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
  // use_global_device_ids = false
} : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>, tensor<4xi64>)
// %result0@(0, 0): [6, 8, 10, 12]
// %result0@(1, 0): [6, 8, 10, 12]
// %result1@(0, 0): [22, 24, 26, 28]
// %result1@(1, 0): [22, 24, 26, 28]

Altri esempi

all_to_all

Semantica

All'interno di ogni gruppo di processi nella griglia di processi StableHLO, suddivide i valori dei tensori operands lungo split_dimension in parti, distribuisce le parti suddivise tra i processi, concatena le parti sparse lungo concat_dimension e produce tensori results. L'operazione divide la griglia di processo StableHLO in process_groups, che è definita come segue:

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

In seguito, all'interno di ogni process_group:

• split_parts...@sender = split(operands...@sender, split_count, split_dimension) per tutti i sender in process_group.
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group] dove receiver_index = process_group.index(receiver).
• results...@process = concatenate(scattered_parts...@process, concat_dimension).

Input

Output

Vincoli

• (C1) 0 <= split_dimension < rank(operands...).
• (C2) dim(operands..., split_dimension) % split_count = 0.
• (C3) 0 <= concat_dimension < rank(operands...).
• (C4) 0 < split_count.
• (C5) is_unique(replica_groups).
• (C6) size(replica_groups) è definito come:
  • num_replicas se si utilizza cross_replica.
  • num_partitions se viene utilizzato cross_partition.
• (C7) 0 <= replica_groups < size(replica_groups).
• (C8) dim(replica_groups, 1) = split_count.
• (C9) type(results...) = type(operands...) tranne se split_dimension != concat_dimension:
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count.
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count.

Esempi

// num_replicas: 2
// num_partitions: 1
// %operand1@(0, 0): [[1, 2, 3, 4],
//                    [5, 6, 7, 8]]
// %operand1@(1, 0): [[9, 10, 11, 12],
//                    [13, 14, 15, 16]]
// %operand2@(0, 0): [[17, 18, 19, 20],
//                    [21, 22, 23, 24]]
// %operand2@(1, 0): [[25, 26, 27, 28],
//                    [29, 30, 31, 32]]
%result:2 = "stablehlo.all_to_all"(%operand1, %operand2) {
  split_dimension = 1 : i64,
  concat_dimension = 0 : i64,
  split_count = 2 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
  // channel_id = 0
} : (tensor<2x4xi64>, tensor<2x4xi64>) -> (tensor<4x2xi64>, tensor<4x2xi64>)
// %result#0@(0, 0): [[1, 2], [5, 6], [9, 10], [13, 14]]
// %result#0@(1, 0): [[3, 4], [7, 8], [11, 12], [15, 16]]
// %result#1@(0, 0): [[17, 18], [21, 22], [25, 26], [29, 30]]
// %result#1@(1, 0): [[19, 20], [23, 24], [27, 28], [31, 32]]

 Altri esempi

e

Semantica

Esegue l'operatore AND a livello di elemento di due tensori lhs e rhs e produce un tensore result. A seconda del tipo di elemento, esegue le seguenti operazioni:

• Per i valori booleani: AND logico.
• Per i numeri interi: AND a livello di bit.

Input

Output

Vincoli

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

Esempi

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

 Altri esempi

atan2

Semantica

Esegue l'operazione atan2 elemento per elemento sui tensori lhs e rhs e produce un result tensore. A seconda del tipo di elemento, esegue le seguenti operazioni:

• Per i numeri in virgola mobile: atan2 dallo standard IEEE-754.
• Per i numeri complessi: atan2 complesso.
• Per i tipi quantizzati: dequantize_op_quantize(atan2, lhs, rhs, type(result)).

Input

Output

Vincoli

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

Esempi

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

 Altri esempi

batch_norm_grad

Semantica

Calcola i gradienti di diversi input di retropropagazione di batch_norm_training da grad_output e produce tensori grad_operand, grad_scale e grad_offset. Più formalmente, questa operazione può essere espressa come una decomposizione nelle operazioni StableHLO esistenti utilizzando la sintassi Python come segue:

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

Per i tipi quantizzati, esegue 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)).

Input

Output

Vincoli

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, mean, variance, grad_output, grad_operand, grad_scale e grad_offset hanno lo stesso baseline_element_type.
• (C3) operand, grad_output e grad_operand hanno la stessa forma.
• (C4) scale, mean, variance, grad_scale e grad_offset hanno la stessa forma.
• (C5) size(scale) = dim(operand, feature_index).

Esempi

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

Semantica

Normalizza il tensore operand in tutte le dimensioni tranne la dimensione feature_index e produce un tensore result. Più formalmente, questa operazione può essere espressa come una decomposizione nelle operazioni StableHLO esistenti utilizzando la sintassi Python come segue:

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)

Per i tipi quantizzati, esegue 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)).

Input

Output

Vincoli

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, mean, variance e result hanno lo stesso 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).

Esempi

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

Semantica

Calcola la media e la varianza in tutte le dimensioni tranne la dimensione feature_index e normalizza il tensore operand producendo i tensori output, batch_mean e batch_var. In modo più formale, questa operazione può essere espressa come una decomposizione delle operazioni StableHLO esistenti utilizzando la sintassi Python come segue:

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

Per i tipi quantizzati, esegue 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)).

Input

Output

Vincoli

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, batch_mean, batch_var e output hanno lo stesso 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).

Esempi

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

Semantica

Esegue un'operazione di bitcast sul tensore operand e produce un tensore result dove i bit dell'intero tensore operand vengono reinterpretati utilizzando il tipo del tensore result.

In modo più formale, dati E = element_type(operand), E' = element_type(result) e R = rank(operand):

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

bits restituisce la rappresentazione in memoria di un determinato valore e il suo comportamento è definito dall'implementazione perché la rappresentazione esatta dei tensori è definita dall'implementazione e anche la rappresentazione esatta dei tipi di elementi è definita dall'implementazione.

Input

Output

Vincoli

• (C1) Dati E = is_quantized(operand) ? storage_type(operand) : element_type(operand), E' = is_quantized(result) ? storage_type(result) : element_type(result) e R = rank(operand):
  • Se num_bits(E') = num_bits(E), shape(result) = shape(operand).
  • Se num_bits(E') < num_bits(E):
  • rank(result) = R + 1.
  • dim(result, i) = dim(operand, i) per tutti i 0 <= i < R.
  • dim(result, R) * num_bits(E') = num_bits(E).
  • Se num_bits(E') > num_bits(E):
  • rank(result) = R - 1.
  • dim(result, i) = dim(operand, i) per tutti i 0 <= i < R.
  • dim(operand, R - 1) * num_bits(E) = num_bits(E').
• (C2) Se is_complex(operand) or is_complex(result), allora is_complex(operand) and is_complex(result).

Esempi

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

 Altri esempi

broadcast_in_dim

Semantica

Espande le dimensioni e/o il ranking di un tensore di input duplicando i dati nel tensore operand e produce un tensore result. Più formalmente, result[result_index] = operand[operand_index] dove per tutti i d in axes(operand):

• operand_index[d] = 0 se dim(operand, d) = 1.
• operand_index[d] = result_index[broadcast_dimensions[d]] in caso contrario.

Input

Output

Vincoli

• (C1) element_type(result) è dato da:
  • element_type(operand), se !is_per_axis_quantized(operand).
  • element_type(operand) ad eccezione dei seguenti casi: quantization_dimension(operand), scales(operand) e zero_points(operand) possono differire da quantization_dimension(result), scales(result) e zero_points(result) risp., altrimenti.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Per tutti i valori d in axes(operand):
  • dim(operand, d) = 1 o
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Se is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Se dim(operand, quantization_dimension(operand)) = 1, then scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))).

Esempi

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

Altri esempi

richiesta

Semantica

Genera l'output dall'esecuzione esattamente di una funzione da branches a seconda del valore di index. Più formalmente, result = selected_branch() dove:

• selected_branch = branches[index] se 0 <= index < size(branches).
• selected_branch = branches[-1] in caso contrario.

Input

Output

Vincoli

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

Esempi

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

 Altri esempi

Cbrt

Semantica

Esegue l'operazione di radice cubica elemento per elemento sul tensore operand e produce un result tensore. A seconda del tipo di elemento, esegue le seguenti operazioni:

• Per i valori float: rootn(x, 3) da IEEE-754.
• Per i numeri complessi: radice cubica complessa.
• Per i tipi quantizzati: dequantize_op_quantize(cbrt, operand, type(result))

Input

Output

Vincoli

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

Esempi

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

Altri esempi

ceil

Semantica

Esegue il tetto elemento per elemento del tensore operand e produce un tensore result. Implementa l'operazione roundToIntegralTowardPositive della specifica IEEE-754. Per i tipi quantizzati, esegue dequantize_op_quantize(ceil, operand, type(result)).

Input

Output

Vincoli

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

Esempi

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

Altri esempi

cholesky

Semantica

Calcola la decomposizione di Cholesky di un batch di matrici.

Più formalmente, per tutti i valori i in index_space(result), result[i0, ..., iR-3, :, :] è una decomposizione di Cholesky di a[i0, ..., iR-3, :, :], sotto forma di matrice triangolare inferiore (se lower è true) o triangolare superiore (se lower è false). I valori di output nel triangolo opposto, ovvero il triangolo superiore rigido o il triangolo inferiore stretto corrispondentemente, sono definiti dall'implementazione.

Se esiste i in cui la matrice di input non è una matrice hermitiana positiva definita, il comportamento non è definito.

Per i tipi quantizzati, esegue dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)).

Input

Output

Vincoli

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

Esempi

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

clampare

Semantica

Fissa ogni elemento del tensore operand tra un valore minimo e massimo e produce un tensore result. In modo più formale, result[result_index] = minimum(maximum(operand[result_index], min_element), max_element), dove min_element = rank(min) = 0 ? min[] : min[result_index], max_element = rank(max) = 0 ? max[] : max[result_index]. Per i tipi quantizzati, esegue dequantize_op_quantize(clamp, min, operand, max, type(result)).

Imporre un ordinamento su numeri complessi comporta una semantica sorprendente, quindi in futuro abbiamo in programma di rimuovere il supporto dei numeri complessi per questa operazione (#560).

Input

Output

Vincoli

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

Esempi

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

Altri esempi

collective_broadcast

Semantica

All'interno di ogni gruppo di processi nella griglia dei processi StableHLO, invia il valore del tensore operand dal processo di origine ai processi di destinazione e produce un tensore result.

L'operazione divide la griglia di processo StableHLO in process_groups, che è definita come segue:

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

Successivamente, result@process è dato da:

• operand@process_groups[i, 0] se esiste un i tale che il processo sia in process_groups[i].
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) in caso contrario.

Input

Output

Vincoli

• (C1) is_unique(replica_groups).
• (C2) 0 <= replica_groups < N dove N è definito come:
  • num_replicas se viene utilizzato cross_replica.
  • num_partitions se viene utilizzato cross_partition.
• (C3) type(result) = type(operand).

Esempi

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

Semantica

All'interno di ogni gruppo di processi nella griglia di processi StableHLO, invia il valore del operand tensore dal processo di origine al processo di destinazione e produce un operand tensore.result

L'operazione suddivide la griglia del processo StableHLO in process_groups, che è definita come segue:

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

Successivamente, result@process è dato da:

• operand@process_groups[i, 0], se esiste un i tale che process_groups[i, 1] = process.
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) in caso contrario.

Input

Output

Vincoli

• (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, dove N è definito come:
  • num_replicas se si utilizza cross_replica.
  • num_partitions se viene utilizzato cross_partition.
• (C5) type(result) = type(operand).

Esempi

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

 Altri esempi

confronta

Semantica

Esegue il confronto elemento per elemento dei tensori lhs e rhs in base a comparison_direction e compare_type e produce un tensore result.

I valori di comparison_direction e compare_type hanno la seguente semantica:

Per i tipi di elementi booleani e interi:

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

Per i tipi di elementi in virgola mobile con compare_type = FLOAT, l'operazione implementa le seguenti operazioni IEEE-754:

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

Per i tipi di elementi in virgola mobile con compare_type = TOTALORDER, l'operatore utilizza la combinazione di operazioni totalOrder e compareQuietEqual da IEEE-754.

Per i tipi di elementi complessi, il confronto lessicografico delle coppie (real, imag) viene eseguito utilizzando gli attributi comparison_direction e compare_type forniti. Imporre un ordinamento su numeri complessi comporta una semantica sorprendente, quindi in futuro abbiamo in programma di rimuovere il supporto per i numeri complessi quando comparison_direction è GE, GT, LE o LT (#560).

Per i tipi quantizzati, esegue dequantize_compare(lhs, rhs, comparison_direction).

Input

Output

Vincoli

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs).
• (C2) shape(lhs) = shape(rhs) = shape(result).
• (C3) compare_type è definito come:
  • SIGNED se is_signed_integer(element_type(lhs)).
  • UNSIGNED se is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)).
  • FLOAT o TOTALORDER se is_float(element_type(lhs)).
  • FLOAT se is_complex(element_type(lhs)).

Esempi

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

 Altri esempi

complesso

Semantica

Esegue la conversione elemento per elemento in un valore complesso da una coppia di valori reali e immaginari, lhs e rhs, e produce un tensore result.

Input

Output

Vincoli

• (C1) type(lhs) = type(rhs).
• (C2) shape(result) = shape(lhs).
• (C3) element_type(result) ha il tipo complex<E> dove E = element_type(lhs).

Esempi

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

 Altri esempi

composito

Semantica

Incapsula un'operazione composta da altre operazioni StableHLO, che prende inputs e composite_attributes e produce results. La semantica dell'operazione è implementata dall'attributo decomposition. L'operazione composite può essere sostituita con la sua decomposizione senza modificare la semantica del programma. Nei casi in cui l'inserimento in linea della decomposizione non fornisca la stessa semantica dell'operatore, preferisci utilizzare custom_call.

Il campo version (il valore predefinito è 0) viene utilizzato per indicare quando cambia la semantica di un composito.

Input

Output

Vincoli

• (C1) is_namespaced_op_name(name)
• (C2) is_defined_in_parent_scope(decomposition)
• (C3) types(inputs...) == input_types(decomposition)
• (C4) types(results...) == output_types(decomposition)

Esempi

%results = "stablehlo.composite"(%input0, %input1) {
  name = "my_namespace.my_op",
  composite_attributes = {
    my_attribute = "my_value"
  },
  decomposition = @my_op,
  version = 1 : i32
} : (tensor<f32>, tensor<f32>) -> tensor<f32>

 Altri esempi

concatenate

Semantica

Concatena inputs lungo la dimensione dimension nello stesso ordine degli argomenti specificati e produce un tensore result. In forma più formale, result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1], dove:

1. id = d0 + ... + dk-1 + kd.
2. d è uguale a dimension e d0, ... sono le da dimensioni delle dimensioni inputs.

Input

Output

Vincoli

• (C1) same(element_type(inputs...)).
• (C2) same(shape(inputs...)) ad eccezione di dim(inputs..., dimension).
• (C3) 0 < size(inputs).
• (C4) 0 <= dimension < rank(inputs[0]).
• (C5) element_type(result) = element_type(inputs[0]).
• (C6) shape(result) = shape(inputs[0]) ad eccezione di:
  • dim(result, dimension) = dim(inputs[0], dimension) + ....

Esempi

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

Altri esempi

costante

Semantica

Produce un tensore output da una costante value.

Input

Output

Vincoli

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

Esempi

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

Altri esempi

convertire

Semantica

Esegue una conversione elemento per elemento da un tipo di elemento a un altro sul operand tensore e produce un result tensore.

Per le conversioni boolean-to-any-supported-type, il valore false viene convertito in zero e il valore true in uno. Per le conversioni any-supported-type-to-boolean, un valore pari a zero viene convertito in false, mentre i valori diversi da zero vengono convertiti in true. Continua a leggere per sapere come funziona per i tipi complessi.

Per le conversioni che coinvolgono intero a intero, intero a virgola mobile o virgola mobile a virgola mobile, se il valore di origine può essere rappresentato esattamente nel tipo di destinazione, il valore del risultato è la rappresentazione esatta. In caso contrario, il comportamento è da decidere (#180).

Per le conversioni che coinvolgono floating-point-to-integer, la parte frazionaria viene troncata. Se il valore troncato non può essere rappresentato nel tipo di destinazione, il comportamento è TBD (#180).

La conversione da complesso a complesso segue lo stesso comportamento delle conversioni da virgola mobile a virgola mobile per la conversione delle parti reali e immaginarie.

Per le conversioni da complesso a qualsiasi altro tipo e da qualsiasi altro tipo a complesso, il valore immaginario di origine viene ignorato o il valore immaginario di destinazione viene impostato su zero. La conversione della parte reale segue le conversioni a virgola mobile.

In linea di principio, questa operazione potrebbe esprimere la dequantizzazione (conversione da tensori quantizzati a tensori regolari), la quantizzazione (conversione da tensori regolari a tensori quantizzati) e la ricontizzazione (conversione tra tensori quantizzati), ma al momento abbiamo operazioni dedicate per questo: uniform_dequantize per il primo caso d'uso e uniform_quantize per il secondo e il terzo caso d'uso. In futuro, queste due operazioni potrebbero essere unite in convert (#1576).

Input

Output

Vincoli

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

Esempi

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

 Altri esempi

convoluzione

Semantica

Calcola i prodotti scalari tra finestre di lhs e sezioni di rhs e produce result. Il seguente diagramma mostra come vengono calcolati gli elementi in result da lhs e rhs utilizzando un esempio concreto.

Più formalmente, considera la seguente riorganizzazione degli input in termini di lhs per poter esprimere finestre di 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).

Questo nuovo inquadramento utilizza le seguenti funzioni di supporto:

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

Se feature_group_count = 1 e batch_group_count = 1, per tutti output_spatial_index in index_space(dim(result, output_spatial_dimensions...)), result[result_shape(:, output_spatial_index, :)] = dot_product dove:

• 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]). Questa funzionalità sembra non essere utilizzata, pertanto in futuro prevediamo di rimuoverla (#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]).

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

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

Per i tipi quantizzati, esegue 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)).

Per i tipi quantizzati ibridi, esegue hybrid_dequantize_then_op( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs).

Input

Output

Vincoli

• (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) Dato 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) Dato 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) In base a 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) Per dim(result, result_dim) si intende:
  • dim(lhs, input_batch_dimension) / batch_group_count se result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) se result_dim = output_feature_dimension.
  • num_windows altrimenti, dove:
  • 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.
• Se l'operazione utilizza tensori non quantizzati:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Se l'operazione utilizza tensori quantizzati:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Se is_per_axis_quantized(rhs), poi quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Se is_per_axis_quantized(result), then quantization_dimension(result) = output_feature_dimension.
  • Se is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Se is_per_tensor_quantized(rhs), then is_per_tensor_quantized(result).
  • Se !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Esempi

// %lhs: [[
//        [
//          [1], [2], [5], [6]
//        ],
//        [
//          [3], [4], [7], [8]
//        ],
//        [
//          [10], [11], [14], [15]
//        ],
//        [
//          [12], [13], [16], [17]
//        ]
//      ]]
//
// %rhs: [
//        [[[1]], [[1]], [[1]]],
//        [[[1]], [[1]], [[1]]],
//        [[[1]], [[1]], [[1]]]
//       ]
%result = "stablehlo.convolution"(%lhs, %rhs) {
  window_strides = array<i64: 4, 4>,
  padding = dense<0> : tensor<2x2xi64>,
  lhs_dilation = array<i64: 2, 2>,
  rhs_dilation = array<i64: 1, 1>,
  window_reversal = array<i1: false, false>,
  // In the StableHLO dialect, dimension numbers are encoded via:
  // `[<input dimensions>]x[<kernel dimensions>]->[output dimensions]`.
  // "b" is batch dimension, "f" is feature dimension,
  // "i" is input feature dimension, "o" is output feature dimension,
  // "0/1/etc" are spatial dimensions.
  dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
  batch_group_count = 1 : i64,
  feature_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>) -> tensor<1x2x2x1xi64>
// %result: [[
//            [[10], [26]],
//            [[46], [62]]
//          ]]

 Altri esempi

coseno

Semantica

Esegue l'operazione di coseno elemento per elemento sul tensore operand e produce un result tensore. A seconda del tipo di elemento, esegue le seguenti operazioni:

• Per i valori float: cos da IEEE-754.
• Per i numeri complessi: coseno complesso.
• Per i tipi quantizzati: dequantize_op_quantize(cosine, operand, type(result)).

Input

Output

Vincoli

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

Esempi

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

Altri esempi

count_leading_zeros

Semantica

Esegue il conteggio elemento per elemento del numero di bit iniziali pari nel tensore operand e produce un tensore result.

Input

Output

Vincoli

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

Esempi

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

 Altri esempi

custom_call

Semantica

Incapsula un'operazione call_target_name definita dall'implementazione che prende inputs e called_computations e produce results. has_side_effect, backend_config e api_version possono essere utilizzati per fornire metadati aggiuntivi definiti dall'implementazione.

Al momento, questa operazione contiene una raccolta di metadati piuttosto disorganizzata che riflette l'evoluzione organica dell'operazione di controparte nel compilatore XLA. In futuro, prevediamo di unificare questi metadati (#741).

Input

Output

Esempi

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

divisione

Semantica

Esegue la divisione elemento per elemento dei tensori dividendo lhs e divisore rhs e produce un tensore result. A seconda del tipo di elemento, esegue le seguenti operazioni:

• Per i numeri interi: divisione di numeri interi che produce il quoziente algebrico con qualsiasi parte frazionaria scartata.
• Per i valori float: division da IEEE-754.
• Per i numeri complessi: divisione complessa.
• Per i tipi quantizzati:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)).

Input

Output

Vincoli

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

Esempi

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

 Altri esempi

dot_general

Semantica

Calcola i prodotti scalari tra i vari slice di lhs e i vari slice di rhs e produce un result tensore.

In forma più formale, result[result_index] = dot_product, dove:

• 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 dove size(result_batching_index) = size(lhs_batching_dimensions), size(result_lhs_index) = size(lhs_result_dimensions) e size(result_rhs_index) = size(rhs_result_dimensions).
• transposed_lhs = transpose(lhs, lhs_batching_dimensions + lhs_result_dimensions + lhs_contracting_dimensions).
• transposed_lhs_slice = slice(transposed_lhs, result_batching_index + result_lhs_index + [:, ..., :]).
• reshaped_lhs_slice = reshape(transposed_lhs_slice, dims(lhs, lhs_contracting_dimensions)).
• transposed_rhs = transpose(rhs, rhs_batching_dimensions + rhs_result_dimensions + rhs_contracting_dimensions).
• transposed_rhs_slice = slice(transposed_rhs, result_batching_index + result_rhs_index + [:, ..., :]).
• reshaped_rhs_slice = reshape(transposed_rhs_slice, dims(rhs, rhs_contracting_dimensions)).
• dot_product = reduce( inputs=[multiply(reshaped_lhs_slice, reshaped_rhs_slice)], init_values=[constant(0, element_type(result))], dimensions=range(size(lhs_contracting_dimensions)), body=lambda x, y: add(x, y)).

Per i tipi quantizzati, esegue 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)).

Per i tipi quantizzati ibridi, esegue hybrid_dequantize_then_op( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs).

precision_config controlla il compromesso tra velocità e accuratezza per i calcoli sui backend degli acceleratori. Può essere uno dei seguenti valori (al momento, la semantica di questi valori enum è sottospecificata, ma prevediamo di risolvere questo problema in #755):

• DEFAULT: calcolo più rapido, ma approssimazione meno accurata del numero originale.
• HIGH: calcolo più lento, ma approssimazione più accurata del numero originale.
• HIGHEST: calcolo più lento, ma approssimazione più accurata del numero originale.

Un DotAlgorithm definisce le proprietà principali dell'algoritmo utilizzato per implementare l'operazione di moltiplicazione, che definisce anche la precisione. Se i campi dell'attributo algoritmo sono impostati, precision_config deve essere DEFAULT. DotAlgorithms non hanno un valore predefinito, poiché i parametri predefiniti sono definiti dall'implementazione. Di conseguenza, tutti i campi dell'algoritmo dei punti possono essere impostati su None per specificare un algoritmo dei punti vuoto, che utilizzerà invece il valore precision_config.

I campi DotAlgorithm includono:

• lhs_precision_type e rhs_precision_type, le precisioni a cui vengono arrotondati il primo e il secondo membro dell'operazione. I tipi di precisione sono indipendenti dai tipi di archiviazione degli input e dell'output.
• accumulation_type la precisione utilizzata per l'accumulo.
• lhs_component_count, rhs_component_count e num_primitive_operations vengono applicati quando utilizziamo un algoritmo che decompone il primo e/o il secondo membro in più componenti ed esegue più operazioni di moltiplicazione "primitive" su questi valori, in genere per emulare una precisione più elevata (ad es. Sfruttare il tipo di dati di IA bfloat16 per calcoli di maggiore precisione: bf16_6x tf32_3x e così via). Per gli algoritmi senza decomposizione, questi valori devono essere impostati su 1.
• allow_imprecise_accumulation per specificare se l'accumulo con una precisione inferiore è consentito per alcuni passaggi (ad es. CUBLASLT_MATMUL_DESC_FAST_ACCUM).

Attributi DotAlgorithm di esempio:

// Inputs are casted to tf32, and then accumulated in f32:
{lhs_precision_type = tf32,
 rhs_precision_type = tf32,
 accumulation_type = f32,
 lhs_component_count = 1,
 rhs_component_count = 1,
 num_primitive_operations = 1,
 allow_imprecise_accumulation = false}


// bf16_6x: each input is decomposed to 3 bf16 components, then 6 dot operations are done on those components, and the result is accumulated in f32.
{lhs_precision_type = bf16,
 rhs_precision_type = bf16,
 accumulation_type = f32,
 lhs_component_count = 3,
 rhs_component_count = 3,
 num_primitive_operations = 6,
 allow_imprecise_accumulation = false}


// Inputs are (casted to) f8e5m2, and we accumulate in f32, but for some steps we may accumulate in lower precision.
{lhs_precision_type = f8e5m2,
 rhs_precision_type = f8e5m2,
 accumulation_type = f32,
 lhs_component_count = 1,
 rhs_component_count = 1,
 num_primitive_operations = 1,
 allow_imprecise_accumulation = true}

Spetta alle implementazioni decidere quali combinazioni sono supportate. In generale, non è garantito che ogni algoritmo sia supportato su ogni tipo di acceleratore dal consumatore di StableHLO. Se un determinato algoritmo non è supportato, deve essere generato un errore anziché passare a un'alternativa. La verifica StableHLO viene eseguita secondo il criterio del massimo impegno, impedendo l'uso di algoritmi non supportati su nessun hardware.

Consulta xla_data.proto > Algorithm per alcuni valori dell'algoritmo supportati. Il ticket 2483 mostra il piano per la creazione di un documento centralizzato sugli algoritmi supportati dal backend.

Input

Output

Vincoli

• (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).
• Se l'operazione utilizza tensori non quantizzati:
  • (C13) element_type(lhs) = element_type(rhs).
• Se l'operazione utilizza tensori quantizzati:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C15) zero_points(rhs) = 0.
  • (C16) Se is_per_axis_quantized(rhs), significa quantization_dimension(rhs) non in rhs_contracting_dimensions.
  • Se is_quantized(lhs):
  • (C17) storage_type(lhs) = storage_type(rhs).
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C19) Se is_per_tensor_quantized(rhs), then is_per_tensor_quantized(result).
  • Se !is_quantized(lhs):
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result).
• Se !is_empty_algorithm(lhs_precision_type, rhs_precision_type, accumulation_type, lhs_component_count, rhs_component_count, num_primitive_operations allow_imprecise_accumulation):
  • (C21) precision_config... = DEFAULT.
  • (C22) 0 < lhs_component_count.
  • (C23) 0 < rhs_component_count.
  • (C24) 0 < num_primitive_operations.

Esempi

// %lhs: [
//        [[1, 2],
//         [3, 4]],
//        [[5, 6],
//         [7, 8]]
//       ]
// %rhs: [
//        [[1, 0],
//         [0, 1]],
//        [[1, 0],
//         [0, 1]]
//       ]
%result = "stablehlo.dot_general"(%lhs, %rhs) {
  dot_dimension_numbers = #stablehlo.dot<
    lhs_batching_dimensions = [0],
    rhs_batching_dimensions = [0],
    lhs_contracting_dimensions = [2],
    rhs_contracting_dimensions = [1]
  >,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>],
  algorithm = #stablehlo.dot_algorithm<
    lhs_precision_type = tf32,
    rhs_precision_type = tf32,
    accumulation_type = f32,
    lhs_component_count = 1,
    rhs_component_count = 1,
    num_primitive_operations = 1,
    allow_imprecise_accumulation = false
  >
} : (tensor<2x2x2xi64>, tensor<2x2x2xi64>) -> tensor<2x2x2xi64>
// %result: [
//           [[1, 2],
//            [3, 4]],
//           [[5, 6],
//            [7, 8]]
//          ]

 Altri esempi

dynamic_broadcast_in_dim

Semantica

Questa operazione è funzionalmente identica all'operazione broadcast_in_dim, ma la forma del risultato viene specificata dinamicamente tramite output_dimensions.

L'operazione accetta anche gli attributi facoltativi known_expanding_dimensions, known_nonexpanding_dimensions per esprimere conoscenze statiche sul comportamento di espansione delle dimensioni. Se non specificato, si presume che tutte le dimensioni possano espandersi.

Input

Output

Vincoli

• (C1) element_type(result) è dato da:
  • element_type(operand), se !is_per_axis_quantized(operand).
  • element_type(operand) ad eccezione dei seguenti casi: quantization_dimension(operand), scales(operand) e zero_points(operand) possono differire da quantization_dimension(result), scales(result) e zero_points(result) risp., altrimenti.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Per tutti i valori d in axes(operand):
  • dim(operand, d) = 1 o
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Se is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Se dim(operand, quantization_dimension(operand)) = 1, then scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))).
• (C7) size(output_dimensions) = rank(result).
• (C8) is_unique(known_expanding_dimensions + known_nonexpanding_dimensions).
• (C9) 0 <= known_expanding_dimensions < rank(operand).
• (C10) 0 <= known_nonexpanding_dimensions < rank(operand).

Esempi

// %operand: [
//            [1, 2, 3]
//           ]
%operand = stablehlo.constant dense<[[1, 2, 3]]> : tensor<1x3xi64>
%output_dimensions = stablehlo.constant dense<[2, 3, 2]> : tensor<3xi64>
%result = "stablehlo.dynamic_broadcast_in_dim"(%operand, %output_dimensions) {
  broadcast_dimensions = array<i64: 2, 1>,
  known_expanding_dimensions = array<i64: 0>,
  known_nonexpanding_dimensions = array<i64: 1>
} : (tensor<1x3xi64>, tensor<3xi64>) -> tensor<2x3x2xi64>
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

 Altri esempi

dynamic_conv

Semantica

Questa operazione è funzionalmente identica all'operazione op di convoluzione, ma il padding viene specificato dinamicamente tramite padding.

Input

Output

Vincoli

• (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) Dato 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) Dato 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) In base a 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) Per dim(result, result_dim) si intende:
  • dim(lhs, input_batch_dimension) / batch_group_count se result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) se result_dim = output_feature_dimension.
  • num_windows altrimenti, dove:
  • 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.
• Se l'operazione utilizza tensori non quantizzati:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Se l'operazione utilizza tensori quantizzati:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Se is_per_axis_quantized(rhs), poi quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Se is_per_axis_quantized(result), then quantization_dimension(result) = output_feature_dimension.
  • Se is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Se is_per_tensor_quantized(rhs), then is_per_tensor_quantized(result).
  • Se !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Esempi

// %lhs: [[
//        [[1], [2], [5], [6]],
//        [[3], [4], [7], [8]],
//        [[10], [11], [14], [15]],
//        [[12], [13], [16], [17]]
//      ]]
//
// %rhs: [
//         [[[1]], [[1]], [[1]]],
//         [[[1]], [[1]], [[1]]],
//         [[[1]], [[1]], [[1]]]
//        ]
// %padding: [[1, 1],
//            [1, 1]]
%result = "stablehlo.dynamic_conv"(%lhs, %rhs, %padding) {
  window_strides = array<i64: 4, 4>,
  lhs_dilation = array<i64: 2, 2>,
  rhs_dilation = array<i64: 1, 1>,
  window_reversal = array<i1: false, false>,
  dimension_numbers = #stablehlo.conv<raw
    input_batch_dimension = 0,
    input_feature_dimension = 3,
    input_spatial_dimensions = [0, 1],
    kernel_input_feature_dimension = 2,
    kernel_output_feature_dimension = 3,
    kernel_spatial_dimensions = [0, 1],
    output_batch_dimension = 0,
    output_feature_dimension = 3,
    output_spatial_dimensions = [1, 2]
  >,
  feature_group_count = 1 : i64,
  batch_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>, tensor<2x2xi64>) -> tensor<1x2x2x1xi64>
// %result: [[
//            [[1], [5]],
//            [[10], [14]]
//          ]]