Specifica StableHLO

StableHLO è un insieme di operazioni per operazioni di alto livello (HLO) nei modelli di machine learning (ML). StableHLO funge da 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 ML (come TensorFlow, JAX e PyTorch) e compilatori ML (come XLA e IREE). A questo scopo, questo 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 che sono costituiti da funzioni StableHLO che a loro volta sono costituite da operazioni StableHLO. All'interno di questa struttura, la sezione Ops specifica la semantica delle singole operazioni. La sezione Esecuzione fornisce la semantica per tutte queste operazioni eseguite insieme all'interno di un programma. Infine, la sezione Notazione illustra la notazione utilizzata in tutta la specifica.

Per visualizzare le specifiche di una release precedente di StableHLO, apri il repository nella release taggata di tuo interesse. Ad esempio, le specifiche di StableHLO v0.19.0. Per visualizzare le modifiche apportate a ogni aggiornamento della versione secondaria di StableHLO, consulta il log delle versioni 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 agli identificatori 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 valori possono essere completamente numerici per semplificare la generazione di programmi StableHLO.

Tipi

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

I tipi StableHLO sono suddivisi in tipi di valori (chiamati anche tipi di prima classe) che rappresentano i valori StableHLO e tipi non di valori che descrivono altri elementi del programma. I tipi StableHLO sono simili a quelli di molti linguaggi di programmazione, con la principale peculiarità di essere specifici per il dominio, il che comporta alcuni risultati insoliti (ad es. 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 array multidimensionali. Hanno una forma e un tipo di elemento, dove una forma rappresenta dimensioni non negative o sconosciute delle dimensioni in ordine crescente delle dimensioni corrispondenti (chiamate anche assi) numerate da 0 a R-1. Il numero di dimensioni R è chiamato rank. 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 ?. Non è possibile rimuovere la classificazione delle forme.

In futuro, prevediamo di esplorare l'estensione dei tipi di tensori oltre le dimensioni e i tipi di elementi, ad esempio per includere layout (#629) e 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 valori interi di un tipo di archiviazione nell'intervallo da storage_min a storage_max (inclusi) che corrispondono a valori in virgola mobile di un tipo espresso. Per un determinato valore intero i, il valore in virgola mobile corrispondente 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 valori predefiniti di min_value(storage_type) e max_value(storage_type) rispettivamente. I tipi di elementi quantizzati presentano 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 in virgola mobile, ma esiste un forte interesse per le scale basate su numeri interi, rappresentate con moltiplicatori e spostamenti. Abbiamo in programma di esplorare questa possibilità nel prossimo futuro (#1404).

È in corso una discussione sulla semantica di QuantizationZeroPoint, inclusi il tipo, i valori e la possibilità di avere uno o più punti zero in un tipo di tensore quantizzato. In base ai risultati di questa discussione, la specifica relativa a zero punti 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, stiamo pianificando di esplorare la rappresentazione di scale e punti zero sconosciuti, in modo simile a come stiamo pianificando di esplorare la rappresentazione di dimensioni sconosciute (#1407).

I tipi di tensori quantizzati rappresentano i tensori con elementi quantizzati. Questi tensori sono esattamente uguali ai tensori normali, tranne per il 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 con un scale e un zero_point per l'intero tensore, oppure per asse, ovvero con più scales e zero_points, una coppia per ogni sezione di una determinata dimensione quantization_dimension. Più formalmente, in un tensore t con quantizzazione per asse, ci sono dim(t, quantization_dimension) sezioni di quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :], ecc. Tutti gli elementi nella i-esima sezione utilizzano scales[i] e zero_points[i] come parametri di quantizzazione. I tipi di tensore quantizzato 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 utilizzati da alcune operazioni. I token vengono utilizzati per imporre l'ordine di esecuzione delle operazioni come descritto nella sezione Esecuzione.

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

I tipi di buffer rappresentano i buffer. Ad esempio, in XLA i buffer sono array multidimensionali con spazio di archiviazione coerente. Analogamente ai tipi di tensore, i tipi di buffer hanno una forma e un tipo di elemento, dove una forma rappresenta dimensioni non negative o sconosciute delle dimensioni in ordine crescente delle dimensioni corrispondenti (chiamate anche assi) numerate da 0 a R-1. Il numero di dimensioni R è chiamato rank. Ad esempio, memref<2x3xf32> è un tipo di buffer 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.

I buffer possono essere allocati utilizzando un valore compreso tra custom_call e CreateBuffer o Pin e deallocati tramite un valore compreso tra custom_call e Unpin. Solo gli operatori custom_call possono leggere e scrivere i contenuti all'interno dei buffer. Per maggiori dettagli, consulta custom_call.

I tipi di tuple rappresentano le tuple, ovvero elenchi eterogenei. Le tuple sono una funzionalità legacy che esiste solo per la compatibilità con HLO. In HLO, le tuple vengono utilizzate per rappresentare input e output variadici. In StableHLO, gli input e gli output variadici sono supportati in modo nativo e l'unico utilizzo delle tuple in StableHLO è quello di rappresentare in modo completo l'ABI HLO, in cui, ad esempio, T, tuple<T> e tuple<tuple<T>> possono essere materialmente diversi a seconda di una particolare implementazione. In futuro, prevediamo di apportare modifiche all'ABI HLO che potrebbero consentirci di rimuovere i tipi di tupla 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 tensore. A differenza di molti linguaggi di programmazione, questi tipi non sono di prima classe in StableHLO. Ciò significa che i programmi StableHLO non possono rappresentare direttamente i valori di questi tipi (di conseguenza, è idiomatico rappresentare i valori scalari di tipo T con tensori di tipo tensor<T> di dimensione 0).

• Il tipo booleano rappresenta i valori booleani true e false.
• I tipi interi possono essere con segno (si) o senza segno (ui) e avere una delle larghezze in 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 in virgola mobile a 8 bit che seguono le convenzioni IEEE-754.
  • 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 Hybrid 8-bit Floating Point (HFP8) Training and Inference for Deep Neural Networks.
  • Tipo bf16 corrispondente al formato bfloat16 descritto in BFloat16: il segreto delle prestazioni elevate su Cloud TPU.
  • f16, f32 e f64 corrispondenti rispettivamente ai formati binary16 ("mezza precisione"), binary32 ("precisione singola") e binary64 ("doppia precisione") descritti nello standard IEEE 754.
  • Il tipo tf32 corrisponde al formato TensorFloat32 ed è supportato in modo limitato in StableHLO.
  • f4E2M1FN, f6E2M3FN, f6E3M2FN e f8E8M0FNU MX (microscaling) descritti nella specifica dei formati di microscaling 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 dei tipi sul lato sinistro di ->) e tipi di output (l'elenco dei tipi sul lato destro di ->). In molti linguaggi di programmazione, i tipi di funzioni sono di prima classe, ma non in StableHLO.

StringType ::= 'string'

Il tipo 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 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 in precedenza, la sintassi di StableHLO è fortemente ispirata a MLIR, che non è necessariamente l'alternativa più ergonomica, ma è probabilmente la più adatta allo 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 StableHLO (chiamate anche ops) hanno un nome, input/output e una firma. Il nome è composto 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 utilizzano input e producono output. Gli input sono suddivisi in valori di input (calcolati durante l'esecuzione), funzioni di input (fornite staticamente, perché in StableHLO le funzioni non sono valori di prima classe) e attributi di input (forniti anche staticamente). Il tipo di input e output utilizzati e prodotti da un'operazione dipende dal relativo mnemonico. Ad esempio, l'operazione add utilizza 2 valori di input e produce 1 valore di output. Al contrario, l'operatore select_and_scatter utilizza 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 cui il nome "anonime "), 2) non dichiarano i tipi di output (i tipi di output vengono dedotti dall'operatore return all'interno della funzione).

La sintassi delle funzioni di input include una parte attualmente inutilizzata (vedi la produzione Unused sopra) che serve 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 alla produzione di Unused, in modo che possano essere distinti l'uno dall'altro. 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 i 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 valori di input. Analogamente, l'operatore 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 natura a volte contengono attributi che non sono descritti in questo documento. In futuro, prevediamo di incorporare questi attributi nell'opset 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 posizione (#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 ->). A rigor di termini, i tipi di input sono ridondanti e anche i tipi di output lo sono quasi sempre (perché per la maggior parte delle operazioni StableHLO, i tipi di output possono essere dedotti dagli input). Tuttavia, la firma dell'operazione 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). Tieni presente che la firma dell'operazione include solo i tipi dei valori di input (ma non i tipi di funzioni e attributi di input forniti inline).

%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 è non ambiguo (ad es. una costante booleana ha inequivocabilmente 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 binaria o ottale, non sono supportate. Le costanti intere presentano 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 in virgola mobile rappresentano valori in 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 in virgola mobile del tipo corrispondente. Le costanti in virgola mobile presentano i seguenti vincoli:

• (C1) Se viene utilizzata una 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 (prima) e una parte immaginaria (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 tensoriali rappresentano i valori tensoriali utilizzando elenchi nidificati specificati tramite la notazione NumPy. Ad esempio, dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> rappresenta un valore tensore con il seguente mapping 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 tensoriali presentano 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 tensoriali quantizzate rappresentano i valori tensoriali quantizzati utilizzando la stessa notazione delle costanti tensoriali, con gli elementi specificati come costanti del tipo di archiviazione. Le costanti tensoriali quantizzate presentano 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))

I valori letterali stringa sono costituiti da byte specificati utilizzando caratteri ASCII e sequenze di escape. Sono indipendenti dalla codifica, quindi l'interpretazione di questi byte è definita dall'implementazione. I valori letterali stringa hanno il tipo string.

Operazioni

abs

Semantica

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

• Per i numeri interi con segno: modulo intero.
• Per i numeri in virgola mobile: 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) altrimenti.

Esempi

// %operand: [-2, 0, 2]
%result = "stablehlo.abs"(%operand)< : (t>ens>or3xi32<) - t>ensor3xi32
// %result: [2, 0, 2]

 Altri esempi

add

Semantica

Esegue l'addizione elemento per 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: OR logico.
• Per i numeri interi: addizione di numeri interi.
• Per i numeri in virgola mobile: 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)< : (ten>sor2x2xi<32, ten>sor>2x2xi32<) - ten>sor2x2xi32
// %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, !stablehl>o.token) - !stablehlo.token

 Altri esempi

all_gather

Semantica

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

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

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

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 viene utilizzato 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...) tranne:
  • 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_grou<ps = den>se[[0, 1]<] : ten>sor1x2xi64,
  // channel_id = 0
  channel_handle = #stablehlo.chan<nel_handlehandle = 0>, type = 0
  // use_global_device_ids = false
}< : (ten>sor2x2xi<64, ten>sor>2x2xi64)< - (ten>sor2x4xi<64, ten>sor2x4xi64)
// %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 dei processi StableHLO, applica una funzione di riduzione computation ai valori dei tensori operands di ogni processo e produce tensori results.

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

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

Successivamente, all'interno di ogni process_group:

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

Input

Output

Vincoli

• (C1) is_unique(replica_groups).
• (C2) size(replica_groups) è definito come:
  • 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(%ar<g0:> tensori64, %ar<g1:> tensori64):
    %0 = "stablehlo.add"(%ar<g0,> %arg1) <: (>ten>sori64,< te>nsori64) - tensori64
    "stable<hlo>.re>turn"(%0) : (tensori64) - ()<
}) {
  >replica_g<roups => dense[[0, 1]] : tensor1x2xi64,
  // channel_id = 0
  channel_hand<le = #stablehlo.chan>nel_handlehandle = 0, type = 0
  // use_global_<devic>e_ids = <false>
} >: (tenso<r4xi6>4, tenso<r4xi6>4) - (tensor4xi64, tensor4xi64)
// %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 dei processi StableHLO, divide i valori dei tensori operands lungo split_dimension in parti, distribuisce le parti divise tra i processi, concatena le parti distribuite lungo concat_dimension e produce tensori results. L'operazione divide la griglia di processi StableHLO in process_groups, che è definito come segue:

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

Successivamente, 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 viene utilizzato 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_grou<ps = den>se[[0, 1]<] : ten>sor1x2xi64
  // channel_id = 0
}< : (ten>sor2x4xi<64, ten>sor>2x4xi64)< - (ten>sor4x2xi<64, ten>sor4x2xi64)
// %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'operazione AND elemento per 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 bit a 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)< : (ten>sor2x2xi<32, ten>sor>2x2xi32<) - ten>sor2x2xi32
// %result: [[1, 2], [3, 0]]

 Altri esempi

atan2

Semantica

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

• Per i numeri in virgola mobile: atan2 da 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)< : (t>ensor3xf<64, t>ens>or3xf64<) - t>ensor3xf64
// %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 batch_norm_training mediante la retropropagazione da grad_output e produce i tensori grad_operand, grad_scale e grad_offset. Più formalmente, questa operazione può essere espressa come una scomposizione in 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
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ensor2xf64,
 <    tenso>r2x>2x2xf64)< - (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf64)
// %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 scomposizione in operazioni StableHLO esistenti utilizzando la sintassi Python nel seguente modo:

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
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ens>or2xf64<) - tenso>r2x2x2xf64
// %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, ad eccezione della dimensione feature_index, e normalizza il tensore operand producendo i tensori output, batch_mean e batch_var. Più formalmente, questa operazione può essere espressa come una scomposizione in 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
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ens>or2xf64) -
 <   (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf64)
// %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 in cui i bit dell'intero tensore operand vengono reinterpretati utilizzando il tipo del tensore result.

Più formalmente, 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, così come la rappresentazione esatta dei tipi di elementi.

Input

Output

Vincoli

• (C1) Dato 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)< : >(te>nsorf64<) - t>ensor4xf16
// %result: [0xCDEF, 0x89AB, 0x4567, 0x0123] // little-endian representation

 Altri esempi

broadcast_in_dim

Semantica

Espande le dimensioni e/o il rango di un tensore di input duplicando i dati nel tensore operand e produce un tensore result. In termini più formali, 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]] altrimenti.

Input

Output

Vincoli

• (C1) element_type(result) è fornito da:
  • element_type(operand), se !is_per_axis_quantized(operand).
  • element_type(operand), tranne che quantization_dimension(operand), scales(operand) e zero_points(operand) potrebbero differire da quantization_dimension(result), scales(result) e zero_points(result) rispettivamente.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Per tutti i 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, allora 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_dimensio<ns = arra>yi64: 2, 1
}< : (ten>sor>1x3xi32<) - tenso>r2x3x2xi32
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

 Altri esempi

richiesta

Semantica

Produce l'output dell'esecuzione di esattamente una funzione da branches a seconda del valore di index. In termini più formali, result = selected_branch() dove:

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

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, %resul<t_bra>nch0) : <(tens>or2>xi64, tensor2xi64) - ()
}, {
  "stablehlo.return"(%result_branc<h1, %>result_b<ranch>1) >: (tensor2xi64, <ten>sor>2xi64) -< ()
}>) : (ten<sori3>2) - (tensor2xi64, tensor2xi64)
// %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 tensore result. A seconda del tipo di elemento, esegue le seguenti operazioni:

• Per i numeri in virgola mobile: 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)< : (t>ens>or4xf64<) - t>ensor4xf64
// %result: [0.0, 1.0, 2.0, 3.0]

 Altri esempi

ceil

Semantica

Esegue l'operazione ceil elemento per elemento del tensore operand e produce un tensore result. Implementa l'operazione roundToIntegralTowardPositive dalla 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)< : (t>ens>or5xf32<) - t>ensor5xf32
// %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 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 stretto o il triangolo inferiore stretto, sono definiti dall'implementazione.

Se esiste i in cui la matrice di input non è una matrice hermitiana definita positiva, 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
}< : (ten>sor>3x3xf32<) - ten>sor3x3xf64
// %result: [
//           [1.0, 0.0, 0.0],
//           [2.0, 4.0, 0.0],
//           [3.0, 5.0, 6.0]
//          ]

clampare

Semantica

Blocca ogni elemento del tensore operand tra un valore minimo e massimo e produce un tensore result. Più formalmente, 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)).

L'imposizione di un ordinamento sui numeri complessi comporta una semantica sorprendente, quindi in futuro prevediamo di rimuovere il supporto per i 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)< : (t>ensor3xi<32, t>ensor3xi<32, t>ens>or3xi32<) - t>ensor3xi32
// %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 genera un tensore result.

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

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

Successivamente, result@process viene fornito da:

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

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_grou<ps = den>se[[2, 1]<] : ten>sor1x2xi64,
  channel_handle = #stablehlo.chan<nel_handlehandle = 0>, type = 0
} : (ten>sor>1x2xi64<) - ten>sor1x2xi64
// %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 dei processi StableHLO, invia il valore del tensore operand dal processo di origine al processo di destinazione e produce un tensore result.

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

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

Successivamente, result@process viene fornito 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)) altrimenti.

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 viene utilizzato 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_pai<rs = dense[[0, 1>], [1, 2]<] : ten>sor2x2xi64,
  channel_handle = #stablehlo.chan<nel_handlehandle = 0>, type = 0
}< : (ten>sor>2x2xi64<) - ten>sor2x2xi64
//
// %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'operazione utilizza la combinazione delle operazioni totalOrder e compareQuietEqual di IEEE-754.

Per i tipi di elementi complessi, il confronto lessicografico delle coppie (real, imag) viene eseguito utilizzando comparison_direction e compare_type forniti. L'imposizione di un ordinamento sui numeri complessi comporta una semantica sorprendente, quindi in futuro prevediamo 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 = <#stablehlocomparison_di>rection LT,
  compare_type = <#stablehlocomparison_>type FLOAT
}< : (t>ensor2xf<32, t>ens>or2xf32<) - >tensor2xi1
// %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) è di tipo complex<E> dove E = element_type(lhs).

Esempi

// %lhs: [1.0, 3.0]
// %rhs: [2.0, 4.0]
%result = "stablehlo.complex"(%lhs, %rhs)< : (t>ensor2xf<64, t>ens>or2xf64<) - tenso<r2x>>complexf64
// %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 viene implementata dall'attributo decomposition. L'operatore composite può essere sostituito con la sua decomposizione senza modificare la semantica del programma. Nei casi in cui l'incorporamento della scomposizione non fornisce la stessa semantica dell'operazione, preferisci utilizzare custom_call.

Il campo version (il valore predefinito è 0) viene utilizzato per indicare quando cambiano le semantiche 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,
 < ve>rsion = <1 :> i3>2
} : (<ten>sorf32, tensorf32) - tensorf32

 Altri esempi

concatenare

Semantica

Concatena inputs lungo la dimensione dimension nello stesso ordine degli argomenti forniti e produce un tensore result. In termini più formali, 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 dimensioni della da dimensione di 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
}< : (ten>sor3x2xi<64, ten>sor>1x2xi64<) - ten>sor4x2xi64
// %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"() {
  val<ue = dense[[0.0, 1.0], [>2.0, 3.0]<] : ten>sor2x2xf3>2
} : (<) - ten>sor2x2xf32
// %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 tensore operand e produce un tensore result.

Per le conversioni boolean-to-any-supported-type, il valore false viene convertito in zero e il valore true viene convertito 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 scoprire come funziona per i tipi complessi.

Per le conversioni che coinvolgono interi, interi 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 definire (#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 è da definire (#180).

La conversione che coinvolge numeri complessi in numeri complessi segue lo stesso comportamento delle conversioni da virgola mobile a virgola mobile per la conversione di 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 azzerato, rispettivamente. La conversione della parte reale segue le conversioni in 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 riquantizzazione (conversione tra tensori quantizzati), ma al momento disponiamo di operazioni dedicate per questo scopo: uniform_dequantize per il primo caso d'uso e uniform_quantize per il secondo e il terzo. 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)< : (t>ens>or3xi64<) - tenso<r3x>>complexf64
// %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.

In termini più formali, considera la seguente riformulazione 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).

Questa riformulazione utilizza le seguenti funzioni helper:

• 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, allora 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 inutilizzata, quindi 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) Dato output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]:
  • is_unique(output_dimensions).
  • 0 <= output_dimensions < N.
• (C21) 0 < feature_group_count.
• (C22) 0 < batch_group_count.
• (C23) feature_group_count = 1 or batch_group_count = 1.
• (C24) size(precision_config) = 2.
• (C25) dim(result, result_dim) è definito come:
  • 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), allora quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Se is_per_axis_quantized(result), allora 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), allora 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_strid<es = arra>yi64: 4, 4,
  paddi<n>g = dense<0 : ten>sor2x2xi64,
  lhs_dilati<on = arra>yi64: 2, 2,
  rhs_dilati<on = arra>yi64: 1, 1,
  window_revers<al = arrayi1: fa>lse, false,
  // In the StableHLO dialect, dimension numbers are encoded vi<a:
  // `[input >dim<ensions]x[kernel >di>mensions]-[output dimensions]`.
  // "b" is batch dimension, "f" is feature dimension,
  // "i" is input feature dimension, "o" is output feature dimension,
  // "0/1/etc" a<re spatial dimensions.
  d>imension_num>bers = #stablehlo.conv[b, 0, 1, f]x[0, 1, i, o]-[b, 0, 1, f],
  batch_group_count = 1 : i64,
  fea<ture_group_count >= 1 : i64,
 < precision_config> = [#stablehl<oprecision >DEFAULT,< #stablehlo>pre>cision <DEFAULT]
} >: (tensor1x4x4x1xi64, tensor3x3x1x1xi64) - tensor1x2x2x1xi64
// %result: [[
//            [[10], [26]],
//            [[46], [62]]
//          ]]

 Altri esempi

coseno

Semantica

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

• Per i numeri in virgola mobile: 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)< : (ten>sor>2x2xf32<) - ten>sor2x2xf32
// %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 zero iniziali 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)< : (ten>sor>2x2xi64<) - ten>sor2x2xi64
// %result: [[64, 63], [56, 0]]

 Altri esempi

custom_call

Semantica

Incapsula un'operazione definita dall'implementazione call_target_name che accetta 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 della sua operazione omologa nel compilatore XLA. In futuro, prevediamo di unificare questi metadati (#741).

Input

Output

(Supporto GPU XLA) Target custom_call speciali

Esistono tre call_target_name speciali correlati ai tipi buffer: CreateBuffer crea un buffer non inizializzato, Pin crea un buffer inizializzato e Unpin dealloca un buffer e restituisce il contenuto del buffer.

%uninitialized_buffer = "stablehlo.custom_call"() {
  call_target_name = "CreateBuffer",
  api_version> = 4 : <i32,
>} : () - memref4xf64

%initialized_buffer = "stablehlo.custom_call"(%init_value) {
  call_target_name = "Pin&quo<t;,
 > ap>i_versi<on = >4 : i32,
} : (tensor4xf64) - memref4xf64

%dealloc_buffer = "stablehlo.custom_call"(%initialized_buffer) {
  call_target_na<me = >&qu>ot;Unpi<n&quo>t;,
  api_version = 4 : i32,
} : (memref4xf64) - tensor4xf64

Alias

Alcune operazioni custom_call potrebbero richiedere che una parte degli output e una parte degli operandi condividano la stessa memoria. Questo valore può essere espresso tramite output_operand_aliases. Una rappresentazione di una coppia di alias è costituita da un elenco di indici di tuple di output che rappresentano la parte di output e da un operand_index insieme a un elenco di indici di tuple di operandi che rappresentano la parte di operando. L'elenco degli indici di output o di tuple di operandi è vuoto se il tipo corrispondente non è un tipo tuple e può essere arbitrariamente lungo per un tipo di tupla arbitrariamente nidificato. Questo è simile alla rappresentazione dell'alias XLA.

La parte di output e la parte di input in una coppia di alias devono avere lo stesso tipo. Per le operazioni personalizzate di chiamata che non sono chiamate a CreateBuffer, Pin e Unpin, un operando buffer può essere visualizzato al massimo in una coppia di alias e un output buffer deve essere visualizzato in una coppia di alias.

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 <= [>@fo>o]
} : <(te>nsorf64) - tensorf64

%updated_buffer = "stablehlo.custom_call"(%buffer) {
  call_target_name = "Update",
  api_version = 4 : i32,
  output_operand_aliases< = [
    #stablehlo.output_operand_aliasoutput_tuple_indices = [],
      operand_ind>ex = 0,
     < oper>and>_tuple_<indic>es = []]
} : (memref4xf64) - memref4xf64

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 intera che produce il quoziente algebrico con qualsiasi parte frazionaria scartata.
• Per i numeri in virgola mobile: 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)< : (t>ensor4xf<32, t>ens>or4xf32<) - t>ensor4xf32
// %result: [5.66666651, -5.66666651, -5.66666651, 5.66666651]

 Altri esempi

dot_general

Semantica

Calcola i prodotti scalari tra le sezioni di lhs e le sezioni di rhs e produce un tensore result.

Più formalmente, 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 precisione per i calcoli sui backend dell'acceleratore. Può essere uno dei seguenti valori (al momento, la semantica di questi valori enum non è specificata, ma prevediamo di risolvere il problema in #755):

• DEFAULT: Calcolo più veloce, ma approssimazione meno precisa del numero originale.
• HIGH: Calcolo più lento, ma approssimazione più accurata al numero originale.
• HIGHEST: Calcolo più lento, ma approssimazione più accurata al numero originale.

Un DotAlgorithm definisce le proprietà principali dell'algoritmo utilizzato per implementare l'operazione dot, che definisce anche la precisione. Se i campi dell'attributo algoritmo sono impostati, precision_config deve essere DEFAULT. DotAlgorithms non hanno un valore predefinito, in quanto i parametri predefiniti sono definiti dall'implementazione. Pertanto, tutti i campi dell'algoritmo punto possono essere impostati su None per specificare un algoritmo punto 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 i lati sinistro e destro 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 eseguiamo un algoritmo che scompone il lato sinistro e/o destro in più componenti ed esegue più operazioni "primitive" sui valori, di solito per emulare una precisione maggiore (ad es. Utilizzo del tipo di dati di intelligenza artificiale 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 a 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 consumer di StableHLO. Se un determinato algoritmo non è supportato, deve essere generato un errore anziché ricorrere a un'alternativa. La verifica StableHLO fornirà una verifica ottimale, impedendo l'utilizzo di algoritmi che non sono noti per essere supportati su qualsiasi hardware.

Consulta xla_data.proto > Algorithm per alcuni valori di algoritmo supportati. Il ticket n. 2483 descrive il piano per creare 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), allora 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), allora 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 = #sta<blehlo.dot
    lhs_batching_dimensions = [0],
    rhs_batching_dimensions = [0],
    lhs_contracting_dimensions = [2],
    rhs_contracting_dimension>s = [1]
  ,
  precision_config = [<#stablehloprecisi>on DEFAULT, <#stablehloprecisi>on 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
  
}< : (tenso>r2x2x2xi<64, tenso>r2x>2x2xi64<) - tenso>r2x2x2xi64
// %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 la conoscenza statica del comportamento di espansione delle dimensioni. Se non specificato, si presume che tutte le dimensioni possano essere in espansione.

Input

Output

Vincoli

• (C1) element_type(result) è fornito da:
  • element_type(operand), se !is_per_axis_quantized(operand).
  • element_type(operand), tranne che quantization_dimension(operand), scales(operand) e zero_points(operand) potrebbero differire da quantization_dimension(result), scales(result) e zero_points(result) rispettivamente.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Per tutti i 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, allora 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_dimensio<ns = arra>yi64: 2, 1,
  known_expanding_dimensio<ns = a>rrayi64: 0,
  known_nonexpanding_dimensio<ns = a>rrayi64: 1
}< : (ten>sor1x3xi<64, t>ens>or3xi64<) - tenso>r2x3x2xi64
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

 Altri esempi

dynamic_conv

Semantica

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

Input