Especificação do StableHLO

O StableHLO é um conjunto de operações para operações de alto nível (HLO, na sigla em inglês) em modelos de machine learning (ML). O StableHLO funciona como uma camada de portabilidade entre diferentes frameworks e compiladores de ML: os frameworks de ML que produzem programas StableHLO são compatíveis com os compiladores de ML que consomem programas StableHLO.

Nosso objetivo é simplificar e acelerar o desenvolvimento de ML, criando mais interoperabilidade entre vários frameworks de ML (como TensorFlow, JAX e PyTorch) e compiladores de ML (como XLA e IREE). Para isso, este documento fornece uma especificação para a linguagem de programação StableHLO.

Esta especificação contém três seções principais. Primeiro, a seção Programs descreve a estrutura dos programas StableHLO, que consistem em funções StableHLO, que por sua vez consistem em operações StableHLO. Dentro dessa estrutura, a seção Ops especifica a semântica de operações individuais. A seção Execução fornece semântica para todas essas operações executadas juntas em um programa. Por fim, a seção Notação discute a notação usada em toda a especificação.

Para ver as especificações de uma versão anterior do StableHLO, abra o repositório na versão marcada (link em inglês) que quiser. Por exemplo, a especificação do StableHLO v0.19.0 (link em inglês). Para conferir as mudanças que ocorreram em cada promoção de versão secundária do StableHLO, consulte o registro da versão em VhloDialect.td (links em inglês).

Programas

Program ::= {Func}

Os programas StableHLO consistem em um número arbitrário de funções StableHLO. Abaixo está um programa de exemplo com uma função @main que tem três entradas (%image, %weights e %bias) e uma saída. O corpo da função tem seis operações.

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

Funções

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

As funções StableHLO (também chamadas de funções nomeadas) têm um identificador, entradas/saídas e um corpo. No futuro, planejamos introduzir mais metadados para funções para alcançar uma melhor compatibilidade com o HLO (#425, #626, #740, #744).

Identificadores

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

Os identificadores StableHLO são semelhantes aos identificadores em muitas linguagens de programação, com duas peculiaridades: 1) todos os identificadores têm selos que diferenciam diferentes tipos de identificadores; 2) os identificadores de valor podem ser completamente numéricos para simplificar a geração de programas StableHLO.

Tipos

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

Os tipos de StableHLO são categorizados em tipos de valor, também chamados de tipos de primeira classe, que representam valores StableHLO e tipos não de valor, que descrevem outros elementos do programa. Os tipos StableHLO são semelhantes aos tipos de muitas linguagens de programação, com a principal peculiaridade sendo a natureza específica do domínio do StableHLO, que resulta em alguns resultados incomuns (por exemplo, tipos escalares não são tipos de valor).

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

Os tipos de tensor representam tensores, ou seja, matrizes multidimensionais. Eles têm uma forma e um tipo de elemento, em que uma forma representa tamanhos de dimensão não negativos ou desconhecidos na ordem crescente das dimensões correspondentes (também chamadas de eixos) numeradas de 0 a R-1. O número de dimensões R é chamado de rank. Por exemplo, tensor<2x3xf32> é um tipo de tensor com forma 2x3 e tipo de elemento f32. Ele tem duas dimensões (ou seja, dois eixos): a dimensão 0 e a 1, cujos tamanhos são 2 e 3. A classificação é 2.

As formas podem ser parcialmente ou totalmente desconhecidas (dinâmicas), por exemplo, tensor<?x2xf64> é parcialmente desconhecida e tensor<?x?xf64> é totalmente desconhecida. Os tamanhos de dimensão dinâmica são representados por um ?. Não é possível classificar formas.

No futuro, planejamos explorar a extensão dos tipos de tensores além dos tamanhos de dimensão e dos tipos de elemento, por exemplo, para incluir layouts (#629) e esparsidade (#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

Os tipos de elementos quantizados representam valores inteiros de um tipo de armazenamento no intervalo de storage_min a storage_max (inclusive) que correspondem a valores de ponto flutuante de um tipo expresso. Para um determinado valor inteiro i, o valor de ponto flutuante correspondente f pode ser calculado como f = (i - zero_point) * scale, em que scale e zero_point são chamados de parâmetros de quantização. O storage_min e o storage_max são opcionais na gramática, mas têm valores padrão de min_value(storage_type) e max_value(storage_type), respectivamente. Os tipos de elementos quantizados têm as seguintes restrições:

• (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), então size(scales) = 1.
• (C11) 0 <= quantization_dimension.

No momento, QuantizationScale é uma constante de ponto flutuante, mas há um grande interesse em escalas baseadas em números inteiros, representadas com multiplicadores e mudanças. Planejamos explorar isso em breve (#1404).

Há uma discussão em andamento sobre a semântica de QuantizationZeroPoint, incluindo o tipo, os valores e se pode haver apenas um ou possivelmente vários pontos nulos em um tipo de tensor quantizado. Com base nos resultados dessa discussão, a especificação em torno de pontos zero pode mudar no futuro (#1405).

Outra discussão em andamento envolve a semântica de QuantizationStorageMin e QuantizationStorageMax para determinar se alguma restrição precisa ser imposta a esses valores e aos valores de tensores quantizados (#1406).

Por fim, planejamos representar escalas e pontos desconhecidos, de forma semelhante à representação de tamanhos de dimensão desconhecidos (#1407).

Os tipos de tensores quantizados representam tensores com elementos quantizados. Esses tensores são exatamente os mesmos que os regulares, exceto que seus elementos têm tipos de elementos quantizados, em vez de tipos de elementos regulares.

Em tensores quantizados, a quantização pode ser por tensor, ou seja, ter um scale e zero_point para todo o tensor, ou pode ser por eixo, ou seja, ter vários scales e zero_points, um par por fração de uma dimensão específica quantization_dimension. Mais formalmente, em um tensor t com quantização por eixo, há fatias dim(t, quantization_dimension) do quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :], etc. Todos os elementos na iª fatia usam scales[i] e zero_points[i] como parâmetros de quantização. Os tipos de tensores quantizados têm as seguintes restrições:

• Para quantização por tensor:
  • Sem restrições adicionais.
• Para quantização por eixo:
  • (C12) quantization_dimension < rank(self).
  • (C13) dim(self, quantization_dimension) = size(scales).

TokenType ::= 'token'

Os tipos de token representam tokens, ou seja, valores opacos produzidos e consumidos por algumas operações. Os tokens são usados para impor uma ordem de execução nas operações, conforme descrito na seção Execução.

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

Os tipos de tuplas representam tuplas, ou seja, listas heterogêneas. As tuplas são um recurso legado que existe apenas para compatibilidade com a HLO. Na HLO, as tuplas são usadas para representar entradas e saídas variadicamente. No StableHLO, as entradas e saídas variadicamente são compatíveis com a ABI HLO, e o único uso de tuplas no StableHLO é representar de forma abrangente a ABI HLO, em que, por exemplo, T, tuple<T> e tuple<tuple<T>> podem ser materialmente diferentes dependendo de uma implementação específica. No futuro, planejamos fazer mudanças na ABI do HLO que podem permitir a remoção de tipos de tupla do 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'

Os tipos de elemento representam elementos de tipos de tensor. Ao contrário de muitas linguagens de programação, esses tipos não são de primeira classe no StableHLO. Isso significa que os programas StableHLO não podem representar diretamente valores desses tipos. Como resultado, é comum representar valores escalares do tipo T com valores de tensor 0-dimensional do tipo tensor<T>.

• O tipo booleano representa valores booleanos true e false.
• Os tipos de inteiro podem ser assinados (si) ou não assinados (ui) e ter uma das larguras de bit aceitas (2, 4, 8, 16, 32 ou 64). Os tipos siN assinados representam valores inteiros de -2^(N-1) a 2^(N-1)-1, e os tipos uiN não assinados representam valores inteiros de 0 a 2^N-1.
• Os tipos de ponto flutuante podem ser um dos seguintes:
  • Números de ponto flutuante de 8 bits f8E3M4, f8E4M3 e f8E5M2 seguindo convenções IEEE-754.
  • Tipos f8E4M3FN e f8E5M2 correspondentes às codificações E4M3 e E5M2 do formato FP8 descritos em Formatos do FP8 para aprendizado profundo.
  • Os tipos f8E4M3FNUZ e f8E5M2FNUZ correspondem às codificações E4M3 e E5M2 dos formatos FP8 descritos em Formatos numéricos de 8 bits para redes neurais profundas.
  • Tipo f8E4M3B11FNUZ correspondente à codificação E4M3 dos formatos FP8 descritos no Treinamento e inferência de ponto flutuante de 8 bits híbridos (HFP8) para redes neurais profundas.
  • O tipo bf16 corresponde ao formato bfloat16 descrito em BFloat16: o segredo do alto desempenho em Cloud TPUs.
  • Os tipos f16, f32 e f64 correspondem, respectivamente, aos formatos binary16 ("meia precisão"), binary32 ("precisão simples") e binary64 ("precisão dupla") descritos no padrão IEEE 754.
  • O tipo tf32 corresponde ao formato TensorFloat32 e tem suporte limitado no StableHLO.
  • Tipos f4E2M1FN, f6E2M3FN, f6E3M2FN e f8E8M0FNU MX (microescalonamento) descritos na Especificação de formatos de microescalonamento do OCP.
• Tipos complexos representam valores complexos que têm uma parte real e uma parte imaginária do mesmo tipo de elemento. Os tipos complexos aceitos são complex<f32> (ambas as partes são do tipo f32) e complex<f64> (ambas as partes são do tipo f64).

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

Os tipos de função representam funções nomeadas e anônimas. Eles têm tipos de entrada (a lista de tipos no lado esquerdo de ->) e tipos de saída (a lista de tipos no lado direito de ->). Em muitos idiomas de programação, os tipos de função são de primeira classe, mas não no StableHLO.

StringType ::= 'string'

O tipo de string representa sequências de bytes. Ao contrário de muitas linguagens de programação, o tipo de string não é de primeira classe no StableHLO e é usado apenas para especificar metadados estáticos para elementos do programa.

Operações

As operações StableHLO (também chamadas de ops) representam um conjunto fechado de operações de alto nível em modelos de aprendizado de máquina. Como discutido acima, a sintaxe do StableHLO é muito inspirada no MLIR, que não é necessariamente a alternativa mais ergonômica, mas é possivelmente a melhor opção para o objetivo do StableHLO de criar mais interoperabilidade entre frameworks e compiladores de ML.

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

As operações do StableHLO (também chamadas de ops) têm um nome, entradas/saídas e uma assinatura. O nome consiste no prefixo stablehlo. e em uma mnemônica que identifica exclusivamente uma das operações com suporte. Confira abaixo uma lista completa de todas as operações com suporte.

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

As operações consomem entradas e produzem saídas. As entradas são categorizadas em valores de entrada (calculados durante a execução), funções de entrada (fornecidas estaticamente, porque nas funções StableHLO não são valores de primeira classe) e atributos de entrada (também fornecidos de forma estática). O tipo de entradas e saídas consumido e produzido por uma operação depende do mnemônico. Por exemplo, a operação add consome dois valores de entrada e produz um valor de saída. Em comparação, a operação select_and_scatter consome três valores de entrada, duas funções de entrada e três atributos de entrada.

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

As funções de entrada (também chamadas de funções anônimas) são muito semelhantes às funções nomeadas, exceto que: 1) elas não têm um identificador (daí o nome "anônimo"), 2) elas não declaram tipos de saída (os tipos de saída são inferidos da operação return na função).

A sintaxe para funções de entrada inclui uma parte atualmente não utilizada (consulte a produção Unused acima), que está presente para compatibilidade com o MLIR. No MLIR, há um conceito mais geral de "regiões", que podem ter vários "blocos" de operações conectados por operações de salto. Esses blocos têm IDs que correspondem à produção de Unused para que possam ser diferenciados. O StableHLO não tem operações de salto. Portanto, a parte correspondente da sintaxe do MLIR não é usada, mas ainda está lá.

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

Os atributos de entrada têm um nome e um valor, que é uma das constantes compatíveis. Eles são a principal maneira de especificar metadados estáticos para elementos de programas. Por exemplo, a operação concatenate usa o atributo dimension para especificar a dimensão em que os valores de entrada são concatenados. Da mesma forma, a operação slice usa vários atributos, como start_indices e limit_indices, para especificar os limites usados para dividir o valor de entrada.

No momento, os programas StableHLO em uso às vezes contêm atributos que não são descritos neste documento. No futuro, planejamos absorver esses atributos no opset StableHLO ou impedir que eles apareçam nos programas StableHLO. Enquanto isso, confira a lista desses atributos:

• layout (#629).
• mhlo.frontend_attributes (628, link em inglês).
• mhlo.sharding (#619).
• output_operand_aliases (#740).
• Metadados de local (#594).

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

A assinatura de operação consiste nos tipos de todos os valores de entrada (a lista de tipos no lado esquerdo de ->) e nos tipos de todos os valores de saída (a lista de tipos no lado direito de ->). Estritamente falando, os tipos de entrada são redundantes, e os tipos de saída quase sempre também são redundantes, porque, para a maioria das operações StableHLO, os tipos de saída podem ser inferidos das entradas. No entanto, a assinatura de operação faz parte deliberadamente da sintaxe StableHLO para compatibilidade com o MLIR.

Confira abaixo um exemplo de operação com mnemônico select_and_scatter. Ela consome três valores de entrada (%operand, %source e %init_value), duas funções de entrada e três atributos de entrada (window_dimensions, window_strides e padding). Observe como a assinatura da operação inclui apenas os tipos dos valores de entrada, mas não os tipos de funções e atributos de entrada fornecidos 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>

Constantes

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

As constantes StableHLO têm um literal e um tipo que, juntos, representam um valor de StableHLO. Geralmente, o tipo faz parte da sintaxe da constante, exceto quando não é ambíguo (por exemplo, uma constante booleana inequivocamente tem o tipo i1, enquanto uma constante de número inteiro pode ter vários tipos possíveis).

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

As constantes booleanas representam valores booleanos true e false. As constantes booleanas têm o tipo i1.

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

As constantes de números inteiros representam valores inteiros por strings que usam notação decimal ou hexadecimal. Outras bases, como binária ou octal, não são aceitas. As constantes de números inteiros têm as seguintes restrições:

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

As constantes de ponto flutuante representam valores de ponto flutuante por strings que usam notação decimal ou científica. Além disso, a notação hexadecimal pode ser usada para especificar diretamente os bits subjacentes no formato de ponto flutuante do tipo correspondente. As constantes de ponto flutuante têm as seguintes restrições:

• (C1) Se uma notação não hexadecimal for usada, is_wellformed(float_literal, float_type).
• (C2) Se a notação hexadecimal for usada, size(hexadecimal_digits) = num_bits(float_type) / 4.

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

As constantes complexas representam valores complexos usando listas de uma parte real (primeiro) e uma parte imaginária (segunda). Por exemplo, (1.0, 0.0) : complex<f32> representa 1.0 + 0.0i e (0.0, 1.0) : complex<f32> representa 0.0 + 1.0i. A ordem em que essas partes são armazenadas na memória é definida pela implementação. As constantes complexas têm as seguintes restrições:

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

As constantes de tensor representam valores de tensor usando listas aninhadas especificadas por meio da notação NumPy. Por exemplo, dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> representa um valor de tensor com o seguinte mapeamento de índices para elementos: {0, 0} => 1, {0, 1} => 2, {0, 2} => 3, {1, 0} => 4, {1, 1} => 5, {1, 2} => 6. A ordem em que esses elementos são armazenados na memória é definida pela implementação. As constantes de tensor têm as seguintes restrições:

• (C1) has_syntax(tensor_literal, element_type(tensor_type)), em que:
  • 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)), em que:
  • 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:]).
  • caso contrário, false.

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

Constantes de tensor quantizadas representam valores de tensor quantizados usando a mesma notação que as constantes de tensor, com elementos especificados como constantes do tipo de armazenamento. As constantes de tensor quantizadas têm as seguintes restrições:

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

Os literals de string consistem em bytes especificados usando caracteres ASCII e sequências de escape. Eles não dependem da codificação, portanto, a interpretação desses bytes é definida pela implementação. Os literais de strings têm o tipo string.

Operações

abs

Semântica

Executa a operação de elemento absoluto no tensor operand e produz um tensor result. Dependendo do tipo de elemento, faça o seguinte:

• Para números inteiros com sinal: módulo de número inteiro.
• Para números flutuantes: abs do IEEE-754.
• Para números complexos: módulo complexo.
• Para tipos quantizados: dequantize_op_quantize(abs, operand, type(result)).

Entradas

Saídas

Restrições

• (C1) shape(result) = shape(operand).
• (C2) baseline_element_type(result) é definido como:
  • complex_element_type(element_type(operand)) se is_complex(operand).
  • baseline_element_type(operand) se não forem.

Exemplos

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

Mais exemplos

adicionar

Semântica

Realiza a adição elementar de dois tensores lhs e rhs e produz um tensor result. Dependendo do tipo de elemento, faça o seguinte:

• Para booleanos: OR lógico.
• Para números inteiros: adição de números inteiros.
• Para pontos flutuantes: addition de IEEE-754.
• Para números complexos: adição complexa.
• Para tipos quantizados: dequantize_op_quantize(add, lhs, rhs, type(result)).

Entradas

Saídas

Restrições

• Se a operação usar tensores não quantizados:
  • (C1) type(lhs) = type(rhs) = type(result).
• Se a operação usa tensores quantizados:
  • (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), então quantization_dimension(lhs) = quantization_dimension(result).
  • (C7) Se is_per_axis_quantized(rhs), então quantization_dimension(rhs) = quantization_dimension(result).

Exemplos

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

 Mais exemplos

after_all

Semântica

Garante que as operações que produzem o inputs sejam executadas antes de qualquer operação que dependa de result. A execução dessa operação não faz nada. Ela existe apenas para estabelecer dependências de dados de result a inputs.

Entradas

Saídas

Exemplos

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

Mais exemplos

all_gather

Semântica

Em cada grupo de processos na grade de processos do StableHLO, concatena os valores dos tensores operands de cada processo junto a all_gather_dim e produz tensores results.

A operação divide a grade do processo StableHLO em process_groups, que é definida da seguinte maneira:

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

Em seguida, em cada process_group:

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

Entradas

Saídas

Restrições

• (C1) 0 <= all_gather_dim < rank(operands...).
• (C2) is_unique(replica_groups).
• (C3) size(replica_groups) é definido como:
  • num_replicas se cross_replica for usado.
  • num_replicas se cross_replica_and_partition for usado.
  • num_processes se flattened_ids for usado.
• (C4) 0 <= replica_groups < size(replica_groups).
• (C5) Se use_global_device_ids = true, então channel_id > 0.
• (C6) type(results...) = type(operands...), exceto:
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1).

Exemplos

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

 Mais exemplos

all_reduce

Semântica

Em cada grupo de processos na grade de processo StableHLO, é aplicada uma função de redução computation aos valores dos tensores operands de cada processo e são gerados tensores results.

A operação divide a grade de processo do StableHLO em process_groups, que é definida da seguinte maneira:

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

Em seguida, em cada process_group:

• results...@process[result_index] = exec(schedule) para alguma árvore binária schedule em que:
  • exec(node) = computation(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value
• schedule é uma árvore binária definida pela implementação, cuja travessia em ordem é to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0])).

Entradas

Saídas

Restrições

• (C1) is_unique(replica_groups).
• (C2) size(replica_groups) é definido como:
  • num_replicas se cross_replica for usado.
  • num_replicas se cross_replica_and_partition for usado.
  • num_processes se flattened_ids for usado.
• (C3) 0 <= replica_groups < size(replica_groups).
• (C4) Se use_global_device_ids = true, então channel_id > 0.
• (C5) computation tem o tipo (tensor<E>, tensor<E>) -> (tensor<E>) em que is_promotable(element_type(operand), E).
• (C6) shape(results...) = shape(operands...).
• (C7) element_type(results...) = E.

Exemplos

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

 Mais exemplos

all_to_all

Semântica

Em cada grupo de processos na grade de processo StableHLO, divide os valores dos tensores operands ao longo de split_dimension em partes, espalha as partes divididas entre os processos, concatena as partes espalhadas ao longo de concat_dimension e produz tensores results. A operação divide a grade do processo StableHLO em process_groups, que é definida da seguinte maneira:

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

Em seguida, em cada process_group:

• split_parts...@sender = split(operands...@sender, split_count, split_dimension) para todos os sender em process_group.
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group] em que receiver_index = process_group.index(receiver).
• results...@process = concatenate(scattered_parts...@process, concat_dimension).

Entradas

Saídas

Restrições

• (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) é definido como:
  • num_replicas se cross_replica for usado.
  • num_partitions se cross_partition for usado.
• (C7) 0 <= replica_groups < size(replica_groups).
• (C8) dim(replica_groups, 1) = split_count.
• (C9) type(results...) = type(operands...), exceto 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.

Exemplos

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

 Mais exemplos

e

Semântica

Realiza a operação AND de dois tensores lhs e rhs e produz um tensor result. Dependendo do tipo de elemento, faz o seguinte:

• Para booleanos: AND lógico.
• Para números inteiros: E bit a bit.

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

atan2

Semântica

Executa a operação atan2 com elementos nos tensores lhs e rhs e produz um tensor result. Dependendo do tipo de elemento, faz o seguinte:

• Para números flutuantes: atan2 do IEEE-754.
• Para números complexos: atan2 complexo.
• Para tipos quantizados: dequantize_op_quantize(atan2, lhs, rhs, type(result)).

Entradas

Saídas

Restrições

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

Exemplos

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

Mais exemplos

batch_norm_grad

Semântica

Calcula gradientes de várias entradas de retropropagação batch_norm_training de grad_output e produz tensores grad_operand, grad_scale e grad_offset. Mais formalmente, essa operação pode ser expressa como uma decomposição para operações StableHLO existentes usando a sintaxe do Python da seguinte maneira:

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

Para tipos quantizados, executa 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)).

Entradas

Saídas

Restrições

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, mean, variance, grad_output, grad_operand, grad_scale e grad_offset têm o mesmo baseline_element_type.
• (C3) operand, grad_output e grad_operand têm a mesma forma.
• (C4) scale, mean, variance, grad_scale e grad_offset têm a mesma forma.
• (C5) size(scale) = dim(operand, feature_index).

Exemplos

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

Semântica

Normaliza o tensor operand em todas as dimensões, exceto a feature_index, e produz um tensor result. De forma mais formal, essa operação pode ser expressa como uma decomposição para operações StableHLO existentes usando a sintaxe do Python da seguinte maneira:

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)

Para tipos quantizados, executa 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)).

Entradas

Saídas

Restrições

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, mean, variance e result têm o mesmo 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).

Exemplos

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

Semântica

Calcula a média e a variância em todas as dimensões, exceto a dimensão feature_index, e normaliza o tensor operand, produzindo os tensores output, batch_mean e batch_var. De forma mais formal, essa operação pode ser expressa como uma decomposição para operações StableHLO existentes usando a sintaxe do Python da seguinte forma:

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

Para tipos quantizados, executa 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)).

Entradas

Saídas

Restrições

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, batch_mean, batch_var e output têm o mesmo 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).

Exemplos

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

Semântica

Realiza uma operação de bitcast no tensor operand e produz um tensor result em que os bits de todo o tensor operand são reinterpretados usando o tipo do tensor result.

Mais formalmente, considerando 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 retorna a representação na memória de um determinado valor, e o comportamento é definido pela implementação porque a representação exata de tensores é definida pela implementação, e a representação exata de tipos de elementos também é definida pela implementação.

Entradas

Saídas

Restrições

• (C1) Considerando 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) para todos os 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) para todos os 0 <= i < R.
  • dim(operand, R - 1) * num_bits(E) = num_bits(E').
• (C2) Se is_complex(operand) or is_complex(result), então is_complex(operand) and is_complex(result).

Exemplos

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

 Mais exemplos

broadcast_in_dim

Semântica

Expande as dimensões e/ou a classificação de um tensor de entrada duplicando os dados no tensor operand e produz um tensor result. Mais formalmente, result[result_index] = operand[operand_index], em que para todos os d em axes(operand):

• operand_index[d] = 0 se dim(operand, d) = 1.
• operand_index[d] = result_index[broadcast_dimensions[d]] se não forem.

Entradas

Saídas

Restrições

• (C1) element_type(result) é dado por:
  • element_type(operand), se !is_per_axis_quantized(operand).
  • element_type(operand), exceto que quantization_dimension(operand), scales(operand) e zero_points(operand) podem ser diferentes de quantization_dimension(result), scales(result) e zero_points(result) resp., caso contrário.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Para todos os d em axes(operand):
  • dim(operand, d) = 1 ou
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Se is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Se for dim(operand, quantization_dimension(operand)) = 1, então scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))).

Exemplos

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

 Mais exemplos

caso

Semântica

Produz a saída da execução de exatamente uma função de branches, dependendo do valor de index. Mais formalmente, result = selected_branch() em que:

• selected_branch = branches[index] se 0 <= index < size(branches).
• Caso contrário, selected_branch = branches[-1].

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

BRT

Semântica

Executa uma operação raiz cúbica elemento a elemento no tensor operand e produz um tensor result. Dependendo do tipo de elemento, faça o seguinte:

• Para pontos flutuantes: rootn(x, 3) de IEEE-754.
• Para números complexos: raiz cúbica complexa.
• Para tipos quantizados: dequantize_op_quantize(cbrt, operand, type(result))

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

ceil

Semântica

Realiza o teto elementar do tensor operand e produz um tensor result. Implementa a operação roundToIntegralTowardPositive da especificação IEEE-754. Para tipos quantizados, executa dequantize_op_quantize(ceil, operand, type(result)).

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

Cholesky

Semântica

Calcula a decomposição de Cholesky de um lote de matrizes.

Mais formalmente, para todos os i em index_space(result), result[i0, ..., iR-3, :, :] é uma decomposição de Cholesky de a[i0, ..., iR-3, :, :], na forma de uma matriz triangular inferior (se lower for true) ou triangular superior (se lower for false). Os valores de saída no triângulo oposto, ou seja, o triângulo superior ou inferior, são definidos pela implementação.

Se houver i em que a matriz de entrada não for uma matriz definida positiva hermitiana, o comportamento será indefinido.

Para tipos quantizados, executa dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)).

Entradas

Saídas

Restrições

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

Exemplos

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

limitar

Semântica

Limita cada elemento do tensor operand entre um valor mínimo e máximo e produz um tensor result. Mais formalmente, result[result_index] = minimum(maximum(operand[result_index], min_element), max_element), em que min_element = rank(min) = 0 ? min[] : min[result_index], max_element = rank(max) = 0 ? max[] : max[result_index]. Para tipos quantizados, realiza dequantize_op_quantize(clamp, min, operand, max, type(result)).

Impor uma ordenação em números complexos envolve semântica surpreendente. Por isso, planejamos remover o suporte a números complexos para essa operação (560).

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

collective_broadcast

Semântica

Em cada grupo de processos na grade de processo StableHLO, envie o valor do tensor operand do processo de origem para os processos de destino e produza um tensor result.

A operação divide a grade de processo do StableHLO em process_groups, que é definida da seguinte maneira:

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

Depois, result@process é dado por:

• operand@process_groups[i, 0] se houver um i de modo que o processo esteja em process_groups[i].
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) se não forem.

Entradas

Saídas

Restrições

• (C1) is_unique(replica_groups).
• (C2) 0 <= replica_groups < N, em que N é definido como:
  • num_replicas se cross_replica for usado.
  • num_partitions se cross_partition for usado.
• (C3) type(result) = type(operand).

Exemplos

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

Semântica

Em cada grupo de processos na grade de processo StableHLO, envia o valor do tensor operand do processo de origem para o processo de destino e produz um tensor result.

A operação divide a grade do processo StableHLO em process_groups, que é definida da seguinte maneira:

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

Depois, result@process é dado por:

• operand@process_groups[i, 0], se houver um i tal que process_groups[i, 1] = process.
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) caso contrário.

Entradas

Saídas

Restrições

• (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, em que N é definido como:
  • num_replicas se cross_replica for usado.
  • num_partitions se cross_partition for usado.
• (C5) type(result) = type(operand).

Exemplos

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

 Mais exemplos

compare

Semântica

Realiza a comparação elemento a elemento dos tensores lhs e rhs de acordo com comparison_direction e compare_type e produz um tensor result.

Os valores de comparison_direction e compare_type têm as seguintes semânticas:

Para tipos de elementos booleanos e inteiros:

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

Para tipos de elementos de ponto flutuante com compare_type = FLOAT, a operação implementa as seguintes operações IEEE-754:

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

Para tipos de elementos de ponto flutuante com compare_type = TOTALORDER, a operação usa a combinação de operações totalOrder e compareQuietEqual do IEEE-754.

Para tipos de elementos complexos, a comparação lexicográfica de pares (real, imag) é realizada usando o comparison_direction e o compare_type fornecidos. Impor uma ordenação em números complexos envolve semântica surpreendente. Por isso, planejamos remover o suporte a números complexos quando comparison_direction for GE, GT, LE ou LT (#560).

Para tipos quantizados, executa dequantize_compare(lhs, rhs, comparison_direction).

Entradas

Saídas

Restrições

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs).
• (C2) shape(lhs) = shape(rhs) = shape(result).
• (C3) compare_type é definido como:
  • SIGNED se is_signed_integer(element_type(lhs)).
  • UNSIGNED se is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)).
  • FLOAT ou TOTALORDER se is_float(element_type(lhs)).
  • FLOAT se is_complex(element_type(lhs)).

Exemplos

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

 Mais exemplos

complexo

Semântica

Executa a conversão por elemento para um valor complexo de um par de valores reais e imaginários, lhs e rhs, e produz um tensor result.

Entradas

Saídas

Restrições

• (C1) type(lhs) = type(rhs).
• (C2) shape(result) = shape(lhs).
• (C3) element_type(result) tem o tipo complex<E> em que E = element_type(lhs).

Exemplos

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

Mais exemplos

composto

Semântica

Encapsula uma operação composta (composta) de outras operações do StableHLO, usando inputs e composite_attributes e produzindo results. A semântica da operação é implementada pelo atributo decomposition. A operação composite pode ser substituída pela decomposição sem alterar a semântica do programa. Nos casos em que a inserção da decomposição não fornece a mesma semântica de operação, prefira usar custom_call.

O campo version (padrão 0) é usado para indicar quando a semântica de um composto muda.

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

concatenate

Semântica

Concatena inputs ao longo da dimensão dimension na mesma ordem dos argumentos fornecidos e produz um tensor result. Mais formalmente, result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1], em que:

1. id = d0 + ... + dk-1 + kd.
2. d é igual a dimension, e d0, ... são os tamanhos da dª dimensão de inputs.

Entradas

Saídas

Restrições

• (C1) same(element_type(inputs...)).
• (C2) same(shape(inputs...)), exceto 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]), exceto:
  • dim(result, dimension) = dim(inputs[0], dimension) + ....

Exemplos

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

Mais exemplos

constante

Semântica

Produz um tensor output de uma value constante.

Entradas

Saídas

Restrições

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

Exemplos

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

Mais exemplos

fazer uma conversão

Semântica

Realiza uma conversão elemento a elemento de um tipo de elemento para outro no tensor operand e produz um tensor result.

Para conversões de booleano para qualquer tipo com suporte, o valor false é convertido em zero, e o valor true é convertido em um. Para conversões any-supported-type-to-boolean, um valor zero é convertido em false, e os valores diferentes de zero são convertidos em true. Confira abaixo como isso funciona para tipos complexos.

Para conversões que envolvem número inteiro para número inteiro, número inteiro para ponto flutuante ou ponto flutuante para ponto flutuante, se o valor de origem puder ser representado exatamente no tipo de destino, o valor do resultado será essa representação exata. Caso contrário, o comportamento será a ser definido (#180).

Para conversões que envolvem ponto flutuante para número inteiro, a parte fracionária é truncada. Se o valor truncado não puder ser representado no tipo de destino, o comportamento será a definir (#180).

A conversão que envolve complexo para complexo segue o mesmo comportamento das conversões de ponto flutuante para ponto flutuante para converter partes reais e imaginárias.

Para conversões complexo-para-qualquer-outro-tipo e qualquer-outro-tipo-para-complexo, o valor imaginário de origem é ignorado ou o valor imaginário de destino é zerado, respectivamente. A conversão da parte real segue as conversões de ponto flutuante.

Em princípio, essa operação pode expressar desquantização (conversão de tensores quânticos para tensores regulares), quantização (conversão de tensores regulares para tensores quânticos) e requantização (conversão entre tensores quânticos), mas no momento temos operações dedicadas para isso: uniform_dequantize para o primeiro caso de uso e uniform_quantize para o segundo e o terceiro casos de uso. No futuro, essas duas operações poderão ser mescladas em convert (#1576).

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

convolução

Semântica

Calcula produtos pontuais entre janelas de lhs e frações de rhs e produz result. O diagrama a seguir mostra como os elementos em result são calculados a partir de lhs e rhs usando um exemplo concreto.

Mais formalmente, considere o seguinte reformulação das entradas em termos de lhs para poder expressar janelas de 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).

Essa reformulação usa as seguintes funções auxiliares:

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

Se feature_group_count = 1 e batch_group_count = 1, para todos output_spatial_index em index_space(dim(result, output_spatial_dimensions...)), result[result_shape(:, output_spatial_index, :)] = dot_product em que:

• 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]). Esse recurso parece não ser usado, então planejamos removê-lo no futuro (#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).

Para tipos quantizados, executa 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)).

Para tipos quantizados híbridos, executa 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).

Entradas

Saídas

Restrições

• (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) Dado o 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) Dado 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) Dado 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) é definido como:
  • 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, em que:
  • 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 a operação usar tensores não quantizados:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Se a operação usa tensores quantizados:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Se is_per_axis_quantized(rhs), então quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Se is_per_axis_quantized(result), então 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), então is_per_tensor_quantized(result).
  • Se !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Exemplos

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

 Mais exemplos

cosseno

Semântica

Executa a operação de cosseno elemento a elemento no tensor operand e produz um tensor result. Dependendo do tipo de elemento, faça o seguinte:

• Para pontos flutuantes: cos de IEEE-754.
• Para números complexos: cosseno complexo.
• Para tipos quantizados: dequantize_op_quantize(cosine, operand, type(result)).

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

count_leading_zeros

Semântica

Executa a contagem por elemento do número de bits zero à esquerda no tensor operand e produz um tensor result.

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

custom_call

Semântica

Encapsula uma operação call_target_name definida pela implementação que usa inputs e called_computations e produz results. has_side_effect, backend_config e api_version podem ser usados para fornecer mais metadados definidos pela implementação.

No momento, essa operação contém uma coleção bastante desorganizada de metadados que reflete a evolução orgânica da operação de contraparte no compilador XLA. No futuro, planejamos unificar esses metadados (#741).

Entradas

Saídas

Exemplos

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

dividir

Semântica

Realiza a divisão elementar dos tensores lhs e rhs do dividendo e produz um tensor result. Dependendo do tipo de elemento, faz o seguinte:

• Para números inteiros: divisão de números inteiros que produz o quociente algébrico com qualquer parte fracionária descartada.
• Para pontos flutuantes: division de IEEE-754.
• Para números complexos: divisão complexa.
• Para tipos quantizados:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)).

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

dot_general

Semântica

Calcula produtos escalares entre fatias de lhs e fatias de rhs e produz um tensor result.

Mais formalmente, result[result_index] = dot_product, em que:

• 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, em que 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)).

Para tipos quantizados, executa 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)).

Para tipos quantizados híbridos, executa 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 controla o equilíbrio entre velocidade e precisão para cálculos em back-ends de aceleradores. Pode ser um dos seguintes (no momento, a semântica desses valores de tipo enumerado está sub-especificada, mas planejamos resolver isso em #755):

• DEFAULT: cálculo mais rápido, mas aproximação menos precisa do número original.
• HIGH: cálculo mais lento, mas aproximação mais precisa do número original.
• HIGHEST: cálculo mais lento, mas a aproximação mais precisa do número original.

Um DotAlgorithm define as principais propriedades do algoritmo usado para implementar a operação de ponto, que também define a precisão. Se os campos de atributo do algoritmo estiverem definidos, o precision_config precisará ser DEFAULT. DotAlgorithms não têm um valor padrão, já que os parâmetros padrão são definidos pela implementação. Assim, todos os campos do algoritmo de pontos podem ser definidos como None para especificar um algoritmo de pontos vazio, que usará o valor precision_config.

Os campos DotAlgorithm incluem:

• lhs_precision_type e rhs_precision_type, as precisões para as quais o LHS e o RHS da operação são arredondados. Os tipos de precisão são independentes dos tipos de armazenamento das entradas e da saída.
• accumulation_type a precisão usada para acúmulo.
• lhs_component_count, rhs_component_count e num_primitive_operations são usados quando estamos fazendo um algoritmo que decompõe o LHS e/ou RHS em vários componentes e faz várias operações de ponto "primitivas" nesses valores, geralmente para emular uma precisão maior (por exemplo, Uso do tipo de dados de inteligência artificial bfloat16 para cálculos de maior precisão: bf16_6x tf32_3x etc.). Para algoritmos sem decomposição, esses valores precisam ser definidos como 1.
• allow_imprecise_accumulation para especificar se o acúmulo com precisão menor é permitido em algumas etapas (por exemplo, CUBLASLT_MATMUL_DESC_FAST_ACCUM).

Exemplos de atributos DotAlgorithm:

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


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


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

Cabe às implementações decidir quais combinações são compatíveis. Em geral, não é garantido que cada algoritmo seja compatível com cada tipo de acelerador pelo consumidor do StableHLO. Se um determinado algoritmo não tiver suporte, um erro será gerado em vez de retornar a uma alternativa. A verificação StableHLO vai oferecer a melhor verificação possível, impedindo que algoritmos que não são conhecidos sejam aceitos em qualquer hardware.

Consulte xla_data.proto > Algorithm para conferir alguns valores de algoritmos compatíveis. O tíquete no 2483 captura o plano para criar um documento centralizado em algoritmos compatíveis por back-end.

Entradas

Saídas

Restrições

• (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 a operação usar tensores não quantizados:
  • (C13) element_type(lhs) = element_type(rhs).
• Se a operação usa tensores quantizados:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C15) zero_points(rhs) = 0.
  • (C16) Se is_per_axis_quantized(rhs), então quantization_dimension(rhs) não está em 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), então 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.

Exemplos

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

 Mais exemplos

dynamic_broadcast_in_dim

Semântica

Essa operação é funcionalmente idêntica à operação broadcast_in_dim, mas a forma do resultado é especificada dinamicamente por output_dimensions.

A operação também aceita atributos opcionais known_expanding_dimensions, known_nonexpanding_dimensions para expressar conhecimento estático sobre o comportamento de expansão das dimensões. Se não for especificado, todas as dimensões vão ser expandidas.

Entradas

Saídas

Restrições

• (C1) element_type(result) é dado por:
  • element_type(operand), se !is_per_axis_quantized(operand).
  • element_type(operand), exceto que quantization_dimension(operand), scales(operand) e zero_points(operand) podem ser diferentes de quantization_dimension(result), scales(result) e zero_points(result) resp., caso contrário.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Para todos os d em axes(operand):
  • dim(operand, d) = 1 ou
  • 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, então 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).

Exemplos

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

Mais exemplos

dynamic_conv

Semântica

Essa operação é funcionalmente idêntica à convolução, mas o padding é especificado dinamicamente por padding.

Entradas

Saídas

Restrições

• (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) Dado o 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) Dado 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) Dado 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) é definido como:
  • 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, em que:
  • 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 a operação usar tensores não quantizados:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Se a operação usa tensores quantizados:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Se is_per_axis_quantized(rhs), então quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Se is_per_axis_quantized(result), então 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), então is_per_tensor_quantized(result).
  • Se !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Exemplos

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