Especificación de StableHLO

StableHLO es un conjunto de operaciones para operaciones de alto nivel (HLO) en modelos de aprendizaje automático (AA). StableHLO funciona como una capa de portabilidad entre diferentes compiladores y frameworks de AA: los frameworks de AA que producen programas de StableHLO son compatibles con los compiladores de AA que consumen programas de StableHLO.

Nuestro objetivo es simplificar y acelerar el desarrollo del AA mediante la creación de más interoperabilidad entre varios frameworks de AA (como TensorFlow, JAX y PyTorch) y compiladores de AA (como IREE y XLA). Para ello, en este documento, se proporciona una especificación para el lenguaje de programación StableHLO.

Esta especificación contiene tres secciones principales. En primer lugar, en la sección Programs, se describe la estructura de los programas de StableHLO, que consisten en funciones de StableHLO que, a su vez, consisten en operaciones de StableHLO. Dentro de esa estructura, la sección Ops especifica la semántica de las operaciones individuales. La sección Ejecución proporciona semántica para todas estas operaciones que se ejecutan juntas dentro de un programa. Por último, en la sección Notation, se analiza la notación que se usa en toda la especificación.

Para ver la especificación de una versión anterior de StableHLO, abre el repositorio en la versión etiquetada de interés. Por ejemplo, la especificación de StableHLO v0.19.0. Para ver los cambios que se produjeron en cada aumento de versión menor de StableHLO, consulta el registro de versiones en VhloDialect.td.

Programas

Program ::= {Func}

Los programas de StableHLO consisten en una cantidad arbitraria de funciones de StableHLO. A continuación, se muestra un programa de ejemplo con una función @main que tiene 3 entradas (%image, %weights y %bias) y 1 salida. El cuerpo de la función tiene 6 operaciones.

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

Funciones

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

Las funciones StableHLO (que también se denominan funciones con nombre) tienen un identificador, entradas/salidas y un cuerpo. En el futuro, planeamos introducir metadatos adicionales para las funciones para lograr una mejor compatibilidad con HLO (#425, #626, #740 y #744).

Identificadores

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

Los identificadores de StableHLO son similares a los identificadores de muchos lenguajes de programación, con dos peculiaridades: 1) todos los identificadores tienen símbolos que distinguen diferentes tipos de identificadores y 2) los identificadores de valor pueden ser completamente numéricos para simplificar la generación de programas de StableHLO.

Tipos

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

Los tipos de StableHLO se clasifican en tipos de valor (que también se denominan tipos de primera clase) que representan valores de StableHLO y tipos que no son de valor que describen otros elementos del programa. Los tipos de StableHLO son similares a los tipos de muchos lenguajes de programación, y la principal peculiaridad es la naturaleza específica del dominio de StableHLO, que genera algunos resultados inusuales (p. ej., los tipos escalares no son tipos de valor).

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

Los tipos de tensores representan tensores, es decir, arreglos multidimensionales. Tienen una forma y un tipo de elemento, en el que una forma representa tamaños de dimensión no negativos o desconocidos en el orden ascendente de las dimensiones correspondientes (que también se denominan ejes) numeradas de 0 a R-1. La cantidad de dimensiones R se denomina rango. Por ejemplo, tensor<2x3xf32> es un tipo de tensor con la forma 2x3 y el tipo de elemento f32. Tiene dos dimensiones (o, en otras palabras, dos ejes): la dimensión 0 y la dimensión 1, cuyos tamaños son 2 y 3. Su clasificación es 2.

Las formas pueden ser parcialmente o completamente desconocidas (dinámicas), p. ej., tensor<?x2xf64> es parcialmente desconocida y tensor<?x?xf64> es completamente desconocida. Los tamaños de las dimensiones dinámicas se representan con un ?. Las formas no pueden tener una clasificación.

En el futuro, planeamos explorar la extensión de los tipos de tensores más allá de los tamaños de dimensión y los tipos de elementos, por ejemplo, para incluir diseños (#629) y dispersión (#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

Los tipos de elementos cuantificados representan valores enteros de un tipo de almacenamiento en el rango de storage_min a storage_max (inclusive) que corresponden a valores de punto flotante de un tipo expresado. Para un valor de número entero i determinado, el valor de punto flotante correspondiente f se puede calcular como f = (i - zero_point) * scale, en el que scale y zero_point se denominan parámetros de cuantización. storage_min y storage_max son opcionales en la gramática, pero tienen valores predeterminados de min_value(storage_type) y max_value(storage_type), respectivamente. Los tipos de elementos cuantificados tienen las siguientes restricciones:

• (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) Si is_empty(quantization_dimension), entonces size(scales) = 1.
• (C11) 0 <= quantization_dimension.

Por el momento, QuantizationScale es una constante de punto flotante, pero hay un gran interés en las escalas basadas en números enteros, representadas con multiplicadores y cambios. Planeamos explorar esto en un futuro cercano (#1404).

Hay un debate en curso sobre la semántica de QuantizationZeroPoint, incluidos el tipo, los valores y si puede haber uno o potencialmente varios puntos cero en un tipo de tensor cuantizado. En función de los resultados de este debate, la especificación en torno a cero puntos podría cambiar en el futuro (#1405).

Otro análisis en curso involucra la semántica de QuantizationStorageMin y QuantizationStorageMax para determinar si se deben imponer restricciones en estos valores y en los valores de los tensores cuantificados (#1406).

Por último, planeamos explorar la representación de escalas y puntos cero desconocidos, de manera similar a como planeamos explorar la representación de tamaños de dimensión desconocidos (#1407).

Los tipos de tensores cuantificados representan tensores con elementos cuantificados. Estos tensores son exactamente los mismos que los regulares, con la excepción de que sus elementos tienen tipos de elementos cuantificados, en lugar de tipos de elementos regulares.

En los tensores cuantificados, la cuantización puede ser por tensor, es decir, tener un scale y un zero_point para todo el tensor o puede ser por eje, es decir, tener varios scales y zero_points, un par por fragmento de una dimensión quantization_dimension particular. De forma más formal, en un tensor t con cuantificación por eje, hay dim(t, quantization_dimension) rebanadas de quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :], etc. Todos los elementos de la iª rebanada usan scales[i] y zero_points[i] como sus parámetros de cuantificación. Los tipos de tensores cuantificados tienen las siguientes restricciones:

• Para la cuantificación por tensor:
  • Sin restricciones adicionales.
• Para la cuantificación por eje, haz lo siguiente:
  • (C12) quantization_dimension < rank(self).
  • (C13) dim(self, quantization_dimension) = size(scales).

TokenType ::= 'token'

Los tipos de tokens representan tokens, es decir, valores opacos que producen y consumen algunas operaciones. Los tokens se usan para imponer el orden de ejecución en las operaciones, como se describe en la sección Ejecución.

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

Los tipos de tupla representan tuplas, es decir, listas heterogéneas. Las tuplas son una función heredada que solo existe para la compatibilidad con HLO. En HLO, se usan tuplas para representar entradas y salidas variadic. En StableHLO, las entradas y salidas variadic se admiten de forma nativa, y el único uso de tuplas en StableHLO es representar de forma integral la ABI de HLO, en la que, p. ej., T, tuple<T> y tuple<tuple<T>> pueden ser muy diferentes según una implementación en particular. En el futuro, planeamos realizar cambios en la ABI de HLO que podrían permitirnos quitar los tipos de tupla de 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'

Los tipos de elementos representan elementos de tipos de tensores. A diferencia de muchos lenguajes de programación, estos tipos no son de primera clase en StableHLO. Esto significa que los programas de StableHLO no pueden representar directamente valores de estos tipos (como resultado, es idiomático representar valores escalares de tipo T con valores de tensor de 0 dimensiones de tipo tensor<T>).

• El tipo booleano representa los valores booleanos true y false.
• Los tipos de números enteros pueden tener firma (si) o no (ui) y tener uno de los anchos de bits admitidos (2, 4, 8, 16, 32 o 64). Los tipos siN con firma representan valores enteros del -2^(N-1) al 2^(N-1)-1 inclusive, y los tipos uiN sin firma representan valores enteros del 0 al 2^N-1 inclusive.
• Los tipos de números de punto flotante pueden ser uno de los siguientes:
  • f8E3M4, f8E4M3 y f8E5M2 son números de punto flotante de 8 bits que siguen las convenciones de IEEE-754.
  • Tipos f8E4M3FN y f8E5M2 que corresponden, respectivamente, a las codificaciones E4M3 y E5M2 del formato FP8 que se describe en Formatos FP8 para el aprendizaje profundo.
  • Los tipos f8E4M3FNUZ y f8E5M2FNUZ correspondientes a las codificaciones E4M3 y E5M2 de los formatos FP8 descritos en Formatos numéricos de 8 bits para redes neuronales profundas.
  • Tipo f8E4M3B11FNUZ que corresponde a la codificación E4M3 de los formatos FP8 que se describen en Inferencia y entrenamiento híbridos de punto flotante de 8 bits (HFP8) para redes neuronales profundas.
  • Tipo bf16 que corresponde al formato bfloat16 que se describe en BFloat16: El secreto del alto rendimiento en las Cloud TPU.
  • Los tipos f16, f32 y f64 corresponden, respectivamente, a los formatos binary16 ("precisión media"), binary32 ("precisión simple") y binary64 ("precisión doble") descritos en el estándar IEEE 754.
  • El tipo tf32 corresponde al formato TensorFloat32 y tiene compatibilidad limitada en StableHLO.
  • Tipos de MX (microescalamiento) f4E2M1FN, f6E2M3FN, f6E3M2FN y f8E8M0FNU que se describen en la Especificación de formatos de microescalamiento OCP.
• Los tipos complejos representan valores complejos que tienen una parte real y una parte imaginaria del mismo tipo de elemento. Los tipos complejos compatibles son complex<f32> (ambas partes son de tipo f32) y complex<f64> (ambas partes son de tipo f64).

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

Los tipos de funciones representan funciones con nombre y anónimas. Tienen tipos de entrada (la lista de tipos en el lado izquierdo de ->) y tipos de salida (la lista de tipos en el lado derecho de ->). En muchos lenguajes de programación, los tipos de funciones son de primera clase, pero no en StableHLO.

StringType ::= 'string'

El tipo de cadena representa secuencias de bytes. A diferencia de muchos lenguajes de programación, el tipo de cadena no es la primera clase en StableHLO y solo se usa para especificar metadatos estáticos para elementos del programa.

Operaciones

Las operaciones de StableHLO (que también se denominan ops) representan un conjunto cerrado de operaciones de alto nivel en modelos de aprendizaje automático. Como se mencionó anteriormente, la sintaxis de StableHLO se inspira en gran medida en MLIR, que no es necesariamente la alternativa más ergonómica, pero es, sin duda, la más adecuada para el objetivo de StableHLO de crear más interoperabilidad entre los frameworks y los compiladores de AA.

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

Las operaciones StableHLO (que también se llaman ops) tienen un nombre, entradas y salidas, y una firma. El nombre consta del prefijo stablehlo. y una mnemotecnia que identifica de forma exclusiva una de las operaciones admitidas. Consulta a continuación una lista completa de todas las operaciones compatibles.

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

Las operaciones consumen entradas y producen salidas. Las entradas se clasifican en valores de entrada (calculados durante la ejecución), funciones de entrada (proporcionadas de forma estática porque, en StableHLO, las funciones no son valores de primera clase) y atributos de entrada (también se proporcionan de forma estática). El tipo de entradas y salidas que consume y produce una operación depende de su mnemónico. Por ejemplo, la operación add consume 2 valores de entrada y produce 1 valor de salida. En comparación, la operación select_and_scatter consume 3 valores de entrada, 2 funciones de entrada y 3 atributos de entrada.

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

Las funciones de entrada (también llamadas funciones anónimas) son muy similares a las funciones con nombre, excepto que: 1) no tienen un identificador (de ahí el nombre “anónimo”) y 2) no declaran tipos de salida (los tipos de salida se infieren de la operación return dentro de la función).

La sintaxis de las funciones de entrada incluye una parte que no se usa actualmente (consulta la producción de Unused anterior) que está allí para la compatibilidad con MLIR. En MLIR, hay un concepto más general de "regiones" que pueden tener varios "bloques" de operaciones conectados entre sí a través de JumpOps. Estos bloques tienen IDs que corresponden a la producción de Unused, de modo que se puedan distinguir entre sí. StableHLO no tiene operaciones de salto, por lo que la parte correspondiente de la sintaxis de MLIR no se usa (pero todavía está allí).

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

Los atributos de entrada tienen un nombre y un valor que es una de las constantes admitidas. Son la forma principal de especificar metadatos estáticos para los elementos de programa. Por ejemplo, la operación concatenate usa el atributo dimension para especificar la dimensión a lo largo de la cual se concatenan sus valores de entrada. De manera similar, la operación slice usa varios atributos, como start_indices y limit_indices, para especificar los límites que se usan para cortar el valor de entrada.

Por el momento, los programas de StableHLO en el entorno real a veces contienen atributos que no se describen en este documento. En el futuro, planeamos absorber estos atributos en el conjunto de operaciones de StableHLO o prohibir que aparezcan en los programas de StableHLO. Mientras tanto, aquí tienes la lista de estos atributos:

• layout (#629).
• mhlo.frontend_attributes (#628).
• mhlo.sharding (#619).
• output_operand_aliases (#740).
• Metadatos de ubicación (#594).

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

La firma de la operación consta de los tipos de todos los valores de entrada (la lista de tipos en el lado izquierdo de ->) y los tipos de todos los valores de salida (la lista de tipos en el lado derecho de ->). Estrictamente hablando, los tipos de entrada son redundantes y los tipos de salida también suelen serlo (porque, para la mayoría de las operaciones de StableHLO, los tipos de salida se pueden inferir a partir de las entradas). Sin embargo, la firma de op forma parte de la sintaxis de StableHLO de forma deliberada para garantizar la compatibilidad con MLIR.

A continuación, se muestra un ejemplo de operación cuyo mnemónico es select_and_scatter. Consume 3 valores de entrada (%operand, %source y %init_value), 2 funciones de entrada y 3 atributos de entrada (window_dimensions, window_strides y padding). Observa cómo la firma de la operación solo incluye los tipos de sus valores de entrada (pero no los tipos de funciones y atributos de entrada que se proporcionan intercalados).

%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

Las constantes de StableHLO tienen un literal y un tipo que, en conjunto, representan un valor de StableHLO. Por lo general, el tipo forma parte de la sintaxis de la constante, excepto cuando es inequívoco (p.ej., una constante booleana tiene el tipo i1 de forma inequívoca, mientras que una constante de número entero puede tener varios tipos posibles).

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

Las constantes booleanas representan los valores booleanos true y false. Las constantes booleanas tienen el tipo i1.

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

Las constantes de números enteros representan valores de números enteros a través de cadenas que usan notación decimal o hexadecimal. No se admiten otras bases, como las binarias o las octal. Las constantes de números enteros tienen las siguientes restricciones:

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

Las constantes de punto flotante representan valores de punto flotante a través de cadenas que usan notación decimal o científica. Además, la notación hexadecimal se puede usar para especificar directamente los bits subyacentes en el formato de punto flotante del tipo correspondiente. Las constantes de punto flotante tienen las siguientes restricciones:

• (C1) Si se usa una notación no hexadecimal, is_wellformed(float_literal, float_type).
• (C2) Si se usa la notación hexadecimal, size(hexadecimal_digits) = num_bits(float_type) / 4.

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

Las constantes complejas representan valores complejos con listas de una parte real (primero) y una parte imaginaria (segundo). Por ejemplo, (1.0, 0.0) : complex<f32> representa 1.0 + 0.0i y (0.0, 1.0) : complex<f32> representa 0.0 + 1.0i. El orden en que estas partes se almacenan en la memoria se define en la implementación. Las constantes complejas tienen las siguientes restricciones:

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

Las constantes de tensor representan valores de tensores con listas anidadas especificadas a través de la notación NumPy. Por ejemplo, dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> representa un valor de tensor con la siguiente asignación de índices a elementos: {0, 0} => 1, {0, 1} => 2, {0, 2} => 3, {1, 0} => 4, {1, 1} => 5 y {1, 2} => 6. El orden en el que se almacenan estos elementos en la memoria se define en la implementación. Las constantes de tensores tienen las siguientes restricciones:

• (C1) has_syntax(tensor_literal, element_type(tensor_type)), en el 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)), donde:
  • 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:]).
  • de lo contrario, false.

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

Las constantes de tensores cuantificados representan valores de tensores cuantificados con la misma notación que las constantes de tensores, con elementos especificados como constantes de su tipo de almacenamiento. Las constantes de tensores cuantificados tienen las siguientes restricciones:

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

Los literales de cadena consisten en bytes especificados con caracteres ASCII y secuencias de escape. Son independientes de la codificación, por lo que la interpretación de estos bytes está definida por la implementación. Los literales de cadena tienen el tipo string.

Ops

abs

Semántica

Realiza la operación abs a nivel de elementos en el tensor operand y produce un tensor result. Según el tipo de elemento, hace lo siguiente:

• Para números enteros firmados: módulo de números enteros
• Para números de punto flotante: abs de IEEE-754.
• Para números complejos: módulo complejo.
• Para tipos cuantificados: dequantize_op_quantize(abs, operand, type(result)).

Entradas

Salidas

Limitaciones

• (C1) shape(result) = shape(operand).
• (C2) baseline_element_type(result) se define de la siguiente manera:
  • complex_element_type(element_type(operand)) si is_complex(operand).
  • baseline_element_type(operand) en caso contrario.

Ejemplos

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

 Más ejemplos

add

Semántica

Realiza la adición a nivel de elementos de dos tensores lhs y rhs y produce un tensor result. Según el tipo de elemento, realiza las siguientes acciones:

• Para valores booleanos: OR lógico.
• Para números enteros: suma de números enteros.
• Para números de punto flotante: addition de IEEE-754.
• Para números complejos: suma compleja.
• Para tipos cuantificados: dequantize_op_quantize(add, lhs, rhs, type(result)).

Entradas

Salidas

Limitaciones

• Si la operación usa tensores no cuantificados, haz lo siguiente:
  • (C1) type(lhs) = type(rhs) = type(result).
• Si la operación usa tensores cuantificados, haz lo siguiente:
  • (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) Si es is_per_axis_quantized(lhs), entonces quantization_dimension(lhs) = quantization_dimension(result).
  • (C7) Si is_per_axis_quantized(rhs), entonces quantization_dimension(rhs) = quantization_dimension(result).

Ejemplos

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

Más ejemplos

after_all

Semántica

Garantiza que las operaciones que producen inputs se ejecuten antes que cualquier operación que dependa de result. La ejecución de esta operación no hace nada, solo existe para establecer dependencias de datos de result a inputs.

Entradas

Salidas

Ejemplos

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

Más ejemplos

all_gather

Semántica

Dentro de cada grupo de procesos en la cuadrícula de procesos StableHLO, concatena los valores de los tensores operands de cada proceso a lo largo de all_gather_dim y produce tensores results.

La operación divide la cuadrícula del proceso de StableHLO en process_groups, que se define de la siguiente manera:

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

Luego, dentro de cada process_group, haz lo siguiente:

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

Entradas

Salidas

Limitaciones

• (C1) 0 <= all_gather_dim < rank(operands...).
• (C2) is_unique(replica_groups)
• (C3) size(replica_groups) se define de la siguiente manera:
  • num_replicas si se usa cross_replica.
  • num_replicas si se usa cross_replica_and_partition.
  • num_processes si se usa flattened_ids.
• (C4) 0 <= replica_groups < size(replica_groups).
• (C5) Si use_global_device_ids = true, entonces channel_id > 0.
• (C6) type(results...) = type(operands...), excepto:
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1).

Ejemplos

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

Más ejemplos

all_reduce

Semántica

Dentro de cada grupo de procesos en la cuadrícula de procesos de StableHLO, aplica una función de reducción computation a los valores de los tensores operands de cada proceso y produce tensores results.

La operación divide la cuadrícula del proceso StableHLO en process_groups, que se define de la siguiente manera:

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

Luego, dentro de cada process_group, haz lo siguiente:

• results...@process[result_index] = exec(schedule) para algún árbol binario schedule en el que:
  • exec(node) = computation(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value.
• schedule es un árbol binario definido por la implementación cuyo recorrido en orden es to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0])).

Entradas

Salidas

Limitaciones

• (C1) is_unique(replica_groups).
• (C2) size(replica_groups) se define de la siguiente manera:
  • num_replicas si se usa cross_replica.
  • num_replicas si se usa cross_replica_and_partition.
  • num_processes si se usa flattened_ids.
• (C3) 0 <= replica_groups < size(replica_groups).
• (C4) Si use_global_device_ids = true, entonces channel_id > 0.
• (C5) computation tiene el tipo (tensor<E>, tensor<E>) -> (tensor<E>) en el que is_promotable(element_type(operand), E).
• (C6) shape(results...) = shape(operands...).
• (C7) element_type(results...) = E.

Ejemplos

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

Más ejemplos

all_to_all

Semántica

Dentro de cada grupo de procesos en la cuadrícula de procesos de StableHLO, divide los valores de los tensores operands a lo largo de split_dimension en partes, dispersa las partes divididas entre los procesos, concatena las partes dispersas a lo largo de concat_dimension y produce tensores results. La operación divide la cuadrícula del proceso de StableHLO en process_groups, que se define de la siguiente manera:

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

Luego, dentro de cada process_group, haz lo siguiente:

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

Entradas

Salidas

Limitaciones

• (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) se define de la siguiente manera:
  • num_replicas si se usa cross_replica.
  • num_partitions si se usa cross_partition.
• (C7) 0 <= replica_groups < size(replica_groups).
• (C8) dim(replica_groups, 1) = split_count.
• (C9) type(results...) = type(operands...), excepto si split_dimension != concat_dimension:
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count.
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count.

Ejemplos

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

 Más ejemplos

y

Semántica

Realiza el operador AND a nivel de elementos de dos tensores, lhs y rhs, y produce un tensor result. Según el tipo de elemento, realiza las siguientes acciones:

• Para valores booleanos: AND lógico.
• Para números enteros: AND binario.

Entradas

Salidas

Limitaciones

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

Ejemplos

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

 Más ejemplos

atan2

Semántica

Realiza la operación atan2 a nivel de elementos en el tensor lhs y rhs y produce un tensor result. Según el tipo de elemento, realiza las siguientes acciones:

• Para números de punto flotante: atan2 de IEEE-754.
• Para números complejos: atan2 complejo.
• Para tipos cuantificados: dequantize_op_quantize(atan2, lhs, rhs, type(result)).

Entradas

Salidas

Limitaciones

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

Ejemplos

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

 Más ejemplos

batch_norm_grad

Semántica

Calcula las gradientes de varias entradas de la propagación hacia atrás de batch_norm_training desde grad_output y produce tensores grad_operand, grad_scale y grad_offset. De manera más formal, esta operación se puede expresar como una descomposición de las operaciones StableHLO existentes mediante la sintaxis de Python de la siguiente manera:

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 los tipos cuantizados, realiza 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

Salidas

Limitaciones

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, mean, variance, grad_output, grad_operand, grad_scale y grad_offset tienen el mismo baseline_element_type.
• (C3) operand, grad_output y grad_operand tienen la misma forma.
• (C4) scale, mean, variance, grad_scale y grad_offset tienen la misma forma.
• (C5) size(scale) = dim(operand, feature_index).

Ejemplos

// %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 el tensor operand en todas las dimensiones, excepto en la dimensión feature_index, y produce un tensor result. De manera más formal, esta operación se puede expresar como una descomposición de las operaciones de StableHLO existentes con la sintaxis de Python de la siguiente manera:

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 los tipos cuantificados, realiza 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

Salidas

Limitaciones

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, mean, variance y result tienen el mismo 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).

Ejemplos

// %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 la media y la varianza en todas las dimensiones, excepto en la dimensión feature_index, y normaliza el tensor operand que produce los tensores output, batch_mean y batch_var. De forma más formal, esta operación se puede expresar como una descomposición de las operaciones de StableHLO existentes con la sintaxis de Python de la siguiente manera:

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 los tipos cuantificados, realiza 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

Salidas

Limitaciones

• (C1) 0 <= feature_index < rank(operand)
• (C2) operand, scale, offset, batch_mean, batch_var y output tienen el mismo 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).

Ejemplos

// %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 una operación de transmisión de bits en el tensor operand y produce un tensor result en el que los bits de todo el tensor operand se reinterpretan con el tipo del tensor result.

De manera más formal, teniendo en cuenta E = element_type(operand), E' = element_type(result) y R = rank(operand):

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

bits muestra la representación en memoria de un valor determinado, y su comportamiento se define en la implementación porque la representación exacta de los tensores y la representación exacta de los tipos de elementos también se definen en la implementación.

Entradas

Salidas

Limitaciones

• (C1) Dados E = is_quantized(operand) ? storage_type(operand) : element_type(operand), E' = is_quantized(result) ? storage_type(result) : element_type(result) y R = rank(operand):
  • Si es num_bits(E') = num_bits(E), shape(result) = shape(operand).
  • Si num_bits(E') < num_bits(E):
  • rank(result) = R + 1.
  • dim(result, i) = dim(operand, i) para todos los 0 <= i < R.
  • dim(result, R) * num_bits(E') = num_bits(E).
  • Si num_bits(E') > num_bits(E):
  • rank(result) = R - 1.
  • dim(result, i) = dim(operand, i) para todos los 0 <= i < R.
  • dim(operand, R - 1) * num_bits(E) = num_bits(E').
• (C2) Si is_complex(operand) or is_complex(result), entonces is_complex(operand) and is_complex(result).

Ejemplos

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

 Más ejemplos

broadcast_in_dim

Semántica

Expande las dimensiones o el rango de un tensor de entrada mediante la duplicación de los datos en el tensor operand y produce un tensor result. De forma más formal, result[result_index] = operand[operand_index], donde para todos los d en axes(operand):

• operand_index[d] = 0 si dim(operand, d) = 1.
• De lo contrario, operand_index[d] = result_index[broadcast_dimensions[d]].

Entradas

Salidas

Limitaciones

• (C1) element_type(result) se obtiene de la siguiente manera:
  • element_type(operand), si !is_per_axis_quantized(operand).
  • element_type(operand), excepto que quantization_dimension(operand), scales(operand) y zero_points(operand) pueden diferir de quantization_dimension(result), scales(result) y zero_points(result) respectivamente, de lo contrario.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions)
• (C5) Para todos los d en axes(operand):
  • dim(operand, d) = 1 o
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Si is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Si dim(operand, quantization_dimension(operand)) = 1, entonces scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))).

Ejemplos

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

 Más ejemplos

caso

Semántica

Produce el resultado de la ejecución de exactamente una función de branches, según el valor de index. De manera más formal, result = selected_branch(), que implica lo siguiente:

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

Entradas

Salidas

Limitaciones

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

Ejemplos

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

 Más ejemplos

cbrt

Semántica

Realiza la operación de raíz cúbica a nivel de elementos en el tensor operand y produce un tensor result. Según el tipo de elemento, realiza las siguientes acciones:

• Para números de punto flotante: rootn(x, 3) de IEEE-754.
• Para números complejos: raíz cúbica compleja.
• Para tipos cuantizados: dequantize_op_quantize(cbrt, operand, type(result))

Entradas

Salidas

Limitaciones

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

Ejemplos

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

 Más ejemplos

ceil

Semántica

Realiza el techo de elementos del tensor operand y produce un tensor result. Implementa la operación roundToIntegralTowardPositive de la especificación IEEE-754. Para los tipos cuantificados, realiza dequantize_op_quantize(ceil, operand, type(result)).

Entradas

Salidas

Limitaciones

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

Ejemplos

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

 Más ejemplos

cholesky

Semántica

Calcula la descomposición de Cholesky de un lote de matrices.

De forma más formal, para todos los i en index_space(result), result[i0, ..., iR-3, :, :] es una descomposición de Cholesky de a[i0, ..., iR-3, :, :], en forma de matriz triangular inferior (si lower es true) o triangular superior (si lower es false). Los valores de salida en el triángulo opuesto, es decir, el triángulo superior estricto o el triángulo inferior estricto, según corresponda, se definen según la implementación.

Si existe i en la que la matriz de entrada no es una matriz hermítica positiva definida, el comportamiento no está definido.

Para los tipos cuantizados, realiza dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)).

Entradas

Salidas

Limitaciones

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

Ejemplos

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

restringir

Semántica

Limita cada elemento del tensor operand entre un valor mínimo y máximo, y produce un tensor result. De forma más formal, result[result_index] = minimum(maximum(operand[result_index], min_element), max_element), donde min_element = rank(min) = 0 ? min[] : min[result_index], max_element = rank(max) = 0 ? max[] : max[result_index]. Para los tipos cuantificados, realiza dequantize_op_quantize(clamp, min, operand, max, type(result)).

Imponer un orden en números complejos implica una semántica sorprendente, por lo que, en el futuro, planeamos quitar la compatibilidad con los números complejos para esta operación (#560).

Entradas

Salidas

Limitaciones

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

Ejemplos

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

Más ejemplos

collective_broadcast

Semántica

Dentro de cada grupo de procesos en la cuadrícula de procesos StableHLO, envía el valor del tensor operand desde el proceso de origen a los procesos de destino y produce un tensor result.

La operación divide la cuadrícula del proceso de StableHLO en process_groups, que se define de la siguiente manera:

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

Luego, result@process se obtiene de la siguiente manera:

• operand@process_groups[i, 0] si existe un i de modo que el proceso esté en process_groups[i].
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) en caso contrario.

Entradas

Salidas

Limitaciones

• (C1) is_unique(replica_groups).
• (C2) 0 <= replica_groups < N, en el que N se define de la siguiente manera:
  • num_replicas si se usa cross_replica.
  • num_partitions si se usa cross_partition.
• (C3) type(result) = type(operand).

Ejemplos

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

Dentro de cada grupo de procesos en la cuadrícula de procesos de StableHLO, envía el valor del tensor operand del proceso de origen al proceso de destino y produce un tensor result.

La operación divide la cuadrícula del proceso de StableHLO en process_groups, que se define de la siguiente manera:

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

Luego, result@process se obtiene de la siguiente manera:

• operand@process_groups[i, 0], si existe un 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)) en caso contrario.

Entradas

Salidas

Limitaciones

• (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, donde N se define de la siguiente manera:
  • num_replicas si se usa cross_replica.
  • num_partitions si se usa cross_partition.
• (C5) type(result) = type(operand).

Ejemplos

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

 Más ejemplos

comparar

Semántica

Realiza una comparación a nivel de elementos de los tensores lhs y rhs según comparison_direction y compare_type, y produce un tensor result.

Los valores de comparison_direction y compare_type tienen la siguiente semántica:

Para los tipos de elementos booleanos y enteros, haz lo siguiente:

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

Para los tipos de elementos de punto flotante con compare_type = FLOAT, la operación implementa las siguientes operaciones IEEE-754:

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

Para los elementos de punto flotante con compare_type = TOTALORDER, la op usa la combinación de operaciones totalOrder y compareQuietEqual de IEEE-754.

Para los tipos de elementos complejos, la comparación lexicográfica de los pares (real, imag) se realiza con los comparison_direction y compare_type proporcionados. Imponer un orden en los números complejos implica una semántica sorprendente, por lo que, en el futuro, planeamos quitar la compatibilidad con los números complejos cuando comparison_direction sea GE, GT, LE o LT (#560).

Para tipos cuantizados, realiza dequantize_compare(lhs, rhs, comparison_direction).

Entradas

Salidas

Limitaciones

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs).
• (C2) shape(lhs) = shape(rhs) = shape(result).
• (C3) compare_type se define de la siguiente manera:
  • SIGNED si is_signed_integer(element_type(lhs)).
  • UNSIGNED si is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)).
  • FLOAT o TOTALORDER si es is_float(element_type(lhs)).
  • FLOAT si is_complex(element_type(lhs)).

Ejemplos

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

Más ejemplos

emergencia compleja,

Semántica

Realiza la conversión de elementos a un valor complejo a partir de un par de valores reales e imaginarios, lhs y rhs, y produce un tensor result.

Entradas

Salidas

Limitaciones

• (C1) type(lhs) = type(rhs).
• (C2) shape(result) = shape(lhs).
• (C3) element_type(result) tiene el tipo complex<E>, en el que E = element_type(lhs).

Ejemplos

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

 Más ejemplos

compuesto

Semántica

Encapsula una operación compuesta por otras operaciones de StableHLO, que toma inputs y composite_attributes y produce results. La semántica de la operación se implementa con el atributo decomposition. La operación composite se puede reemplazar por su descomposición sin cambiar la semántica del programa. En los casos en que la incorporación de la descomposición no proporciona la misma semántica de operación, prefiere usar custom_call.

El campo version (el valor predeterminado es 0) se usa para indicar cuándo cambia la semántica de un compuesto.

Entradas

Salidas

Limitaciones

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

Ejemplos

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

Más ejemplos

concatenate

Semántica

Concatena inputs junto con la dimensión dimension en el mismo orden que los argumentos dados y produce un tensor result. De forma más formal, result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1], en la que:

1. id = d0 + ... + dk-1 + kd.
2. d es igual a dimension, y d0, ... son los tamaños de da dimensión de inputs.

Entradas

Salidas

Limitaciones

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

Ejemplos

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

 Más ejemplos

constante

Semántica

Produce un tensor output a partir de una constante value.

Entradas

Salidas

Limitaciones

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

Ejemplos

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

 Más ejemplos

generar una conversión

Semántica

Realiza una conversión por elemento de un tipo de elemento a otro en el tensor operand y produce un tensor result.

Para las conversiones de boolean-to-any-supported-type, el valor false se convierte en cero y el valor true se convierte en uno. En el caso de las conversiones de any-supported-type-to-boolean, un valor cero se convierte en false y los valores distintos de cero se convierten en true. A continuación, se explica cómo funciona esto para los tipos complejos.

En el caso de las conversiones que involucran número entero a número entero, número entero a número de punto flotante o número de punto flotante a número de punto flotante, si el valor de origen se puede representar exactamente en el tipo de destino, el valor del resultado es esa representación exacta. De lo contrario, el comportamiento está por definir (#180).

En el caso de las conversiones que involucran números de punto flotante a números enteros, la parte fraccionaria se trunca. Si el valor truncado no se puede representar en el tipo de destino, el comportamiento es por definir (#180).

Las conversiones que involucran números complejos a números complejos siguen el mismo comportamiento de las conversiones de números de punto flotante a números de punto flotante para convertir partes reales e imaginarias.

Para las conversiones complex-to-any-other-type y any-other-type-to-complex, se ignora el valor imaginario de origen o se establece en cero el valor imaginario de destino, respectivamente. La conversión de la parte real sigue las conversiones de punto flotante.

En principio, esta operación podría expresar la decuantización (conversión de tensores cuantificados a tensores normales), la cuantificación (conversión de tensores normales a tensores cuantificados) y la recantidad (conversión entre tensores cuantificados), pero, por el momento, tenemos operaciones dedicadas para eso: uniform_dequantize para el primer caso de uso y uniform_quantize para el segundo y el tercer caso de uso. En el futuro, es posible que estas dos operaciones se combinen en convert (#1576).

Entradas

Salidas

Limitaciones

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

Ejemplos

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

 Más ejemplos

convolución

Semántica

Calcula los productos escalares entre ventanas de lhs y rebanadas de rhs y produce result. En el siguiente diagrama, se muestra cómo se calculan los elementos de result a partir de lhs y rhs con un ejemplo concreto.

De manera más formal, considera el siguiente replanteamiento de las entradas en términos de lhs para poder expresar ventanas 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).

Este replanteamiento usa las siguientes funciones 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] en el que j[d] = i[permutation[d]].

Si es feature_group_count = 1 y batch_group_count = 1, entonces para todos los output_spatial_index en index_space(dim(result, output_spatial_dimensions...)), result[result_shape(:, output_spatial_index, :)] = dot_product, donde:

• 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]). Parece que esta función no se usa, por lo que planeamos quitarla en el 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]).

Si es 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).

Si 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 los tipos cuantificados, realiza 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 los tipos cuantificados híbridos, realiza 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

Salidas

Limitaciones

• (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 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) se define de la siguiente manera:
  • dim(lhs, input_batch_dimension) / batch_group_count si result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) si result_dim = output_feature_dimension.
  • num_windows de lo contrario, donde:
  • 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.
• Si la operación usa tensores no cuantificados, haz lo siguiente:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Si la operación usa tensores cuantificados, haz lo siguiente:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Si es is_per_axis_quantized(rhs), entonces quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Si is_per_axis_quantized(result), entonces quantization_dimension(result) = output_feature_dimension.
  • Si is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Si is_per_tensor_quantized(rhs), entonces is_per_tensor_quantized(result).
  • Si !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Ejemplos

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

 Más ejemplos

coseno

Semántica

Realiza una operación de coseno en términos de elementos en el tensor operand y produce un tensor result. Según el tipo de elemento, hace lo siguiente:

• Para números de punto flotante: cos de IEEE-754.
• Para números complejos: coseno complejo.
• Para tipos cuantificados: dequantize_op_quantize(cosine, operand, type(result)).

Entradas

Salidas

Limitaciones

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

Ejemplos

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

 Más ejemplos

count_leading_zeros

Semántica

Realiza un recuento por elemento de la cantidad de bits de ceros iniciales en el tensor operand y produce un tensor result.

Entradas

Salidas

Limitaciones

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

Ejemplos

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

 Más ejemplos

custom_call

Semántica

Encapsula una operación call_target_name definida por la implementación que toma inputs y called_computations y produce results. Se pueden usar has_side_effect, backend_config y api_version para proporcionar metadatos adicionales definidos por la implementación.

Por el momento, esta operación contiene una colección bastante desorganizada de metadatos que refleja la evolución orgánica de su operación equivalente en el compilador de XLA. En el futuro, planeamos unificar estos metadatos (#741).

Entradas

Salidas

Ejemplos

%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 la división por elementos de los tensores lhs del dividendo y rhs del divisor, y produce un tensor result. Según el tipo de elemento, hace lo siguiente:

• Para números enteros: División de números enteros que produce el cociente algebraico con cualquier parte fraccionaria descartada.
• Para números de punto flotante: division de IEEE-754.
• Para números complejos: división compleja.
• Para tipos cuantizados:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)).

Entradas

Salidas

Limitaciones

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

Ejemplos

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

 Más ejemplos

dot_general

Semántica

Calcula los productos punto entre las porciones de lhs y las porciones de rhs, y produce un tensor result.

Más formalmente, result[result_index] = dot_product, donde:

• 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 en el que size(result_batching_index) = size(lhs_batching_dimensions), size(result_lhs_index) = size(lhs_result_dimensions) y 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 los tipos cuantificados, realiza 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 híbridos cuantizados, realiza 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 el equilibrio entre la velocidad y la precisión para los procesamientos en los backends del acelerador. Puede ser uno de los siguientes (por el momento, la semántica de estos valores de enum no está especificada, pero planeamos abordar esto en #755):

• DEFAULT: Es el cálculo más rápido, pero la aproximación menos precisa al número original.
• HIGH: Cálculo más lento, pero aproximación más precisa al número original.
• HIGHEST: Es el cálculo más lento, pero la aproximación más precisa al número original.

Un DotAlgorithm define las propiedades principales del algoritmo que se usa para implementar la operación punto, que también define la precisión. Si se establecen los campos de atributos del algoritmo, precision_config debe ser DEFAULT. DotAlgorithms no tiene un valor predeterminado, ya que los parámetros predeterminados se definen en la implementación. Por lo tanto, todos los campos del algoritmo de punto se pueden establecer en None para especificar un algoritmo de punto vacío, que usará el valor precision_config.

Los campos DotAlgorithm incluyen lo siguiente:

• lhs_precision_type y rhs_precision_type, las precisiones a las que se redondean el LHS y el RHS de la operación. Los tipos de precisión son independientes de los tipos de almacenamiento de las entradas y la salida.
• accumulation_type es la precisión que se usa para la acumulación.
• lhs_component_count, rhs_component_count y num_primitive_operations se aplican cuando realizamos un algoritmo que descompone el lado izquierdo o derecho en varios componentes y realiza varias operaciones de punto "primitivo" en esos valores, por lo general, para emular una precisión más alta (p. ej., Aprovecha el tipo de datos de inteligencia artificial bfloat16 para realizar cálculos de mayor precisión: bf16_6x tf32_3x, etc.). Para los algoritmos sin descomposición, estos valores deben establecerse en 1.
• allow_imprecise_accumulation para especificar si se permite la acumulación en una precisión más baja para algunos pasos (p. ej., CUBLASLT_MATMUL_DESC_FAST_ACCUM).

Atributos DotAlgorithm de ejemplo:

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

Depende de las implementaciones decidir qué combinaciones son compatibles. En general, no se garantiza que el consumidor de StableHLO admita cada algoritmo en cada tipo de acelerador. Si no se admite un algoritmo determinado, se debe generar un error en lugar de recurrir a una alternativa. La verificación de StableHLO proporcionará una verificación del mejor esfuerzo, lo que evitará que se usen algoritmos que no se sepan compatibles con ningún hardware.

Consulta xla_data.proto > Algorithm para ver algunos valores de algoritmos compatibles. En el ticket n° 2483, se captura el plan para crear un doc centralizado sobre los algoritmos compatibles por backend.

Entradas

Salidas

Limitaciones

• (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).
• Si la operación usa tensores no cuantificados, haz lo siguiente:
  • (C13) element_type(lhs) = element_type(rhs).
• Si la operación usa tensores cuantificados, haz lo siguiente:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C15) zero_points(rhs) = 0.
  • (C16) Si is_per_axis_quantized(rhs), entonces quantization_dimension(rhs) no está en rhs_contracting_dimensions.
  • Si is_quantized(lhs):
  • (C17) storage_type(lhs) = storage_type(rhs).
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C19) Si is_per_tensor_quantized(rhs), entonces is_per_tensor_quantized(result).
  • Si !is_quantized(lhs):
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result).
• Si !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.

Ejemplos

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

 Más ejemplos

dynamic_broadcast_in_dim

Semántica

Esta operación es funcionalmente idéntica a la operación broadcast_in_dim, pero la forma del resultado se especifica de forma dinámica a través de output_dimensions.

La operación también acepta atributos opcionales known_expanding_dimensions y known_nonexpanding_dimensions para expresar el conocimiento estático sobre el comportamiento de expansión de las dimensiones. Si no se especifica, se supone que todas las dimensiones pueden expandirse.

Entradas

Salidas

Limitaciones

• (C1) element_type(result) se obtiene de la siguiente manera:
  • element_type(operand), si !is_per_axis_quantized(operand).
  • element_type(operand), excepto que quantization_dimension(operand), scales(operand) y zero_points(operand) pueden diferir de quantization_dimension(result), scales(result) y zero_points(result) respectivamente, de lo contrario.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions)
• (C5) Para todos los d en axes(operand):
  • dim(operand, d) = 1 o
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Si is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Si es dim(operand, quantization_dimension(operand)) = 1, entonces 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).

Ejemplos

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

 Más ejemplos

dynamic_conv

Semántica

Esta operación es funcionalmente idéntica a la operación de convolución, pero el padding se especifica de forma dinámica a través de padding.

Entradas

Salidas

Limitaciones

• (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 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) se define de la siguiente manera:
  • dim(lhs, input_batch_dimension) / batch_group_count si result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) si result_dim = output_feature_dimension.
  • num_windows de lo contrario, donde:
  • 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.
• Si la operación usa tensores no cuantificados, haz lo siguiente:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Si la operación usa tensores cuantificados, haz lo siguiente:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Si es is_per_axis_quantized(rhs), entonces quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Si is_per_axis_quantized(result), entonces quantization_dimension(result) = output_feature_dimension.
  • Si is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Si is_per_tensor_quantized(rhs), entonces is_per_tensor_quantized(result).
  • Si !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Ejemplos

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