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 frameworks y compiladores 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 creando más interoperabilidad entre varios frameworks de AA (como TensorFlow, JAX y PyTorch) y compiladores de AA (como XLA y IREE). Para ello, este documento proporciona una especificación del lenguaje de programación StableHLO.

Esta especificación contiene tres secciones principales. Primero, la sección Programas describe la estructura de los programas de StableHLO, que constan de funciones de StableHLO, que, a su vez, constan de 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 la semántica para todas estas operaciones que se ejecutan juntas dentro de un programa. Por último, en la sección Notación, 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 repo en la versión etiquetada que te interese. Por ejemplo, la especificación de StableHLO v0.19.0. Para ver los cambios que se produjeron en cada incremento de versión secundaria de StableHLO, consulta el registro de versiones en VhloDialect.td.

Programas

Program ::= {Func}

Los programas de StableHLO constan de una cantidad arbitraria de funciones de StableHLO. A continuación, se muestra un ejemplo de programa 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 de StableHLO (también llamadas funciones con nombre) tienen un identificador, entradas y salidas, y un cuerpo. En el futuro, planeamos introducir metadatos adicionales para las funciones y lograr una mejor compatibilidad con HLO (#425, #626, #740, #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 sigilos que distinguen diferentes tipos de identificadores y 2) los identificadores de valores pueden ser completamente numéricos para simplificar la generación de programas de StableHLO.

Tipos

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

Los tipos de StableHLO se clasifican en tipos de valor (también llamados 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 de muchos lenguajes de programación, y su 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, arrays multidimensionales. Tienen una forma y un tipo de elemento, en los que una forma representa tamaños de dimensión no negativos o desconocidos en orden ascendente de las dimensiones correspondientes (que también se denominan ejes) numeradas del 0 al R-1. La cantidad de dimensiones R se denomina rango. Por ejemplo, tensor<2x3xf32> es un tipo de tensor con forma 2x3 y 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 rango es 2.

Las formas pueden ser parcial o completamente desconocidas (dinámicas), p.ej., tensor<?x2xf64> es parcialmente desconocida y tensor<?x?xf64> es completamente desconocida. Los tamaños de dimensiones dinámicas se representan con un ?. No se pueden quitar las clasificaciones de las formas.

En el futuro, planeamos explorar la posibilidad de extender los tipos de tensores más allá de los tamaños de las dimensiones 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 entero determinado i, el valor de punto flotante correspondiente f se puede calcular como f = (i - zero_point) * scale, donde scale y zero_point se denominan parámetros de cuantificación. Los parámetros 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 existe un gran interés en las escalas basadas en números enteros, representadas con multiplicadores y desplazamientos. Tenemos previsto explorar esta opción en el futuro cercano (#1404).

Hay un debate en curso sobre la semántica de QuantizationZeroPoint, incluido el tipo, los valores y si puede haber solo un punto cero o potencialmente varios en un tipo de tensor cuantificado. Según los resultados de esta conversación, es posible que la especificación sobre los puntos nulos cambie en el futuro (#1405).

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

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

Los tipos de tensores cuantificados representan tensores con elementos cuantificados. Estos tensores son exactamente iguales a los tensores regulares, excepto que sus elementos tienen tipos de elementos cuantificados, en lugar de tipos de elementos regulares.

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

• Para la cuantificación por tensor, haz lo siguiente:
  • 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 búfer representan búferes. Por ejemplo, en XLA, los búferes son arrays multidimensionales con almacenamiento coherente. Al igual que los tipos de tensor, los tipos de búfer tienen una forma y un tipo de elemento, en los que una forma representa tamaños de dimensión no negativos o desconocidos en orden ascendente de las dimensiones correspondientes (que también se denominan ejes) numeradas del 0 al R-1. La cantidad de dimensiones R se denomina rango. Por ejemplo, memref<2x3xf32> es un tipo de búfer con forma 2x3 y 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 rango es 2.

Los búferes se pueden asignar con un custom_call a CreateBuffer o Pin, y se pueden liberar con un custom_call a Unpin. Solo las operaciones de custom_call pueden leer y escribir el contenido dentro de los búferes. Consulta custom_call para obtener más detalles.

Los tipos de tuplas representan tuplas, es decir, listas heterogéneas. Las tuplas son una función heredada que solo existe para la compatibilidad con HLO. En HLO, las tuplas se usan para representar entradas y salidas variádicas. En StableHLO, las entradas y salidas variádicas se admiten de forma nativa, y el único uso de tuplas en StableHLO es para 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, lo que podría permitirnos quitar los tipos de tuplas 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 (por lo tanto, 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 ser con signo (si) o sin signo (ui) y tener uno de los anchos de bits admitidos (2, 4, 8, 16, 32 o 64). Los tipos siN con signo representan valores enteros desde -2^(N-1) hasta 2^(N-1)-1 inclusive, y los tipos uiN sin signo representan valores enteros desde 0 hasta 2^N-1 inclusive.
• Los tipos 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 FP8 Formats for Deep Learning.
  • Tipos f8E4M3FNUZ y f8E5M2FNUZ que corresponden a las codificaciones E4M3 y E5M2 de los formatos FP8 que se describen 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 Hybrid 8-bit Floating Point (HFP8) Training and Inference for Deep Neural Networks.
  • Tipo bf16 que corresponde al formato bfloat16 descrito en BFloat16: El secreto del alto rendimiento en las Cloud TPU.
  • Tipos f16, f32 y f64 que corresponden, respectivamente, a los formatos binary16 (“precisión media”), binary32 (“precisión simple”) y binary64 (“precisión doble”) que se describen en el estándar IEEE 754.
  • El tipo tf32 corresponde al formato TensorFloat32 y tiene asistencia limitada en StableHLO.
  • Tipos de MX (microescalado) f4E2M1FN, f6E2M3FN, f6E3M2FN y f8E8M0FNU que se describen en la Especificación de formatos de microescalado de OCP.
• Los tipos complejos representan valores complejos que tienen una parte real y una parte imaginaria del mismo tipo de elemento. Los tipos complejos admitidos 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 a la izquierda de ->) y tipos de salida (la lista de tipos a la derecha 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 de primera clase en StableHLO y solo se usa para especificar metadatos estáticos para los elementos del programa.

Operaciones

Las operaciones de StableHLO (también llamadas ops) representan un conjunto cerrado de operaciones de alto nivel en los 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 podría decirse que es la que mejor se adapta al objetivo de StableHLO de crear más interoperabilidad entre los frameworks de AA y los compiladores de AA.

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

Las operaciones de StableHLO (también llamadas ops) tienen un nombre, entradas y salidas, y una firma. El nombre consta del prefijo stablehlo. y una mnemotecnia que identifica de forma única una de las operaciones admitidas. A continuación, se incluye 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, ya que en StableHLO las funciones no son valores de primera clase) y atributos de entrada (también proporcionados 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 por lo siguiente: 1) No tienen un identificador (de ahí el nombre “anónimas”). 2) No declaran tipos de salida (los tipos de salida se infieren a partir 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 Unused anterior), que está allí para la compatibilidad con MLIR. En MLIR, existe un concepto más general de "regiones" que pueden tener varios "bloques" de operaciones conectados a través de operaciones de salto. 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 sigue 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 del 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. Del mismo modo, la operación slice usa varios atributos, como start_indices y limit_indices, para especificar los límites que se usan para segmentar el valor de entrada.

Por el momento, los programas de StableHLO en la naturaleza 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 ->). En sentido estricto, los tipos de entrada son redundantes, y los tipos de salida también lo son casi siempre (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 la operación forma parte deliberadamente de la sintaxis de StableHLO para garantizar la compatibilidad con MLIR.

A continuación, se muestra un ejemplo de una 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). Ten en cuenta que 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 de forma intercalada).

%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, juntos, representan un valor de StableHLO. En general, el tipo forma parte de la sintaxis de la constante, excepto cuando no es ambiguo (p.ej., una constante booleana tiene de forma no ambigua el tipo i1, 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 la binaria o la 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, se puede usar la notación hexadecimal 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, se debe incluir 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 está definido por 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 tensor con listas anidadas especificadas a través de la notación de 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, {1, 2} => 6. El orden en el que estos elementos se almacenan en la memoria está definido por la implementación. Las constantes de tensor tienen las siguientes restricciones:

• (C1) has_syntax(tensor_literal, element_type(tensor_type)), donde:
  • 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 tensor cuantificado representan valores de tensor cuantificado con la misma notación que las constantes de tensor, 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 constan de bytes especificados con caracteres ASCII y secuencias de escape. No dependen de la codificación, por lo que la interpretación de estos bytes se define en la implementación. Los literales de cadena tienen el tipo string.

Operaciones

abs

Semántica

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

• Para números enteros con signo, es el 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).
  • En caso contrario, baseline_element_type(operand).

Ejemplos

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

 Más ejemplos

add

Semántica

Realiza la suma de dos tensores lhs y rhs elemento por elemento, y produce un tensor result. Según el tipo de elemento, hace lo siguiente:

• 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 cuantizados, se aplica 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 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)< : (ten>sor2x2xi<32, ten>sor>2x2xi32<) - ten>sor2x2xi32
// %result: [[6, 8], [10, 12]]

 Más ejemplos

after_all

Semántica

Garantiza que las operaciones que producen el 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, !stablehl>o.token) - !stablehlo.token

 Más ejemplos

all_gather

Semántica

Dentro de cada grupo de procesos en la cuadrícula de procesos de 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 de procesos 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_grou<ps = den>se[[0, 1]<] : ten>sor1x2xi64,
  // channel_id = 0
  channel_handle = #stablehlo.chan<nel_handlehandle = 0>, type = 0
  // use_global_device_ids = false
}< : (ten>sor2x2xi<64, ten>sor>2x2xi64)< - (ten>sor2x4xi<64, ten>sor2x4xi64)
// %result0@(0, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result0@(1, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result1@(0, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]
// %result1@(1, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]

 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 de procesos 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:

• results...@process[result_index] = exec(schedule) para algún árbol binario schedule donde:
  • 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>), donde 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(%ar<g0:> tensori64, %ar<g1:> tensori64):
    %0 = "stablehlo.add"(%ar<g0,> %arg1) <: (>ten>sori64,< te>nsori64) - tensori64
    "stable<hlo>.re>turn"(%0) : (tensori64) - ()<
}) {
  >replica_g<roups => dense[[0, 1]] : tensor1x2xi64,
  // channel_id = 0
  channel_hand<le = #stablehlo.chan>nel_handlehandle = 0, type = 0
  // use_global_<devic>e_ids = <false>
} >: (tenso<r4xi6>4, tenso<r4xi6>4) - (tensor4xi64, tensor4xi64)
// %result0@(0, 0): [6, 8, 10, 12]
// %result0@(1, 0): [6, 8, 10, 12]
// %result1@(0, 0): [22, 24, 26, 28]
// %result1@(1, 0): [22, 24, 26, 28]

 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 de procesos 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, dentro de cada process_group, haz lo siguiente:

• split_parts...@sender = split(operands...@sender, split_count, split_dimension) para todo sender en process_group.
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group] donde receiver_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_grou<ps = den>se[[0, 1]<] : ten>sor1x2xi64
  // channel_id = 0
}< : (ten>sor2x4xi<64, ten>sor>2x4xi64)< - (ten>sor4x2xi<64, ten>sor4x2xi64)
// %result#0@(0, 0): [[1, 2], [5, 6], [9, 10], [13, 14]]
// %result#0@(1, 0): [[3, 4], [7, 8], [11, 12], [15, 16]]
// %result#1@(0, 0): [[17, 18], [21, 22], [25, 26], [29, 30]]
// %result#1@(1, 0): [[19, 20], [23, 24], [27, 28], [31, 32]]

 Más ejemplos

y

Semántica

Realiza la operación AND a nivel del elemento de dos tensores lhs y rhs, y produce un tensor result. Según el tipo de elemento, hace lo siguiente:

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

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)< : (ten>sor2x2xi<32, ten>sor>2x2xi32<) - ten>sor2x2xi32
// %result: [[1, 2], [3, 0]]

 Más ejemplos

atan2

Semántica

Realiza una operación atan2 a nivel del elemento en los tensores lhs y rhs, y produce un tensor result. Según el tipo de elemento, hace lo siguiente:

• Para números de punto flotante: atan2 de IEEE-754.
• Para números complejos: atan2 compleja.
• 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)< : (t>ensor3xf<64, t>ens>or3xf64<) - t>ensor3xf64
// %result: [0.0, 1.57079637, -1.57079637] // [0.0, pi/2, -pi/2]

 Más ejemplos

batch_norm_grad

Semántica

Calcula los gradientes de varias entradas de batch_norm_training a través de la retropropagación 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 en operaciones existentes de StableHLO con la siguiente sintaxis de Python:

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 cuantificados, 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
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ensor2xf64,
 <    tenso>r2x>2x2xf64)< - (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf64)
// %grad_operand: [
//                 [[0.0, 0.0], [0.0, 0.0]],
//                 [[0.0, 0.0], [0.0, 0.0]]
//                ]
// %grad_scale:  [0.0, 0.0]
// %grad_offset: [0.4, 0.4]

batch_norm_inference

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 en operaciones existentes de StableHLO con la siguiente sintaxis de Python:

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
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ens>or2xf64<) - tenso>r2x2x2xf64
// %result: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]

batch_norm_training

Semántica

Calcula la media y la varianza en todas las dimensiones, excepto la dimensión feature_index, y normaliza el tensor operand para producir los tensores output, batch_mean y batch_var. De manera más formal, esta operación se puede expresar como una descomposición en operaciones existentes de StableHLO con la siguiente sintaxis de Python:

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
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ens>or2xf64) -
 <   (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf64)
// %output: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]
// %batch_mean: [2.0, 3.0]
// %batch_var: [1.0, 1.0]

bitcast_convert

Semántica

Realiza una operación de bitcast 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, dadas las condiciones E = element_type(operand), E' = element_type(result) y R = rank(operand):

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

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

Entradas

Salidas

Limitaciones

• (C1) Dadas 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), haz lo siguiente:
  • 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), haz lo siguiente:
  • 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)< : >(te>nsorf64<) - t>ensor4xf16
// %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 duplicando los datos en el tensor operand y produce un tensor result. De manera más formal, result[result_index] = operand[operand_index], donde para todo d en axes(operand), se cumple lo siguiente:

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

Entradas

Salidas

Limitaciones

• (C1) element_type(result) se calcula 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, en otros casos.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Para todo 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_dimensio<ns = arra>yi64: 2, 1
}< : (ten>sor>1x3xi32<) - tenso>r2x3x2xi32
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

 Más ejemplos

caso

Semántica

Produce el resultado de ejecutar exactamente una función de branches, según el valor de index. De manera más formal, result = selected_branch(), donde:

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

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, %resul<t_bra>nch0) : <(tens>or2>xi64, tensor2xi64) - ()
}, {
  "stablehlo.return"(%result_branc<h1, %>result_b<ranch>1) >: (tensor2xi64, <ten>sor>2xi64) -< ()
}>) : (ten<sori3>2) - (tensor2xi64, tensor2xi64)
// %result0: [1, 1]
// %result1: [1, 1]

 Más ejemplos

cbrt

Semántica

Realiza una operación de raíz cúbica a nivel de los 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: rootn(x, 3) de IEEE-754.
• Para números complejos: raíz cúbica compleja.
• Para los tipos cuantificados: 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)< : (t>ens>or4xf64<) - t>ensor4xf64
// %result: [0.0, 1.0, 2.0, 3.0]

 Más ejemplos

ceil

Semántica

Realiza la operación ceil para cada elemento 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)< : (t>ens>or5xf32<) - t>ensor5xf32
// %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 manera más formal, para todo i en index_space(result), result[i0, ..., iR-3, :, :] es una descomposición de Cholesky de a[i0, ..., iR-3, :, :], en forma de una 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 estrictamente superior o el triángulo estrictamente inferior, respectivamente, se definen según la implementación.

Si existe i donde la matriz de entrada no es una matriz hermitiana definida positiva, el comportamiento no está definido.

Para los tipos cuantificados, 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
}< : (ten>sor>3x3xf32<) - ten>sor3x3xf64
// %result: [
//           [1.0, 0.0, 0.0],
//           [2.0, 4.0, 0.0],
//           [3.0, 5.0, 6.0]
//          ]

restringir

Semántica

Ajusta cada elemento del tensor operand entre un valor mínimo y uno máximo, y produce un tensor result. De manera 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 los 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)< : (t>ensor3xi<32, t>ensor3xi<32, t>ens>or3xi32<) - t>ensor3xi32
// %result: [5, 13, 20]

 Más ejemplos

collective_broadcast

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 a los procesos de destino y genera un tensor result.

La operación divide la cuadrícula de procesos 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 calcula de la siguiente manera:

• operand@process_groups[i, 0] si existe un i tal que el proceso se encuentra 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, donde 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_grou<ps = den>se[[2, 1]<] : ten>sor1x2xi64,
  channel_handle = #stablehlo.chan<nel_handlehandle = 0>, type = 0
} : (ten>sor>1x2xi64<) - ten>sor1x2xi64
// %result@(0, 0): [[0, 0]]
// %result@(1, 0): [[5, 6]]
// %result@(2, 0): [[5, 6]]
// %result@(3, 0): [[0, 0]]

collective_permute

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 de procesos 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 calcula 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_pai<rs = dense[[0, 1>], [1, 2]<] : ten>sor2x2xi64,
  channel_handle = #stablehlo.chan<nel_handlehandle = 0>, type = 0
}< : (ten>sor>2x2xi64<) - ten>sor2x2xi64
//
// %result@(0, 0): [[0, 0], [0, 0]]
// %result@(1, 0): [[1, 2], [3, 4]]
// %result@(2, 0): [[5, 6], [7, 8]]

 Más ejemplos

comparar

Semántica

Realiza una comparación de cada elemento 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 de números 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 tipos de elementos de punto flotante con compare_type = TOTALORDER, la operación 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 cuantificados, 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 = <#stablehlocomparison_di>rection LT,
  compare_type = <#stablehlocomparison_>type FLOAT
}< : (t>ensor2xf<32, t>ens>or2xf32<) - >tensor2xi1
// %result: [true, false]

 Más ejemplos

emergencia compleja,

Semántica

Realiza una conversión de cada elemento en 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>, donde E = element_type(lhs).

Ejemplos

// %lhs: [1.0, 3.0]
// %rhs: [2.0, 4.0]
%result = "stablehlo.complex"(%lhs, %rhs)< : (t>ensor2xf<64, t>ens>or2xf64<) - tenso<r2x>>complexf64
// %result: [(1.0, 2.0), (3.0, 4.0)]

 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 la operación, se prefiere usar custom_call.

El campo version (el valor predeterminado es 0) se usa para indicar cuándo cambian las semánticas de un elemento 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,
 < ve>rsion = <1 :> i3>2
} : (<ten>sorf32, tensorf32) - tensorf32

 Más ejemplos

concatenate

Semántica

Concatena inputs a lo largo de la dimensión dimension en el mismo orden que los argumentos proporcionados y produce un tensor result. De manera más formal, result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1], donde:

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

Entradas

Salidas

Limitaciones

• (C1) same(element_type(inputs...)).
• (C2) same(shape(inputs...)), excepto 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
}< : (ten>sor3x2xi<64, ten>sor>1x2xi64<) - ten>sor4x2xi64
// %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"() {
  val<ue = dense[[0.0, 1.0], [>2.0, 3.0]<] : ten>sor2x2xf3>2
} : (<) - ten>sor2x2xf32
// %output: [[0.0, 1.0], [2.0, 3.0]]

 Más ejemplos

generar una conversión

Semántica

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

En el caso de 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. Consulta a continuación cómo funciona para los tipos complejos.

En el caso de las conversiones que involucran números enteros a números enteros, números enteros a números de punto flotante o números de punto flotante a números 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 definirse (#180).

En el caso de las conversiones que involucran floating-point-to-integer, se trunca la parte fraccionaria. Si el valor truncado no se puede representar en el tipo de destino, el comportamiento está pendiente de definición (TBD) (#180).

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

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

En principio, esta operación podría expresar la desantificación (conversión de tensores cuantificados a tensores regulares), la cuantificación (conversión de tensores regulares a tensores cuantificados) y la recuantificación (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 fusionen en convert (#1576).

Entradas

Salidas

Limitaciones

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

Ejemplos

// %operand: [-1, 0, 1]
%result = "stablehlo.convert"(%operand)< : (t>ens>or3xi64<) - tenso<r3x>>complexf64
// %result: [(-1.0, 0.0), (0.0, 0.0), (1.0, 0.0)]

 Más ejemplos

convolución

Semántica

Calcula los productos punto entre las ventanas de lhs y los segmentos 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] donde j[d] = i[permutation[d]].

Si feature_group_count = 1 y batch_group_count = 1, entonces para todo 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]). Esta función parece no usarse, por lo que, en el futuro, planeamos quitarla (#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 feature_group_count > 1, haz lo siguiente:

• 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, haz lo siguiente:

• 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 en otros casos, 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 cuantizados, se aplica lo siguiente:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Si 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), haz lo siguiente:
  • (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), haz lo siguiente:
  • (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_strid<es = arra>yi64: 4, 4,
  paddi<n>g = dense<0 : ten>sor2x2xi64,
  lhs_dilati<on = arra>yi64: 2, 2,
  rhs_dilati<on = arra>yi64: 1, 1,
  window_revers<al = arrayi1: fa>lse, false,
  // In the StableHLO dialect, dimension numbers are encoded vi<a:
  // `[input >dim<ensions]x[kernel >di>mensions]-[output dimensions]`.
  // "b" is batch dimension, "f" is feature dimension,
  // "i" is input feature dimension, "o" is output feature dimension,
  // "0/1/etc" a<re spatial dimensions.
  d>imension_num>bers = #stablehlo.conv[b, 0, 1, f]x[0, 1, i, o]-[b, 0, 1, f],
  batch_group_count = 1 : i64,
  fea<ture_group_count >= 1 : i64,
 < precision_config> = [#stablehl<oprecision >DEFAULT,< #stablehlo>pre>cision <DEFAULT]
} >: (tensor1x4x4x1xi64, tensor3x3x1x1xi64) - tensor1x2x2x1xi64
// %result: [[
//            [[10], [26]],
//            [[46], [62]]
//          ]]

 Más ejemplos

coseno

Semántica

Realiza una operación de coseno a nivel de los elementos en el tensor operand y genera 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)< : (ten>sor>2x2xf32<) - ten>sor2x2xf32
// %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 cero 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)< : (ten>sor>2x2xi64<) - ten>sor2x2xi64
// %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 de metadatos bastante desorganizada 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

(Compatibilidad con GPU de XLA) Destinos de custom_call especiales

Existen tres call_target_name especiales relacionados con los tipos buffer: CreateBuffer crea un buffer sin inicializar, Pin crea un buffer inicializado y Unpin desasigna un buffer y devuelve el contenido del buffer.

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

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

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

Alias

Algunas operaciones custom_call pueden requerir que una parte de las salidas y una parte de los operandos compartan la misma memoria. Esto se puede expresar a través de output_operand_aliases. Una representación de par de alias consta de una lista de índices de tuplas de salida que representan la parte de salida y un operand_index junto con una lista de índices de tuplas de operandos que representan la parte del operando. La lista de índices de tuplas de salida o de operandos está vacía si el tipo correspondiente no es un tipo tuple y puede ser arbitrariamente larga para un tipo de tupla anidada arbitrariamente. Esto es similar a la representación del alias de XLA.

La parte de salida y la parte de entrada de un par de alias deben tener el mismo tipo. En el caso de las operaciones de custom_call que no son llamadas a CreateBuffer, Pin y Unpin, un operando buffer puede aparecer en, como máximo, un par de alias, y una salida buffer debe aparecer en un par de alias.

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

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

dividir

Semántica

Realiza la división de cada elemento de los tensores de dividendo lhs y divisor rhs, 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 los tipos cuantificados:
  • 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)< : (t>ensor4xf<32, t>ens>or4xf32<) - t>ensor4xf32
// %result: [5.66666651, -5.66666651, -5.66666651, 5.66666651]

 Más ejemplos

dot_general

Semántica

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

De manera más formal, 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, donde 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 los tipos cuantificados híbridos, 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 la compensación entre velocidad y precisión para los cálculos en los backends del acelerador. Puede ser uno de los siguientes (por el momento, la semántica de estos valores de enumeración no está especificada, pero planeamos abordar este problema 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 de punto, que también define la precisión. Si se configuran los campos de atributos del algoritmo, precision_config debe ser DEFAULT. DotAlgorithms no tienen un valor predeterminado, ya que los parámetros predeterminados se definen en la implementación. Por lo tanto, todos los campos del algoritmo de puntos pueden establecerse en None para especificar un algoritmo de puntos vacío, que, en cambio, usará el valor precision_config.

Los campos de DotAlgorithm incluyen lo siguiente:

• lhs_precision_type y rhs_precision_type, las precisiones a las que se redondean el LHD y el RHD de la operación. Los tipos de precisión son independientes de los tipos de almacenamiento de las entradas y la salida.
• accumulation_type la precisión que se usa para la acumulación.
• lhs_component_count, rhs_component_count y num_primitive_operations se aplican cuando usamos un algoritmo que descompone el LHD o el LHD en varios componentes y realiza varias operaciones de punto "primitivas" en esos valores, por lo general, para emular una mayor precisión (p. ej., Leveraging the bfloat16 Artificial Intelligence Datatype For Higher-Precision Computations: bf16_6x tf32_3x, etc.). En el caso de los algoritmos sin descomposición, estos valores deben establecerse en 1.
• allow_imprecise_accumulation para especificar si se permite la acumulación con una precisión más baja para algunos pasos (p.ej., CUBLASLT_MATMUL_DESC_FAST_ACCUM).

Ejemplo de atributos DotAlgorithm:

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


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


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

Las implementaciones deciden qué combinaciones se admiten. En general, no se garantiza que cada algoritmo sea compatible con cada tipo de acelerador por el consumidor de StableHLO. 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á la mejor verificación posible, lo que evitará algoritmos que no se sabe que sean compatibles con ningún hardware.

Consulta xla_data.proto > Algorithm para ver algunos valores de algoritmos admitidos. El ticket núm. 2483 registra el plan para crear un documento centralizado sobre los algoritmos admitidos por el 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 cuantizados, se aplica 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), haz lo siguiente:
  • (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), haz lo siguiente:
  • (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), haz lo siguiente:
  • (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 = #sta<blehlo.dot
    lhs_batching_dimensions = [0],
    rhs_batching_dimensions = [0],
    lhs_contracting_dimensions = [2],
    rhs_contracting_dimension>s = [1]
  ,
  precision_config = [<#stablehloprecisi>on DEFAULT, <#stablehloprecisi>on DEFAULT],
  algorithm = #stablehlo.dot<_algorithm
    lhs_precision_type = tf32,
    rhs_precision_type = tf32,
    accumulation_type = f32,
    lhs_component_count = 1,
    rhs_component_count = 1,
    num_primitive_operations = 1,
    allow_imprecise_accumulation >= false
  
}< : (tenso>r2x2x2xi<64, tenso>r2x>2x2xi64<) - tenso>r2x2x2xi64
// %result: [
//           [[1, 2],
//            [3, 4]],
//           [[5, 6],
//            [7, 8]]
//          ]

 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 se pueden expandir.

Entradas

Salidas

Limitaciones

• (C1) element_type(result) se calcula 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, en otros casos.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Para todo 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))).
• (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_dimensio<ns = arra>yi64: 2, 1,
  known_expanding_dimensio<ns = a>rrayi64: 0,
  known_nonexpanding_dimensio<ns = a>rrayi64: 1
}< : (ten>sor1x3xi<64, t>ens>or3xi64<) - tenso>r2x3x2xi64
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

 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