Стабильная спецификацияHLO

StableHLO — это набор операций для операций высокого уровня (HLO) в моделях машинного обучения (ML). StableHLO работает как уровень переносимости между различными платформами машинного обучения и компиляторами машинного обучения: платформы машинного обучения, создающие программы StableHLO, совместимы с компиляторами машинного обучения, которые используют программы StableHLO.

Наша цель — упростить и ускорить разработку машинного обучения за счет большей совместимости между различными платформами машинного обучения (такими как TensorFlow, JAX и PyTorch) и компиляторами машинного обучения (такими как XLA и IREE). С этой целью в этом документе представлена ​​спецификация языка программирования StableHLO.

Данная спецификация содержит три основных раздела. Во-первых, раздел «Программы» описывает структуру программ StableHLO, которые состоят из функций StableHLO, которые сами состоят из операций StableHLO. В этой структуре раздел Ops определяет семантику отдельных операций. Раздел «Выполнение» предоставляет семантику для всех этих операций, выполняемых вместе в программе. Наконец, в разделе «Обозначения» обсуждаются обозначения, используемые в спецификации.

Чтобы просмотреть спецификацию предыдущего выпуска StableHLO, откройте репозиторий интересующего выпуска с тегом . Например, StableHLO v0.19.0 Spec . Чтобы просмотреть изменения, произошедшие при каждом второстепенном обновлении версии StableHLO, обратитесь к журналу версий в VhloDialect.td .

Программы

Program ::= {Func}

Программы StableHLO состоят из произвольного количества функций StableHLO. Ниже приведен пример программы с функцией @main , которая имеет 3 входа ( %image , %weights и %bias ) и 1 выход. Тело функции имеет 6 операций.

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

Функции

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

Функции StableHLO (которые также называются именованными функциями ) имеют идентификатор, входы/выходы и тело. В будущем мы планируем ввести дополнительные метаданные для функций для достижения лучшей совместимости с HLO ( #425 , #626 , #740 , #744 ).

Идентификаторы

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

Идентификаторы StableHLO похожи на идентификаторы во многих языках программирования, но имеют две особенности: 1) все идентификаторы имеют символы, которые различают разные типы идентификаторов, 2) идентификаторы значений могут быть полностью числовыми, чтобы упростить создание программ StableHLO.

Типы

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

Типы StableHLO подразделяются на типы значений (которые также называются типами первого класса ), которые представляют значения StableHLO, и типы, не являющиеся значениями , которые описывают другие элементы программы. Типы StableHLO похожи на типы во многих языках программирования, при этом основной особенностью является предметно-ориентированный характер StableHLO, что приводит к некоторым необычным результатам (например, скалярные типы не являются типами значений).

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

Тензорные типы представляют собой тензоры, т.е. многомерные массивы. У них есть форма и тип элемента , где форма представляет собой неотрицательные или неизвестные размеры размеров в порядке возрастания соответствующих размеров (которые также называются осями ), пронумерованных от 0 до R-1 . Число измерений R называется рангом . Например, tensor<2x3xf32> — это тип тензора с формой 2x3 и типом элемента f32 . Он имеет два измерения (или, другими словами, две оси) — 0-е измерение и 1-е измерение — размеры которых равны 2 и 3. Его ранг равен 2.

Формы могут быть частично или полностью неизвестными (динамическими), например, tensor<?x2xf64> частично неизвестен, а tensor<?x?xf64> полностью неизвестен. Размеры динамических размеров обозначаются знаком ? . Фигуры не могут быть лишены ранжирования.

В будущем мы планируем изучить расширение типов тензоров за пределы размеров размеров и типов элементов, например, включив макеты ( #629 ) и разреженность ( #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

Типы квантованных элементов представляют собой целочисленные значения типа хранения в диапазоне от storage_min до storage_max (включительно), которые соответствуют значениям с плавающей запятой выраженного типа . Для данного целочисленного значения i соответствующее значение f с плавающей запятой может быть вычислено как f = (i - zero_point) * scale , где scale и zero_point называются параметрами квантования . storage_min и storage_max не являются обязательными в грамматике, но имеют значения по умолчанию min_value(storage_type) и max_value(storage_type) соответственно. Типы квантованных элементов имеют следующие ограничения:

• (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) Если is_empty(quantization_dimension) , то size(scales) = 1 .
• (C11) 0 <= quantization_dimension .

На данный момент QuantizationScale представляет собой константу с плавающей запятой, но существует большой интерес к целочисленным шкалам, представленным множителями и сдвигами. Мы планируем изучить это в ближайшем будущем ( #1404 ).

Продолжается обсуждение семантики QuantizationZeroPoint , включая тип, значения и может ли быть только одна или потенциально несколько нулевых точек в квантованном тензорном типе. По результатам этого обсуждения спецификация нулевых точек может измениться в будущем ( #1405 ).

Другое продолжающееся обсуждение касается семантики QuantizationStorageMin и QuantizationStorageMax , чтобы определить, следует ли накладывать какие-либо ограничения на эти значения и на значения квантованных тензоров ( #1406 ).

Наконец, мы планируем изучить представление неизвестных масштабов и нулевых точек аналогично тому, как мы планируем изучить представление неизвестных размеров измерений ( #1407 ).

Типы квантованных тензоров представляют собой тензоры с квантованными элементами. Эти тензоры точно такие же, как и обычные тензоры, за исключением того, что их элементы имеют типы квантованных элементов вместо обычных типов элементов.

В квантованных тензорах квантование может быть потензорным , то есть иметь один scale и zero_point для всего тензора, или может быть поосевым , то есть иметь несколько scales и zero_points , одну пару на срез определенного измерения quantization_dimension . Более формально, в тензоре t с поосевым квантованием существуют dim(t, quantization_dimension) срезы quantization_dimension : t[:, ..., 0, ..., :], t[:, ..., 1, ..., :] и т. д. Все элементы в i -м срезе используют scales[i] и zero_points[i] в ​​качестве параметров квантования. Типы квантованных тензоров имеют следующие ограничения:

• Для потензорного квантования:
  • Никаких дополнительных ограничений.
• Для поосевого квантования:
  • (C12) quantization_dimension < rank(self) .
  • (C13) dim(self, quantization_dimension) = size(scales) .

TokenType ::= 'token'

Типы токенов представляют собой токены, т.е. непрозрачные значения, создаваемые и потребляемые некоторыми операциями. Токены используются для установления порядка выполнения операций, как описано в разделе «Выполнение» .

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

Типы кортежей представляют собой кортежи, т. е. гетерогенные списки. Кортежи — это устаревшая функция, которая существует только для совместимости с HLO. В HLO кортежи используются для представления переменных входных и выходных данных. В StableHLO изначально поддерживаются переменные входные и выходные данные, и единственное использование кортежей в StableHLO — это всестороннее представление HLO ABI, где, например, T , tuple<T> и tuple<tuple<T>> могут существенно отличаться в зависимости от конкретной реализации. . В будущем мы планируем внести изменения в HLO ABI, которые могут позволить нам удалить типы кортежей из 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 ::= 'f8E3M4' | 'f8E4M3' | 'f8E4M3FN' | 'f8E4M3FNUZ' | 'f8E4M3B11FNUZ'
            | 'f8E5M2' | 'f8E5M2FNUZ' | 'bf16' | 'f16' | 'f32' | 'f64'
TensorFloat32 ::= 'tf32'
ComplexType ::= 'complex' '<' ComplexElementType '>'
ComplexElementType ::= 'f32' | 'f64'

Типы элементов представляют собой элементы тензорных типов. В отличие от многих языков программирования, эти типы не являются первоклассными в StableHLO. Это означает, что программы StableHLO не могут напрямую представлять значения этих типов (в результате идиоматично представлять скалярные значения типа T с помощью 0-мерных тензорных значений типа tensor<T> ).

• Тип Boolean представляет логические значения true и false .
• Целочисленные типы могут быть знаковыми ( si ) или беззнаковыми ( ui ) и иметь одну из поддерживаемых разрядностей ( 2 , 4 , 8 , 16 , 32 или 64 ). Знаковые типы siN представляют целочисленные значения от -2^(N-1) до 2^(N-1)-1 включительно, а беззнаковые типы uiN представляют целочисленные значения от 0 до 2^N-1 включительно.
• Типы с плавающей запятой могут быть одним из следующих:
  • f8E3M4 , f8E4M3 и f8E5M2 8-битные числа с плавающей запятой в соответствии с соглашениями IEEE-754.
  • Типы f8E4M3FN и f8E5M2 , соответствующие соответственно кодировкам E4M3 и E5M2 формата FP8, описанным в разделе «Форматы FP8 для глубокого обучения» .
  • Типы f8E4M3FNUZ и f8E5M2FNUZ , соответствующие кодировкам E4M3 и E5M2 форматов FP8, описанных в 8-битных числовых форматах для глубоких нейронных сетей .
  • Тип f8E4M3B11FNUZ , соответствующий кодировке E4M3 форматов FP8, описанных в разделе «Обучение и вывод гибридных 8-битных чисел с плавающей запятой (HFP8) для глубоких нейронных сетей» .
  • Тип bf16 , соответствующий формату bfloat16 , описанному в BFloat16: Секрет высокой производительности на Cloud TPU .
  • Типы f16 , f32 и f64 , соответствующие соответственно binary16 («половинная точность»), binary32 («одинарная точность») binary64 («двойная точность»), описанным в стандарте IEEE 754 .
  • Тип tf32 соответствует формату TensorFloat32 и имеет ограниченную поддержку в StableHLO.
• Сложные типы представляют собой комплексные значения, которые имеют действительную и мнимую части одного и того же типа элемента . Поддерживаемые сложные типы: complex<f32> (обе части имеют тип f32 ) и complex<f64> (обе части имеют тип f64 ).

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

Типы функций представляют как именованные, так и анонимные функции. У них есть типы ввода (список типов в левой части -> ) и типы вывода (список типов в правой части -> ). Во многих языках программирования типы функций являются первоклассными, но не в StableHLO.

StringType ::= 'string'

Тип String представляет собой последовательность байтов. В отличие от многих языков программирования, строковый тип не является первым классом в StableHLO и используется только для указания статических метаданных для элементов программы.

Операции

Операции StableHLO (которые также называются ops ) представляют собой закрытый набор операций высокого уровня в моделях машинного обучения. Как обсуждалось выше, синтаксис StableHLO во многом основан на MLIR, который не обязательно является наиболее эргономичной альтернативой, но, возможно, лучше всего подходит для цели StableHLO по созданию большей совместимости между платформами ML и компиляторами ML.

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

Операции StableHLO (которые также называются ops ) имеют имя, входы/выходы и подпись. Название состоит из stablehlo. префикс и мнемоника , которая однозначно идентифицирует одну из поддерживаемых операций. Ниже приведен полный список всех поддерживаемых операций.

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

Операционные операторы потребляют входные данные и производят выходные данные . Входные данные подразделяются на входные значения (вычисляемые во время выполнения), входные функции (предоставляемые статически, поскольку в StableHLO функции не являются значениями первого класса) и входные атрибуты (также предоставляемые статически). Вид входных и выходных данных, потребляемых и производимых операцией, зависит от ее мнемоники. Например, операция add потребляет 2 входных значения и производит 1 выходное значение. Для сравнения, операция select_and_scatter использует 3 входных значения, 2 входные функции и 3 входных атрибута.

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

Функции ввода (которые также называются анонимными функциями ) очень похожи на именованные функции, за исключением того, что: 1) они не имеют идентификатора (отсюда и название «анонимные»), 2) они не объявляют типы вывода (типы вывода выводится из операции return внутри функции).

Синтаксис функций ввода включает в себя неиспользуемую в настоящее время часть (см. раздел Unused » выше), которая предназначена для совместимости с MLIR. В MLIR существует более общая концепция «регионов», которая может состоять из нескольких «блоков» операций, соединенных вместе посредством операций перехода. Эти блоки имеют идентификаторы, соответствующие Unused продукции, чтобы их можно было отличить друг от друга. В StableHLO нет прыжковых операций, поэтому соответствующая часть синтаксиса MLIR не используется (но все еще существует).

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

Входные атрибуты имеют имя и значение, которое является одной из поддерживаемых констант. Они являются основным способом указания статических метаданных для элементов программы. Например, операция concatenate использует dimension атрибута, чтобы указать измерение, по которому объединяются его входные значения. Аналогично, операция slice использует несколько атрибутов, таких как start_indices и limit_indices для указания границ, которые используются для среза входного значения.

На данный момент существующие программы StableHLO иногда содержат атрибуты, не описанные в этом документе. В будущем мы планируем либо включить эти атрибуты в опсет StableHLO, либо запретить их появление в программах StableHLO. А пока вот список этих атрибутов:

• layout ( #629 ).
• mhlo.frontend_attributes ( #628 ).
• mhlo.sharding ( #619 ).
• output_operand_aliases ( #740 ).
• Метаданные местоположения ( #594 ).

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

Сигнатура операции состоит из типов всех входных значений (список типов в левой части -> ) и типов всех выходных значений (список типов в правой части -> ). Строго говоря, входные типы избыточны, и выходные типы также почти всегда избыточны (поскольку для большинства операций StableHLO типы выходных данных могут быть выведены из входных данных). Тем не менее, подпись op намеренно является частью синтаксиса StableHLO для совместимости с MLIR.

Ниже приведен пример операции, мнемоника которой — select_and_scatter . Он использует 3 входных значения ( %operand , %source и %init_value ), 2 входные функции и 3 входных атрибута ( window_dimensions , window_strides и padding ). Обратите внимание, что подпись операции включает только типы входных значений (но не типы входных функций и атрибутов, которые предоставляются в строке).

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

Константы

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

Константы StableHLO имеют литерал и тип, которые вместе представляют значение StableHLO. Обычно тип является частью синтаксиса константы, за исключением случаев, когда он однозначен (например, логическая константа однозначно имеет тип i1 , тогда как целочисленная константа может иметь несколько возможных типов).

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

Булевы константы представляют логические значения true и false . Булевы константы имеют тип i1 .

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

Целочисленные константы представляют целочисленные значения посредством строк, в которых используется десятичная или шестнадцатеричная запись. Другие системы счисления, например двоичная или восьмеричная, не поддерживаются. Целочисленные константы имеют следующие ограничения:

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

Константы с плавающей запятой представляют значения с плавающей запятой в виде строк, в которых используется десятичная или экспоненциальная запись. Кроме того, шестнадцатеричная запись может использоваться для непосредственного указания базовых битов в формате с плавающей запятой соответствующего типа. Константы с плавающей запятой имеют следующие ограничения:

• (C1) Если используется нешестнадцатеричная система записи, is_wellformed(float_literal, float_type) .
• (C2) Если используется шестнадцатеричная запись, size(hexadecimal_digits) = num_bits(float_type) / 4 .

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

Комплексные константы представляют комплексные значения с использованием списков вещественной части (идет первой) и мнимой части (идет второй). Например, (1.0, 0.0) : complex<f32> представляет 1.0 + 0.0i , а (0.0, 1.0) : complex<f32> представляет 0.0 + 1.0i . Порядок, в котором эти части затем сохраняются в памяти, определяется реализацией. Комплексные константы имеют следующие ограничения:

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

Тензорные константы представляют значения тензора с использованием вложенных списков, заданных с помощью нотации NumPy. Например, dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> представляет значение тензора со следующим сопоставлением индексов с элементами: {0, 0} => 1 , {0, 1} => 2 , {0, 2} => 3 , {1, 0} => 4 , {1, 1} => 5 , {1, 2} => 6 . Порядок, в котором эти элементы затем сохраняются в памяти, определяется реализацией. Тензорные константы имеют следующие ограничения:

• (C1) has_syntax(tensor_literal, element_type(tensor_type)) , где:
  • 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)) , где:
  • 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:]) .
  • в противном случае false .

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

Квантованные тензорные константы представляют квантованные тензорные значения с использованием тех же обозначений, что и тензорные константы, с элементами, заданными как константы их типа хранения. Квантованные тензорные константы имеют следующие ограничения:

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

Строковые литералы состоят из байтов, заданных с помощью символов ASCII и escape-последовательностей. Они не зависят от кодировки, поэтому интерпретация этих байтов определяется реализацией. Строковые литералы имеют тип string .

Операции

пресс

Семантика

Выполняет поэлементную операцию абс над тензором operand и создает тензор result . В зависимости от типа элемента выполняет следующие действия:

• Для целых чисел со знаком: целочисленный модуль.
• Для поплавков: abs из IEEE-754.
• Для комплексных чисел: комплексный модуль.
• Для квантованных типов: dequantize_op_quantize(abs, operand, type(result)) .

Входы

Выходы

Ограничения

• (C1) shape(result) = shape(operand) .
• (C2) baseline_element_type(result) определяется как:
  • complex_element_type(element_type(operand)) if is_complex(operand) .
  • baseline_element_type(operand) в противном случае.

Примеры

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

Больше примеров

добавлять

Семантика

Выполняет поэлементное сложение двух тензоров lhs и rhs и создает result тензор. В зависимости от типа элемента выполняет следующие действия:

• Для логических значений: логическое ИЛИ.
• Для целых чисел: сложение целых чисел.
• Для поплавков: addition из IEEE-754.
• Для комплексных чисел: комплексное сложение.
• Для квантованных типов: dequantize_op_quantize(add, lhs, rhs, type(result)) .

Входы

Выходы

Ограничения

• Если в операции используются неквантованные тензоры:
  • (C1) type(lhs) = type(rhs) = type(result) .
• Если в операции используются квантованные тензоры:
  • (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) Если is_per_axis_quantized(lhs) , то quantization_dimension(lhs) = quantization_dimension(result) .
  • (C7) Если is_per_axis_quantized(rhs) , то quantization_dimension(rhs) = quantization_dimension(result) .

Примеры

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

Больше примеров

после всего

Семантика

Гарантирует, что операции, производящие inputs , выполняются до выполнения любых операций, зависящих от result . Выполнение этой операции ничего не делает, она существует только для того, чтобы установить зависимости данных от result до inputs .

Входы

Выходы

Примеры

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

Больше примеров

все_собрать

Семантика

Внутри каждой группы процессов в сетке процессов StableHLO объединяет значения тензоров operands каждого процесса по all_gather_dim и создает тензоры results .

Эта операция разбивает сетку процессов StableHLO на process_groups , которые определяются следующим образом:

• cross_replica(replica_groups) если channel_id <= 0 and use_global_device_ids = false .
• cross_replica_and_partition(replica_groups) если channel_id > 0 and use_global_device_ids = false .
• flattened_ids(replica_groups) если channel_id > 0 and use_global_device_ids = true .

Затем внутри каждой process_group :

• operands...@receiver = [operand@sender for sender in process_group] для всех receiver в process_group .
• results...@process = concatenate(operands...@process, all_gather_dim) для всех process в process_group .

Входы

Выходы

Ограничения

• (C1) 0 <= all_gather_dim < rank(operands...) .
• (C2) is_unique(replica_groups) .
• (C3) size(replica_groups) определяется как:
  • num_replicas , если используется cross_replica .
  • num_replicas , если используется cross_replica_and_partition .
  • num_processes , если используется flattened_ids .
• (C4) 0 <= replica_groups < size(replica_groups) .
• (C5) Если use_global_device_ids = true , то channel_id > 0 .
• (C6) type(results...) = type(operands...) за исключением:
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1) .

Примеры

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

Больше примеров

all_reduce

Семантика

В каждой группе процессов в сетке процессов StableHLO применяет computation функции сокращения к значениям тензоров operands каждого процесса и создает тензоры results .

Эта операция разбивает сетку процессов StableHLO на process_groups , которые определяются следующим образом:

• cross_replica(replica_groups) если channel_id <= 0 and use_global_device_ids = false .
• cross_replica_and_partition(replica_groups) если channel_id > 0 and use_global_device_ids = false .
• flattened_ids(replica_groups) если channel_id > 0 and use_global_device_ids = true .

Затем внутри каждой process_group :

• results...@process[result_index] = exec(schedule) для некоторого schedule двоичного дерева, где:
  • exec(node) = computation(exec(node.left), exec(node.right)) .
  • exec(leaf) = leaf.value .
• schedule — это двоичное дерево, определяемое реализацией, порядок обхода которого равен to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0])) .

Входы

Выходы

Ограничения

• (C1) is_unique(replica_groups) .
• (C2) size(replica_groups) определяется как:
  • num_replicas , если используется cross_replica .
  • num_replicas , если используется cross_replica_and_partition .
  • num_processes , если используется flattened_ids .
• (C3) 0 <= replica_groups < size(replica_groups) .
• (C4) Если use_global_device_ids = true , то channel_id > 0 .
• (C5) computation имеет тип (tensor<E>, tensor<E>) -> (tensor<E>) где is_promotable(element_type(operand), E) .
• (C6) shape(results...) = shape(operands...) .
• (C7) element_type(results...) = E .

Примеры

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

Больше примеров

all_to_all

Семантика

Внутри каждой группы процессов в сетке процессов StableHLO значения тензоров operands разделяются по split_dimension на части, распределяются разделенные части между процессами, объединяются разбросанные части по concat_dimension и выдаются тензоры results . Эта операция разбивает сетку процессов StableHLO на process_groups , которые определяются следующим образом:

• cross_replica(replica_groups) если channel_id <= 0 .
• cross_partition(replica_groups) если channel_id > 0 .

Затем внутри каждой process_group :

• split_parts...@sender = split(operands...@sender, split_count, split_dimension) для всех sender в process_group .
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group] где receiver_index = process_group.index(receiver) .
• results...@process = concatenate(scattered_parts...@process, concat_dimension) .

Входы

Выходы

Ограничения

• (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) определяется как:
  • num_replicas , если используется cross_replica .
  • num_partitions , если используется cross_partition .
• (C7) 0 <= replica_groups < size(replica_groups) .
• (C8) dim(replica_groups, 1) = split_count .
• (C9) type(results...) = type(operands...) кроме случая, когда split_dimension != concat_dimension :
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count .
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count .

Примеры

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

Больше примеров

и

Семантика

Выполняет поэлементное «И» двух тензоров lhs и rhs и создает result тензор. В зависимости от типа элемента выполняет следующие действия:

• Для логических значений: логическое И.
• Для целых чисел: побитовое И.

Входы

Выходы

Ограничения

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

Примеры

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

Больше примеров

Атан2

Семантика

Выполняет поэлементную операцию atan2 над тензорами lhs и rhs и создает result тензор. В зависимости от типа элемента выполняет следующие действия:

• Для чисел с плавающей запятой: atan2 из IEEE-754.
• Для комплексных чисел: комплекс atan2.
• Для квантованных типов: dequantize_op_quantize(atan2, lhs, rhs, type(result)) .

Входы

Выходы

Ограничения

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

Примеры

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

Больше примеров

пакет_норм_град

Семантика

Вычисляет градиенты нескольких входных данных batch_norm_training с обратным распространением ошибки от grad_output и создает тензоры grad_operand , grad_scale и grad_offset . Более формально эту операцию можно выразить как декомпозицию существующих операций StableHLO с использованием синтаксиса 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

Для квантованных типов выполняет 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)) .

Входы

Выходы

Ограничения

• (C1) 0 <= feature_index < rank(operand) .
• (C2) operand , scale , mean , variance , grad_output , grad_operand , grad_scale и grad_offset имеют одинаковый baseline_element_type .
• (C3) operand , grad_output и grad_operand имеют одинаковую форму.
• (C4) scale , mean , variance , grad_scale и grad_offset имеют одинаковую форму.
• (C5) size(scale) = dim(operand, feature_index) .

Примеры

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

пакетная норма_inference

Семантика

Нормализует тензор operand по всем измерениям, кроме измерения feature_index , и создает result тензор. Более формально эту операцию можно выразить как декомпозицию существующих операций StableHLO с использованием синтаксиса 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)

Для квантованных типов выполняет 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)) .

Входы

Выходы

Ограничения

• (C1) 0 <= feature_index < rank(operand) .
• (C2) operand , scale , offset , mean , variance и result имеют один и тот же 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) .

Примеры

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

Семантика

Вычисляет среднее значение и дисперсию по всем измерениям, за исключением измерения feature_index , и нормализует тензор operand , создавая output , тензоры batch_mean и batch_var . Более формально эту операцию можно выразить как декомпозицию существующих операций StableHLO с использованием синтаксиса 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

Для квантованных типов выполняет 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)) .

Входы

Выходы

Ограничения

• (C1) 0 <= feature_index < rank(operand) .
• (C2) operand , scale , offset , batch_mean , batch_var и output имеют один и тот же 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) .

Примеры

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

Семантика

Выполняет операцию побитового преобразования тензора operand и создает тензор result , в котором биты всего тензора operand переинтерпретируются с использованием типа тензора result .

Более формально, учитывая E = element_type(operand) , E' = element_type(result) и R = rank(operand) :

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

bits возвращает представление данного значения в памяти, и его поведение определяется реализацией, поскольку точное представление тензоров определяется реализацией, а точное представление типов элементов также определяется реализацией.

Входы

Выходы

Ограничения

• (C1) Учитывая E = is_quantized(operand) ? storage_type(operand) : element_type(operand) , E' = is_quantized(result) ? storage_type(result) : element_type(result) и R = rank(operand) :
  • Если num_bits(E') = num_bits(E) , shape(result) = shape(operand) .
  • Если num_bits(E') < num_bits(E) :
  • rank(result) = R + 1 .
  • dim(result, i) = dim(operand, i) для всех 0 <= i < R .
  • dim(result, R) * num_bits(E') = num_bits(E) .
  • Если num_bits(E') > num_bits(E) :
  • rank(result) = R - 1 .
  • dim(result, i) = dim(operand, i) для всех 0 <= i < R .
  • dim(operand, R - 1) * num_bits(E) = num_bits(E') .
• (C2) Если is_complex(operand) or is_complex(result) , то is_complex(operand) and is_complex(result) .

Примеры

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

Больше примеров

Broadcast_in_dim

Семантика

Расширяет размеры и/или ранг входного тензора, дублируя данные в тензоре operand и дает тензор result . Более формально, result[result_index] = operand[operand_index] где для всех d в axes(operand) :

• operand_index[d] = 0 если dim(operand, d) = 1 .
• operand_index[d] = result_index[broadcast_dimensions[d]] в противном случае.

Входы

Выходы

Ограничения

• (C1) element_type(result) дается:
  • element_type(operand) , if !is_per_axis_quantized(operand) .
  • element_type(operand) за исключением того, что quantization_dimension(operand) , scales(operand) и zero_points(operand) может отличаться от quantization_dimension(result) , scales(result) и zero_points(result) соответственно, в противном случае.
• (C2) size(broadcast_dimensions) = rank(operand) .
• (C3) 0 <= broadcast_dimensions < rank(result) .
• (C4) is_unique(broadcast_dimensions) .
• (C5) Для всех d на axes(operand) :
  • dim(operand, d) = 1 или
  • dim(operand, d) = dim(result, broadcast_dimensions[d]) .
• (C6) Если is_per_axis_quantized(result) :
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)] .
  • Если dim(operand, quantization_dimension(operand)) = 1 , то scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))) .

Примеры

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

Больше примеров

случай

Семантика

Производит выход из выполнения ровно одной функции из branches в зависимости от значения index . Более формально, result = selected_branch() где:

• selected_branch = branches[index] если 0 <= index < size(branches) .
• selected_branch = branches[-1] в противном случае.

Входы

Выходы

Ограничения

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

Примеры

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

Больше примеров

CBRT

Семантика

Выполняет элементную работу кубического корня на тензоре operand и получает тензор result . В зависимости от типа элемента, делает следующее:

• Для поплавков: rootn(x, 3) от IEEE-754.
• Для сложных чисел: сложный кубический корень.
• Для квантовых типов: dequantize_op_quantize(cbrt, operand, type(result))

Входы

Выходы

Ограничения

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

Примеры

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

Больше примеров

клетка

Семантика

Выполняет элементный Ceil из operand Tensor и дает тензор result . Реализует операцию roundToIntegralTowardPositive из спецификации IEEE-754. Для квантовых типов выполняет dequantize_op_quantize(ceil, operand, type(result)) .

Входы

Выходы

Ограничения

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

Примеры

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

Больше примеров

Чолский

Семантика

Вычисляет размесщение хоузского партии матриц.

Более формально, для всех i в index_space(result) , result[i0, ..., iR-3, :, :] -это хоулесское разложение a[i0, ..., iR-3, :, :] , в форме любой из нижней триангулярной (если lower true или верхнего триангулярного (если lower false ) матрица. Выходные значения в противоположном треугольнике, то есть соответственно строгий верхний треугольник или строгий нижний треугольник, определяются реализацией.

Если существует i , где входная матрица не является матрицей с положительным определением гермиты, то поведение не определен.

Для квантованных типов выполняет dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)) .

Входы

Выходы

Ограничения

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

Примеры

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

зажим

Семантика

Зажимая каждый элемент тензора operand между минимальным и максимальным значением и дает тензор result . Более формально, result[result_index] = minimum(maximum(operand[result_index], min_element), max_element) , где min_element = rank(min) = 0 ? min[] : min[result_index] , max_element = rank(max) = 0 ? max[] : max[result_index] . Для квантованных типов выполняет dequantize_op_quantize(clamp, min, operand, max, type(result)) .

Введение заказа на сложные числа включает в себя удивительную семантику, поэтому в будущем мы планируем удалить поддержку комплексных чисел для этой операции ( #560 ).

Входы

Выходы

Ограничения

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

Примеры

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

Больше примеров

Collective_broadcast

Семантика

В рамках каждой группы процессов в сетке процесса StableHlo отправьте значение тензора operand из исходного процесса в целевые процессы и дают тензор result .

Операция разбивает сетку процесса stablehlo в process_groups , которая определяется следующим образом:

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

После этого result@process определяется:

• operand@process_groups[i, 0] если существует i так что процесс находится в process_groups[i] .
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) в противном случае.

Входы

Выходы

Ограничения

• (C1) is_unique(replica_groups) .
• (C2) 0 <= replica_groups < N где N определяется как:
  • num_replicas , если используется cross_replica .
  • num_partitions , если используется cross_partition .
• (C3) type(result) = type(operand) .

Примеры

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

Семантика

В рамках каждой группы процессов в сетке процесса StableHLO отправляет значение тензора operand из исходного процесса в целевой процесс и дает тензор result .

Операция разбивает сетку процесса stablehlo в process_groups , которая определяется следующим образом:

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

После этого result@process определяется:

• operand@process_groups[i, 0] , если существует i , что process_groups[i, 1] = process .
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) в противном случае.

Входы

Выходы

Ограничения

• (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 , где N определяется как:
  • num_replicas , если используется cross_replica .
  • num_partitions , если используется cross_partition .
• (C5) type(result) = type(operand) .

Примеры

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

Больше примеров

сравнивать

Семантика

Проводит элементное сравнение тензоров lhs и rhs в соответствии с comparison_direction и compare_type и дает result тензора.

Значения comparison_direction и compare_type имеют следующую семантику:

Для логических и целочисленных типов элементов:

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

Для типов элементов с плавающей точкой с compare_type = FLOAT , OP реализует следующие операции IEEE-754:

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

Для типов элементов с плавающей точкой с compare_type = TOTALORDER в OP используется комбинация totalOrder и compareQuietEqual операций с IEEE-754.

Для комплексных типов элементов лексикографическое сравнение (real, imag) пар выполняется с использованием предоставленного comparison_direction и compare_type . Введение заказа на комплексные числа включает в себя удивительную семантику, поэтому в будущем мы планируем удалить поддержку комплексных чисел, когда comparison_direction - GE , GT , LE или LT ( #560 ).

Для квантованных типов. выполняет dequantize_compare(lhs, rhs, comparison_direction) .

Входы

Выходы

Ограничения

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs) .
• (C2) shape(lhs) = shape(rhs) = shape(result) .
• (C3) compare_type определяется как:
  • SIGNED , если is_signed_integer(element_type(lhs)) .
  • UNSIGNED if is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)) .
  • FLOAT или TOTALORDER If is_float(element_type(lhs)) .
  • FLOAT If is_complex(element_type(lhs)) .

Примеры

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

Больше примеров

сложный

Семантика

Выполняет элементное преобразование в сложное значение из пары реальных и воображаемых значений, lhs и rhs , и дает тензор result .

Входы

Выходы

Ограничения

• (C1) type(lhs) = type(rhs) .
• (C2) shape(result) = shape(lhs) .
• (C3) element_type(result) имеет complex<E> где E = element_type(lhs) .

Примеры

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

Больше примеров

композитный

Семантика

Инкапсулирует операцию, выполненную (составленную) других операций StableHlo, принимая inputs и composite_attributes и дает results . Семантика OP реализована атрибутом decomposition . composite OP может быть заменен его разложением без изменения семантики программы. В тех случаях, когда внедрение разложения не обеспечивает ту же операционную семантику, предпочитаю использовать custom_call .

Поле version (по умолчанию к 0 ) используется для обозначения при изменении семантики композита.

Входы

Выходы

Ограничения

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

Примеры

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

Больше примеров

объединять

Семантика

Конкатенаты inputs вдоль измерения dimension в том же порядке, что и заданные аргументы, и дает тензор result . Более формально, result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1] , где:

1. id = d0 + ... + dk-1 + kd .
2. d равен dimension и d0 , ... являются d inputs размеров.

Входы

Выходы

Ограничения

• (C1) same(element_type(inputs...)) .
• (C2) same(shape(inputs...)) за исключением 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]) за исключением:
  • dim(result, dimension) = dim(inputs[0], dimension) + ...

Примеры

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

Больше примеров

постоянный

Семантика

Производит output тензор из постоянного value .

Входы

Выходы

Ограничения

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

Примеры

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

Больше примеров

конвертировать

Семантика

Выполняет элементное преобразование от одного типа элемента в другой на operand тензоре и дает тензор result .

Для конверсий с логическим и всем, поддерживаемым всем , значение false преобразуется в ноль, а значение true преобразуется в одно. Для конверсий с любым поддерживаемым типом в булавицах нулевое значение преобразуется в false , а ненулевые значения преобразуются в true . Смотрите ниже, как это работает для сложных типов.

Для преобразования, включающих целое число в интеллектуальную точку , точку с целым числом или точка с плавающей точкой с плавающей точкой , если исходное значение может быть точно представлено в типе назначения, значение результата-это точное представление. В противном случае поведение - TBD ( #180 ).

Для преобразования, связанных с плавающей точкой в ​​интеграцию , дробная часть усечена. Если усеченное значение не может быть представлено в типе назначения, поведение - TBD ( #180 ).

Преобразование, включающее комплекс в комплекс, следуют тому же поведению с плавающей точкой в ​​точку, которая для преобразования реальных и воображаемых частей.

Для преобразований из комплекса-и-другого и другого типа в любое другое тип в-закол в переносном значении игнорируется или воображаемое значение назначения, соответственно. Преобразование реальной части следует переоборудованию с плавающей точкой.

В принципе, эта операция может выразить декантизацию (превращение из квантованных тензоров в обычные тензоры), квантование (превращение от обычных тензоров в квантовые тензоры) и регистрация (конверсия между квантованными тензорами), но в настоящее время мы uniform_dequantize операции для этого Первый вариант использования и uniform_quantize для второго и третьего варианта использования. В будущем эти два оператора могут быть объединены в convert ( #1576 ).

Входы

Выходы

Ограничения

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

Примеры

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

Больше примеров

сверток

Семантика

Вычисляет точечные продукты между окнами lhs и ломтиками rhs и дает result . На следующей диаграмме показано, как элементы в result вычисляются из lhs и rhs с использованием конкретного примера.

Более формально рассмотрите следующее переосмысление входов с точки зрения lhs , чтобы иметь возможность выражать окна 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) .

В этом переосмыслении используются следующие вспомогательные функции:

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

If feature_group_count = 1 и batch_group_count = 1 , тогда для всех output_spatial_index в index_space(dim(result, output_spatial_dimensions...)) , result[result_shape(:, output_spatial_index, :)] = dot_product где:

• 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]) . Эта функция, по -видимому, не используется, поэтому в будущем мы планируем ее удалить ( #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]) .

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

Если 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) .

Для квантованных типов выполняет 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)) .

Для гибридных квантовых типов выполняет 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) .

Входы

Выходы

Ограничения

• (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) Учитывает 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) Учитывая 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) Учитывает 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) определяется как:
  • dim(lhs, input_batch_dimension) / batch_group_count if result_dim = output_batch_dimension .
  • dim(rhs, kernel_output_feature_dimension) если result_dim = output_feature_dimension .
  • num_windows иначе, где:
  • 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 .
• Если операция использует некванизированные тензоры:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result) .
• Если операция использует квантовые тензоры:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs) .
  • (C29) if is_per_axis_quantized(rhs) , то quantization_dimension(rhs) = kernel_output_feature_dimension .
  • (C30) Если is_per_axis_quantized(result) , то quantization_dimension(result) = output_feature_dimension .
  • Если is_quantized(lhs) :
  • (C31) storage_type(lhs) = storage_type(rhs) .
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C33) Если is_per_tensor_quantized(rhs) , то is_per_tensor_quantized(result) .
  • Если !is_quantized(lhs) :
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result) .

Примеры

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

Больше примеров

косинус

Семантика

Выполняет элементную эксплуатацию косинус на тензоре operand и дает тензор result . В зависимости от типа элемента, делает следующее:

• Для поплавков: cos от IEEE-754.
• Для сложных чисел: сложный косинус.
• Для квантовых типов: dequantize_op_quantize(cosine, operand, type(result)) .

Входы

Выходы

Ограничения

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

Примеры

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

Больше примеров

count_leading_zeros

Семантика

Выполняет количество элементов количества ведущих нулевых бит в тензоре operand и дает тензор result .

Входы

Выходы

Ограничения

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

Примеры

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

Больше примеров

custom_call

Семантика

Инкапсулирует определяемую реализацией операцию call_target_name , который принимает inputs и called_computations и дает results . has_side_effect , backend_config и api_version могут использоваться для обеспечения дополнительных определяемых реализации метаданных.

На данный момент эта операция содержит довольно неорганизованную коллекцию метаданных, которая отражает органическую эволюцию своей операции аналога в компиляторе XLA. В будущем мы планируем объединить эти метаданные ( #741 ).

Входы

Выходы

Примеры

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

разделять

Семантика

Выполняет элементное разделение дивидендных lhs и делительных тензоров rhs и дает тензор result . В зависимости от типа элемента, делает следующее:

• Для целых чисел: целочисленное разделение, которое производит алгебраический коэффициент с любой дробной частью отброшенной.
• Для поплавков: division от IEEE-754.
• Для сложных чисел: сложное разделение.
• Для квантованных типов:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)) .

Входы

Выходы

Ограничения

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

Примеры

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

Больше примеров

dot_general

Семантика

Вычисляет точечные продукты между ломтиками lhs и ломтиками rhs и дает тензор result .

Более формально, result[result_index] = dot_product , где:

• 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 где size(result_batching_index) = size(lhs_batching_dimensions) , size(result_lhs_index) = size(lhs_result_dimensions) и 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)) .

Для квантованных типов выполняет 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))

For hybrid quantized types, performs 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 контролирует компромисс между скоростью и точностью для вычислений на бэкэндах ускорителя. Это может быть одно из следующих (в настоящее время семантика этих значений перечисления недостаточно, но мы планируем решить это в #755 ):

• DEFAULT : самый быстрый расчет, но наименьшее точное приближение к исходному номеру.
• HIGH : более медленный расчет, но более точное приближение к исходному числу.
• HIGHEST : самый медленный расчет, но наиболее точное приближение к исходному числу.

DotAlgorithm определяет основные свойства алгоритма, используемого для реализации операции DOT, что также определяет точность. Если установлены поля атрибута алгоритма, то precision_config должен быть DEFAULT . DotAlgorithms не имеют значения по умолчанию, так как параметры по умолчанию определяются реализацией. Таким образом, все поля алгоритма тота могут быть установлены на None чтобы указать пустой точечный алгоритм, который вместо этого будет использовать значение precision_config .

Поля DotAlgorithm включают:

• lhs_precision_type и rhs_precision_type , точности, к которым приводятся LHS и RHS операции. Типы точности не зависят от типов хранения входов и вывода.
• accumulation_type точность, используемая для накопления.
• lhs_component_count , rhs_component_count и num_primitive_operations применяются, когда мы делаем алгоритм, который разлагает LHS и/или RHS в несколько компонентов и выполняет множественные «примитивные» точечные операции на этих значениях - обычно для эмуляции более высокой точности (EG LEARVERGEATERSTIONSIOL INTLATIONENTION DATATIPERENTIOVERENTIONENTIONERENTIOVERENTIONSTIONSIO Для вычислений с более высокой конкретной позицией : BF16_6X TF32_3X и т. Д.). Для алгоритмов без разложения эти значения должны быть установлены на 1 .
• allow_imprecise_accumulation для указания того, разрешено ли накопление в более низкой точности для некоторых шагов (например, CUBLASLT_MATMUL_DESC_FAST_ACCUM ).

Пример атрибутов 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}

Реализации должны решить, какие комбинации поддерживаются. В целом, не гарантируется, что каждый алгоритм поддерживается на каждом типе ускорителя потребителем стабильного. Если заданный алгоритм не поддерживается, должна быть вызвана ошибка, в отличие от падения к альтернативе. Проверка StableHLO обеспечит наилучшую проверку усилий, предотвращая алгоритмы, которые, как известно, не поддерживаются на любом оборудовании.

См. xla_data.proto > Algorithm для некоторых поддерживаемых значений алгоритма. Билет № 2483 фиксирует план по созданию централизованного документа на алгоритмах поддерживаемых по бэкэнд.

Входы

Выходы

Ограничения

• (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) .
• Если операция использует некванизированные тензоры:
  • (C13) element_type(lhs) = element_type(rhs) .
• Если операция использует квантовые тензоры:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs) .
  • (C15) zero_points(rhs) = 0 .
  • (C16) If is_per_axis_quantized(rhs) , then quantization_dimension(rhs) not in rhs_contracting_dimensions .
  • If is_quantized(lhs) :
  • (C17) storage_type(lhs) = storage_type(rhs) .
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C19) If is_per_tensor_quantized(rhs) , then is_per_tensor_quantized(result) .
  • If !is_quantized(lhs) :
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result) .
• If !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 .

Примеры

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

More Examples

dynamic_broadcast_in_dim

Семантика

This operation is functionally identical to broadcast_in_dim op, but the result shape is specified dynamically via output_dimensions .

The operation also accepts optional attributes known_expanding_dimensions , known_nonexpanding_dimensions to express static knowledge about the expanding behavior of dimensions. If not specified, all dimensions are assumed to be possibly expanding.

Входы

Выходы

Constraints

• (C1) element_type(result) is given by:
  • element_type(operand) , if !is_per_axis_quantized(operand) .
  • element_type(operand) except that quantization_dimension(operand) , scales(operand) , and zero_points(operand) may differ from quantization_dimension(result) , scales(result) , and zero_points(result) resp., otherwise.
• (C2) size(broadcast_dimensions) = rank(operand) .
• (C3) 0 <= broadcast_dimensions < rank(result) .
• (C4) is_unique(broadcast_dimensions) .
• (C5) For all d in axes(operand) :
  • dim(operand, d) = 1 or
  • dim(operand, d) = dim(result, broadcast_dimensions[d]) .
• (C6) If is_per_axis_quantized(result) :
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)] .
  • If dim(operand, quantization_dimension(operand)) = 1 , then scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))) .
• (C7) size(output_dimensions) = rank(result) .
• (C8) is_unique(known_expanding_dimensions + known_nonexpanding_dimensions) .
• (C9) 0 <= known_expanding_dimensions < rank(operand) .
• (C10) 0 <= known_nonexpanding_dimensions < rank(operand) .

Примеры

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

More Examples

dynamic_conv

Семантика

This operation is functionally identical to convolution op, but the padding is specified dynamically via padding .

Входы

Выходы

Constraints

• (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) Given 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) Given 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) Given 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) is defined as:
  • dim(lhs, input_batch_dimension) / batch_group_count if result_dim = output_batch_dimension .
  • dim(rhs, kernel_output_feature_dimension) if result_dim = output_feature_dimension .
  • num_windows otherwise, where:
  • 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 .
• If the operation uses non-quantized tensors:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result) .
• If the operation uses quantized tensors:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs) .
  • (C29) If is_per_axis_quantized(rhs) , then quantization_dimension(rhs) = kernel_output_feature_dimension .
  • (C30) If is_per_axis_quantized(result) , then quantization_dimension(result) = output_feature_dimension .
  • If is_quantized(lhs) :
  • (C31) storage_type(lhs) = storage_type(rhs) .
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C33) If is_per_tensor_quantized(rhs) , then is_per_tensor_quantized(result) .
  • If !is_quantized(lhs) :
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result) .

Примеры

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