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

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

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

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

Программы

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 ::= IntegerConstant
QuantizationStorageMax ::= IntegerConstant
QuantizationExpressedType ::= FloatType
QuantizationDimension ::= IntegerConstant
QuantizationParameters ::= QuantizationParameter
                         | '{' QuantizationParameter {',' QuantizationParameter} '}'
QuantizationParameter ::= QuantizationScale ':' QuantizationZeroPoint
QuantizationScale ::= FloatConstant
QuantizationZeroPoint ::= IntegerConstant

Типы квантованных элементов представляют собой целочисленные значения типа хранения в диапазоне от 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] в ​​качестве параметров квантования. Типы квантованных тензоров имеют следующие ограничения:

• Для потензорного квантования:
  • Никаких дополнительных ограничений.
• Для поосевого квантования:
  • (C13) quantization_dimension < rank(self) .
  • (C14) 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 ::= 'si4' | 'si8' | 'si16' | 'si32' | 'si64'
UnsignedIntegerType ::= 'ui4' | 'ui8' | 'ui16' | 'ui32' | 'ui64'
FloatType ::= 'f8E4M3FN' | 'f8E5M2' | 'f8E4M3FNUZ' | 'f8E5M2FNUZ'
            | 'f8E4M3B11FNUZ' | 'bf16' | 'f16' | 'f32' | 'f64'
ComplexType ::= 'complex' '<' ComplexElementType '>'
ComplexElementType ::= 'f32' | 'f64'

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

• Тип Boolean представляет логические значения true и false .
• Целочисленные типы могут быть знаковыми ( si ) или беззнаковыми ( ui ) и иметь одну из поддерживаемых разрядностей ( 4 , 8 , 16 , 32 или 64 ). Знаковые типы siN представляют целые значения от -2^(N-1) до 2^(N-1)-1 включительно, а беззнаковые типы uiN представляют целочисленные значения от 0 до 2^N-1 включительно.
• Типы с плавающей запятой могут быть одним из следующих:
  • Типы 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 .
• Сложные типы представляют собой комплексные значения, которые имеют действительную и мнимую части одного и того же типа элемента . Поддерживаемые сложные типы: 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. префикс и мнемоника , которая однозначно идентифицирует одну из поддерживаемых операций. Ниже приведен полный список всех поддерживаемых операций.

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

• builtin.module , func.func , func.call и func.return ( #425 ).
• операции chlo ( #602 ).
• Категория операций StableHLO «Не в HLO» — изначально они были частью опсета StableHLO, но позже было сочтено, что они ему не подходят: broadcast , create_token , cross-replica-sum , dot , einsum , torch_index_select , unary_einsum ( #3 ). .
• Категория «Динамизм» операций StableHLO — они были загружены из MHLO, и мы находимся в процессе их спецификации: real_dynamic_slice , set_dimension_size . ( #8 ).
• Вычисления формы, включая arith , shape и tensor операции ( #8 ).

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 .

Операции

пресс

Семантика

Выполняет поэлементную операцию abs над тензором 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)) если 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 объединяет значения тензора operand каждого процесса по all_gather_dim и создает тензор result .

Эта операция разбивает сетку процессов 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 .
• result@process = concatenate(operands@process, all_gather_dim) для всех process process_group .

Входы

Выходы

Ограничения

• (C1) 0 <= all_gather_dim < rank(operand) .
• (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(result) = type(operand) за исключением:
  • dim(result, all_gather_dim) = dim(operand, all_gather_dim) * dim(process_groups, 1) .

Примеры

// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [[1, 2], [3, 4]]
// %operand@(1, 0): [[5, 6], [7, 8]]
%result = "stablehlo.all_gather"(%operand) {
  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<2x4xi64>
// %result@(0, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result@(1, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]

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

all_reduce

Семантика

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

Эта операция разбивает сетку процессов 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 :

• result@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(result) = shape(operand) .
• (C7) element_type(result) = E .

Примеры

// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [1, 2, 3, 4]
// %operand@(1, 0): [5, 6, 7, 8]
%result = "stablehlo.all_reduce"(%operand) ({
  ^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_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor<4xi64>) -> tensor<4xi64>
// %result@(0, 0): [6, 8, 10, 12]
// %result@(1, 0): [6, 8, 10, 12]

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

all_to_all

Семантика

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

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

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

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

• split_parts@sender = split(operand@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) .
• result@process = concatenate(scattered_parts@process, concat_dimension) .

Входы

Выходы

Ограничения

• (C1) 0 <= split_dimension < rank(operand) .
• (C2) dim(operand, split_dimension) % split_count = 0 .
• (C3) 0 <= concat_dimension < rank(operand) .
• (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(result) = type(operand) за исключением:
  • dim(result, split_dimension) = dim(operand, split_dimension) / split_count .
  • dim(result, concat_dimension) = dim(operand, concat_dimension) * split_count .

Примеры

// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [[1, 2, 3, 4],
//                   [5, 6, 7, 8]]
// %operand@(1, 0): [[9, 10, 11, 12],
//                   [13, 14, 15, 16]]
%result = "stablehlo.all_to_all"(%operand) {
  split_dimension = 1 : i64,
  concat_dimension = 0 : i64,
  split_count = 2 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
} : (tensor<2x4xi64>) -> tensor<4x2xi64>
// %result@(0, 0): [[1, 2],
//                  [5, 6],
//                  [9, 10],
//                  [13, 14]]
// %result@(1, 0): [[3, 4],
//                  [7, 8],
//                  [11, 12],
//                  [15, 16]]

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

и

Семантика

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

Для комплексных типов элементов лексикографическое сравнение (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",
  api_version = 1 : 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))

Для гибридных квантовых типов выполняет 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 : самый медленный расчет, но наиболее точное приближение к исходному числу.

Входы

Выходы

Ограничения

• (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) Если is_per_axis_quantized(rhs) , то quantization_dimension(rhs) не в rhs_contracting_dimensions .
  • Если is_quantized(lhs) :
  • (C17) storage_type(lhs) = storage_type(rhs) .
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C19) Если is_per_tensor_quantized(rhs) , то is_per_tensor_quantized(result) .
  • Если !is_quantized(lhs) :
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result) .

Примеры

// %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>]
} : (tensor<2x2x2xi64>, tensor<2x2x2xi64>) -> tensor<2x2x2xi64>
// %result: [
//           [[1, 2],
//            [3, 4]],
//           [[5, 6],
//            [7, 8]]
//          ]

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

dynamic_broadcast_in_dim

Семантика

Эта операция функционально идентична Broadcast_in_dim OP, но форма результата определяется динамически через output_dimensions .

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

Входы

Выходы

Ограничения

• (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))) .
• (C7) size(output_dimensions) = rank(result) .
• (C8) is_unique(known_expanding_dimensions + known_non_expanding_dimensions) .
• (C9) 0 <= known_expanding_dimensions < rank(operand) .
• (C10) 0 <= known_non_expanding_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_non_expanding_dimensions = array<i64: 1>
} : (tensor<1x3xi64>, tensor<3xi64>) -> tensor<2x3x2xi64>
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

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

Dynamic_conv

Семантика

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

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

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

If feature_group_count = 1 and batch_group_count = 1 , then for all output_spatial_index in index_space(dim(result, output_spatial_dimensions...)) , result[result_shape(:, output_spatial_index, :)] = dot_product where:

• 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]) . This feature appears to be unused, so in the future we are planning to remove it ( #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) .

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

For quantized types, performs dequantize_op_quantize( lambda lhs, rhs: convolution(lhs, rhs, d_padding, window_strides, 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)) .

For hybrid quantized types, performs hybrid_dequantize_then_op( lambda lhs, rhs: convolution(lhs, rhs, d_padding, window_strides, 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(d_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) rank(result) = N .
• If the operation uses non-quantized tensors:
  • (C26) element_type(lhs) = element_type(rhs) = element_type(result) .
• If the operation uses quantized tensors:
  • (C27) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs) .
  • (C28) If is_per_axis_quantized(rhs) , then quantization_dimension(rhs) = kernel_output_feature_dimension .
  • (C29) If is_per_axis_quantized(result) , then quantization_dimension(result) = output_feature_dimension .
  • If is_quantized(lhs) :
  • (C30) storage_type(lhs) = storage_type(rhs) .
  • (C31) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C32) If is_per_tensor_quantized(rhs) , then is_per_tensor_quantized(result) .
  • If !is_quantized(lhs) :
  • (C33) 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]]
//          ]]

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

dynamic_gather

Семантика

This operation is functionally identical to iota op, with the slice_sizes specified dynamically as a value.

Входы

Выходы

Ограничения

• (C1) rank(operand) = size(offset_dims) + size(collapsed_slice_dims) .
• (C2) 0 <= index_vector_dim <= rank(start_indices) .
• (C3) size(start_index_map) = index_vector_dim < rank(start_indices) ? dim(start_indices, index_vector_dim) : 1 .
• (C4) is_unique(offset_dims) and is_sorted(offset_dims) .
• (C5) 0 <= offset_dims < rank(result) .
• (C6) is_unique(collapsed_slice_dims) and is_sorted(collapsed_slice_dims) .
• (C7) 0 <= collapsed_slice_dims < rank(operand) .
• (C8) slice_sizes[collapsed_slice_dims...] <= 1 .
• (C9) is_unique(start_index_map) .
• (C10) 0 <= start_index_map < rank(operand) .
• (C11) size(slice_sizes) = rank(operand) .
• (C12) 0 <= slice_sizes <= shape(operand) .
• (C13) shape(result) = combine(batch_dim_sizes, offset_dim_sizes) where:
  • batch_dim_sizes = shape(start_indices) except that the dimension size of start_indices corresponding to index_vector_dim is not included.
  • offset_dim_sizes = shape(slice_sizes) except that the dimension sizes in slice_sizes corresponding to collapsed_slice_dims are not included.
  • combine puts batch_dim_sizes at axes corresponding to batch_dims and offset_dim_sizes at axes corresponding to offset_dims .
• (C14) element_type(operand) = element_type(result) .

Примеры

// %operand: [
//            [[1, 2], [3, 4], [5, 6], [7, 8]],
//            [[9, 10],[11, 12], [13, 14], [15, 16]],
//            [[17, 18], [19, 20], [21, 22], [23, 24]]
//           ]
// %start_indices: [
//                  [[0, 0], [1, 0], [2, 1]],
//                  [[0, 1], [1, 1], [0, 2]]
//                 ]
// %slize_sizes: [1, 2, 2]
%result = "stablehlo.dynamic_gather"(%operand, %start_indices, %slize_sizes) {
  dimension_numbers = #stablehlo.gather<
    offset_dims = [2, 3],
    collapsed_slice_dims = [0],
    start_index_map = [1, 0],
    index_vector_dim = 2>,
  indices_are_sorted = false
} : (tensor<3x4x2xi64>, tensor<2x3x2xi64>, tensor<3xi64>) -> tensor<2x3x2x2xi64>
// %result: [
//            [
//              [[1, 2], [3, 4]],
//              [[3, 4], [5, 6]],
//              [[13, 14], [15, 16]]
//            ],
//            [
//              [[9, 10], [11, 12]],
//              [[11, 12], [13, 14]],
//              [[17, 18], [19, 20]]
//            ]
//          ]

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

dynamic_iota

Семантика

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

Входы

Выходы

Ограничения

• (C1) 0 <= iota_dimension < size(output_shape) .
• (C2) rank(result) = size(output_shape) .

Примеры

%output_shape = stablehlo.constant dense<[4, 5]> : tensor<2xi64>
%result = "stablehlo.dynamic_iota"(%output_shape) {
  iota_dimension = 0 : i64
} : (tensor<2xi64>) -> tensor<4x5xi64>
// %result: [
//           [0, 0, 0, 0, 0],
//           [1, 1, 1, 1, 1],
//           [2, 2, 2, 2, 2],
//           [3, 3, 3, 3, 3]
//          ]

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

dynamic_pad

Семантика

This operation is functionally identical to pad op, but with edge_padding_low , edge_padding_high , and interior_padding specified dynamically as values.

Входы

Выходы

Ограничения

• (C1) element_type(operand) = element_type(padding_value) = element_type(result) .
• (C2) size(edge_padding_low) = size(edge_padding_high) = size(interior_padding) = rank(operand) .
• (C3) 0 <= interior_padding .
• (C4) shape(result) = shape(operand) + edge_padding_low + max(shape(operand) - 1, 0) * interior_padding + edge_padding_high .

Примеры

// %operand: [
//            [1, 2, 3],
//            [4, 5, 6]
//           ]
// %padding_value: 0
// %edge_padding_low: [0, 1]
// %edge_padding_high: [2, 1]
// %interior_padding: [1, 2]
%result = "stablehlo.dynamic_pad"(%operand, %padding_value,
  %edge_padding_low, %edge_padding_high, %interior_padding
) : (tensor<2x3xi64>, tensor<i64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> tensor<5x9xi64>
// %result: [
//           [0, 1, 0, 0, 2, 0, 0, 3, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0],
//           [0, 4, 0, 0, 5, 0, 0, 6, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0]
//          ]

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

dynamic_reshape

Семантика

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

Входы

Выходы

Ограничения

• (C1) element_type(result) is given by:
  • element_type(operand) , if !is_per_axis_quantized(operand) .
  • element_type(operand) except that quantization_dimension(operand) and quantization_dimension(result) may differ, otherwise.
• (C2) size(operand) = size(result) .
• (C3) If is_per_axis_quantized(operand) :
  • reduce(dims(operand, [0, 1, ..., quantization_dimension(operand) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) = reduce(dims(result, [0, 1, ..., quantization_dimension(result) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) .
  • dim(operand, quantization_dimension(operand)) = dim(result, quantization_dimension(result)) .
  • reduce(dims(operand, [quantization_dimension(operand) + 1, ..., rank(operand) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) = reduce(dims(result, [quantization_dimension(result) + 1, ..., rank(result) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) .
• (C4) size(output_shape) = rank(result) .

Примеры

// %operand: [[1, 2, 3], [4, 5, 6]]
%output_shape = stablehlo.constant dense<[3, 2]> : tensor<2xi64>
%result = "stablehlo.dynamic_reshape"(%operand, %output_shape) : (tensor<2x3xi64>, tensor<2xi64>) -> tensor<3x2xi64>
// %result: [[1, 2], [3, 4], [5, 6]]

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

dynamic_slice

Семантика

Extracts a slice from the operand using dynamically-computed starting indices and produces a result tensor. start_indices contain the starting indices of the slice for each dimension subject to potential adjustment, and slice_sizes contain the sizes of the slice for each dimension. More formally, result[result_index] = operand[operand_index] where:

• adjusted_start_indices = clamp(0, start_indices, shape(operand) - slice_sizes) .
• operand_index = adjusted_start_indices + result_index .

Входы

Выходы

Ограничения

• (C1) element_type(operand) = element_type(result) .
• (C2) size(start_indices) = size(slice_sizes) = rank(operand) .
• (C3) same(type(start_indices...)) .
• (C4) 0 <= slice_sizes <= shape(operand) .
• (C5) shape(result) = slice_sizes .

Примеры

// %operand: [
//            [0, 0, 1, 1],
//            [0, 0, 1, 1],
//            [0, 0, 0, 0],
//            [0, 0, 0, 0]
//           ]
// %start_indices0: -1
// %start_indices1: 3
%result = "stablehlo.dynamic_slice"(%operand, %start_indices0, %start_indices1) {
  slice_sizes = array<i64: 2, 2>
} : (tensor<4x4xi32>, tensor<i64>, tensor<i64>) -> tensor<2x2xi32>
// %result: [
//           [1, 1],
//           [1, 1]
//          ]

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

dynamic_update_slice

Семантика

Produces a result tensor which is equal to the operand tensor except that the slice starting at start_indices is updated with the values in update . More formally, result[result_index] is defined as:

• update[update_index] if 0 <= update_index < shape(update) where:
  • adjusted_start_indices = clamp(0, start_indices, shape(operand) - shape(update)) .
  • update_index = result_index - adjusted_start_indices .
• operand[result_index] otherwise.

Входы

Выходы

Ограничения

• (C1) type(operand) = type(result) .
• (C2) element_type(update) = element_type(operand) .
• (C3) rank(update) = rank(operand) .
• (C4) size(start_indices) = rank(operand) .
• (C5) same(type(start_indices...)) .
• (C6) 0 <= shape(update) <= shape(operand) .

Примеры

// %operand: [
//            [1, 1, 0, 0],
//            [1, 1, 0, 0],
//            [1, 1, 1, 1],
//            [1, 1, 1, 1]
//           ]
// %update: [
//           [1, 1],
//           [1, 1]
//          ]
// %start_indices0: -1
// %start_indices1: 3
%result = "stablehlo.dynamic_update_slice"(%operand, %update, %start_indices0, %start_indices1)
  : (tensor<4x4xi32>, tensor<2x2xi32>, tensor<i64>, tensor<i64>) -> tensor<4x4xi32>
// %result: [
//           [1, 1, 1, 1],
//           [1, 1, 1, 1],
//           [1, 1, 1, 1],
//           [1, 1, 1, 1]
//          ]

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

exponential

Семантика

Performs element-wise exponential operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: exp from IEEE-754.
• For complex numbers: complex exponential.
• For quantized types: dequantize_op_quantize(exponential, operand, type(result)) .

Входы

Выходы

Ограничения

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

Примеры

// %operand: [[0.0, 1.0], [2.0, 3.0]]
%result = "stablehlo.exponential"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
// %result: [[1.0, 2.7182818284590451], [7.3890560989306504, 20.085536923187668]]

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

exponential_minus_one

Семантика

Performs element-wise exponential minus one operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: expm1 from IEEE-754.
• For complex numbers: complex exponential minus one.
• For quantized types: dequantize_op_quantize(exponential_minus_one, operand, type(result)) .

Входы

Выходы

Ограничения

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

Примеры

// %operand: [0.0, 1.0]
%result = "stablehlo.exponential_minus_one"(%operand) : (tensor<2xf64>) -> tensor<2xf64>
// %result: [0.0, 1.71828187]

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

fft

Семантика

Performs the forward and inverse Fourier transforms for real and complex inputs/outputs.

fft_type is one of the following:

• FFT : Forward complex-to-complex FFT.
• IFFT : Inverse complex-to-complex FFT.
• RFFT : Forward real-to-complex FFT.
• IRFFT : Inverse real-to-complex FFT (ie takes complex, returns real).

More formally, given the function fft which takes 1-dimensional tensors of complex types as input, produces 1-dimensional tensors of same types as output and computes the discrete Fourier transform:

For fft_type = FFT , result is defined as the final result of a series of L computations where L = size(fft_length) . For example, for L = 3 :

• result1[i0, ..., :] = fft(operand[i0, ..., :]) .
• result2[i0, ..., :, iR-1] = fft(result1[i0, ..., :, iR-1]) .
• result[i0, ..., :, iR-2, iR-1] = fft(result2[i0, ..., :, iR-2, iR-1]) .

Furthermore, given the function ifft which has the same type signature and computes the inverse of fft :

For fft_type = IFFT , result is defined as the inverse of the computations for fft_type = FFT . For example, for L = 3 :

• result1[i0, ..., :, iR-2, iR-1] = ifft(operand[i0, ..., :, iR-2, iR-1]) .
• result2[i0, ..., :, iR-1] = ifft(result1[i0, ..., :, iR-1]) .
• result[i0, ..., :] = ifft(result2[i0, ..., :]) .

Furthermore, given the function rfft which takes 1-dimensional tensors of floating-point types, produces 1-dimensional tensors of complex types of the same floating-point semantics and works as follows:

• rfft(real_operand) = truncated_result where
• complex_operand... = (real_operand..., 0.0) .
• complex_result = fft(complex_operand) .
• truncated_result = complex_result[:(rank(complex_result) / 2 + 1)] .

(When the discrete Fourier transform is computed for real operands, the first N/2 + 1 elements of the result unambiguously define the rest of the result, so the result of rfft is truncated to avoid computing redundant elements).

For fft_type = RFFT , result is defined as the final result of a series of L computations where L = size(fft_length) . For example, for L = 3 :

• result1[i0, ..., :] = rfft(operand[i0, ..., :]) .
• result2[i0, ..., :, iR-1] = fft(result1[i0, ..., :, iR-1]) .
• result[i0, ..., :, iR-2, iR-1] = fft(result2[i0, ..., :, iR-2, iR-1]) .

Finally, given the function irfft which has the same type signature and computes the inverse of rfft :

For fft_type = IRFFT , result is defined as the inverse of the computations for fft_type = RFFT . For example, for L = 3 :

• result1[i0, ..., :, iR-2, iR-1] = ifft(operand[i0, ..., :, iR-2, iR-1]) .
• result2[i0, ..., :, iR-1] = ifft(result1[i0, ..., :, iR-1]) .
• result[i0, ..., :] = irfft(result2[i0, ..., :]) .

Входы

Выходы

Ограничения

• (C1) size(fft_length) <= rank(operand) .
• (C2) The relationship between operand and result element types varies:
  • If fft_type = FFT , element_type(operand) and element_type(result) have the same complex type.
  • If fft_type = IFFT , element_type(operand) and element_type(result) have the same complex type.
  • If fft_type = RFFT , element_type(operand) is a floating-point type and element_type(result) is a complex type of the same floating-point semantics.
  • If fft_type = IRFFT , element_type(operand) is a complex type and element_type(result) is a floating-point type of the same floating-point semantics.
• (C3) 1 <= size(fft_length) <= 3 .
• (C4) If among operand and result , there is a tensor real of a floating-point type, then shape(real)[-size(fft_length):] = fft_length .
• (C5) shape(result) = shape(operand) except for:
  • If fft_type = RFFT , dim(result, -1) = dim(operand, -1) = 0 ? 0 : dim(operand, -1) / 2 + 1 .
  • If fft_type = IRFFT , dim(operand, -1) = dim(result, -1) = 0 ? 0 : dim(result, -1) / 2 + 1 .

Примеры

// %operand: [(1.0, 0.0), (0.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
%result = "stablehlo.fft"(%operand) {
  fft_type = #stablehlo<fft_type FFT>,
  fft_length = array<i64: 4>
} : (tensor<4xcomplex<f32>>) -> tensor<4xcomplex<f32>>
// %result: [(1.0, 0.0), (1.0, 0.0), (1.0, 0.0), (1.0, 0.0)]

пол

Семантика

Performs element-wise floor of operand tensor and produces a result tensor. Implements the roundToIntegralTowardNegative operation from the IEEE-754 specification. For quantized types, performs dequantize_op_quantize(floor, operand, type(result)) .

Входы

Выходы

Ограничения

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

Примеры

// %operand: [-0.8166, -0.2530, 0.2530, 0.8166, 2.0]
%result = "stablehlo.floor"(%operand) : (tensor<5xf32>) -> tensor<5xf32>
// %result: [-1.0, -1.0, 0.0, 0.0, 2.0]

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

собирать

Семантика

Gathers slices from operand tensor from offsets specified in start_indices and produces a result tensor.

The following diagram shows how elements in result map on elements in operand using a concrete example. The diagram picks a few example result indices and explains in detail which operand indices they correspond to.

More formally, result[result_index] = operand[operand_index] where:

• batch_dims = [d for d in axes(result) and d not in offset_dims] .
• batch_index = result_index[batch_dims...] .
• start_index is defined as:
  • start_indices[bi0, ..., :, ..., biN] where bi are individual elements in batch_index and : is inserted at the index_vector_dim index, if index_vector_dim < rank(start_indices) .
  • [start_indices[batch_index]] otherwise.
• For d_operand in axes(operand) ,
  • full_start_index[d_operand] = clamp(start_index[d_start], 0, dim(operand, d_operand) - slice_sizes[d_operand]) if d_operand = start_index_map[d_start] .
  • full_start_index[d_operand] = 0 otherwise.
• offset_index = result_index[offset_dims...] .
• full_offset_index = [oi0, ..., 0, ..., oiN] where oi are individual elements in offset_index , and 0 is inserted at indices from collapsed_slice_dims .
• operand_index = full_start_index + full_offset_index .

If indices_are_sorted is true then the implementation can assume that start_indices are sorted with respect to start_index_map , otherwise the behavior is undefined. More formally, for all i1 < i2 from indices(result) , full_start_index(i1) <= full_start_index(i2) .

Входы

Выходы

Ограничения

• (C1) rank(operand) = size(offset_dims) + size(collapsed_slice_dims) .
• (C2) 0 <= index_vector_dim <= rank(start_indices) .
• (C3) size(start_index_map) = index_vector_dim < rank(start_indices) ? dim(start_indices, index_vector_dim) : 1 .
• (C4) is_unique(offset_dims) and is_sorted(offset_dims) .
• (C5) 0 <= offset_dims < rank(result) .
• (C6) is_unique(collapsed_slice_dims) and is_sorted(collapsed_slice_dims) .
• (C7) 0 <= collapsed_slice_dims < rank(operand) .
• (C8) slice_sizes[collapsed_slice_dims...] <= 1 .
• (C9) is_unique(start_index_map) .
• (C10) 0 <= start_index_map < rank(operand) .
• (C11) size(slice_sizes) = rank(operand) .
• (C12) 0 <= slice_sizes <= shape(operand) .
• (C13) shape(result) = combine(batch_dim_sizes, offset_dim_sizes) where:
  • batch_dim_sizes = shape(start_indices) except that the dimension size of start_indices corresponding to index_vector_dim is not included.
  • offset_dim_sizes = shape(slice_sizes) except that the dimension sizes in slice_sizes corresponding to collapsed_slice_dims are not included.
  • combine puts batch_dim_sizes at axes corresponding to batch_dims and offset_dim_sizes at axes corresponding to offset_dims .
• (C14) element_type(operand) = element_type(result) .

Примеры

// %operand: [
//            [[1, 2], [3, 4], [5, 6], [7, 8]],
//            [[9, 10],[11, 12], [13, 14], [15, 16]],
//            [[17, 18], [19, 20], [21, 22], [23, 24]]
//           ]
// %start_indices: [
//                  [[0, 0], [1, 0], [2, 1]],
//                  [[0, 1], [1, 1], [0, 2]]
//                 ]
%result = "stablehlo.gather"(%operand, %start_indices) {
  dimension_numbers = #stablehlo.gather<
    offset_dims = [2, 3],
    collapsed_slice_dims = [0],
    start_index_map = [1, 0],
    index_vector_dim = 2>,
  slice_sizes = array<i64: 1, 2, 2>,
  indices_are_sorted = false
} : (tensor<3x4x2xi32>, tensor<2x3x2xi64>) -> tensor<2x3x2x2xi32>
// %result: [
//            [
//              [[1, 2], [3, 4]],
//              [[3, 4], [5, 6]],
//              [[13, 14], [15, 16]]
//            ],
//            [
//              [[9, 10], [11, 12]],
//              [[11, 12], [13, 14]],
//              [[17, 18], [19, 20]]
//            ]
//          ]

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

get_dimension_size

Семантика

Produces the size of the given dimension of the operand . More formally, result = dim(operand, dimension) . The Semantics concerns only with the shape component of the type. The element-type could be anything.

Входы

Выходы

Ограничения

• (C1) 0 <= dimension < rank(operand) .

Примеры

// %operand: [[1, 2, 3], [4, 5, 6]]
%result = "stablehlo.get_dimension_size"(%operand) {
  dimension = 1 : i64
} : (tensor<2x3xi64>) -> tensor<i32>
// %result: 3

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

get_tuple_element

Семантика

Extracts element at index position of the operand tuple and produces a result . More formally, result = operand[index] .

Входы

Выходы

Ограничения

• (C1) 0 <= index < size(operand) .
• (C2) type(result) = tuple_element_types(operand)[index] .

Примеры

// %operand: ([1.0, 2.0], (3))
%result = "stablehlo.get_tuple_element"(%operand) {
  index = 0 : i32
} : (tuple<tensor<2xf32>, tuple<tensor<i32>>>) -> tensor<2xf32>
// %result: [1.0, 2.0]

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

если

Семантика

Produces the output from executing exactly one function from true_branch or false_branch depending on the value of pred . More formally, result = pred ? true_branch() : false_branch() .

Входы

Выходы

Constraints

• (C1) input_types(true_branch) = input_types(false_branch) = [] .
• (C2) output_types(true_branch) = output_types(false_branch) .
• (C3) type(results...) = output_types(true_branch) .

Примеры

// %result_true_branch: 10
// %result_false_branch: 11
// %pred: true
%result = "stablehlo.if"(%pred) ({
  "stablehlo.return"(%result_true_branch) : (tensor<i32>) -> ()
}, {
  "stablehlo.return"(%result_false_branch) : (tensor<i32>) -> ()
}) : (tensor<i1>) -> tensor<i32>
// %result: 10

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

imag

Семантика

Extracts the imaginary part, element-wise, from the operand and produces a result tensor. More formally, for each element x : imag(x) = is_complex(x) ? imaginary_part(x) : constant(0, element_type(result)) .

Входы

Выходы

Ограничения

• (C1) shape(result) = shape(operand) .
• (C2) element_type(result) is defined as:
  • complex_element_type(element_type(operand)) if is_complex(operand) .
  • element_type(operand) otherwise.

Примеры

// %operand: [(1.0, 2.0), (3.0, 4.0)]
%result = "stablehlo.imag"(%operand) : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
// %result: [2.0, 4.0]

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

infeed

Семантика

Reads data from the infeed and produces results .

Semantics of infeed_config is implementation-defined.

results consist of payload values which come first and a token which comes last. In the future, we are planning to split the payload and the token into two separate outputs to improve clarity ( #670 ).

Входы

Выходы

Constraints

• (C1) 0 < size(results) .
• (C2) is_empty(result[:-1]) or is_tensor(type(results[:-1])) .
• (C3) is_token(type(results[-1])) .

Примеры

// %token: !stablehlo.token
// infeed_queue[0]: [[1, 2], [3, 4]]
// infeed_queue[1]: [[5, 6], [7, 8]]
%results0:2 = "stablehlo.infeed"(%token) {
  infeed_config = ""
} : (!stablehlo.token) -> (tensor<2x2xi64>, !stablehlo.token)
// results0#0: [[1, 2], [3, 4]]
%results1:2 = "stablehlo.infeed"(%token) {
  infeed_config = ""
} : (!stablehlo.token) -> (tensor<2x2xi64>, !stablehlo.token)
// results1#0: [[5, 6], [7, 8]]

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

йота

Семантика

Fills an output tensor with values in increasing order starting from zero along the iota_dimension dimension. More formally,

output[output_index] = constant(is_quantized(output) ? quantize(output_index[iota_dimension], element_type(output)) : output_index[iota_dimension], element_type(output)) .

Входы

Выходы

Constraints

• (C1) 0 <= iota_dimension < rank(output) .

Примеры

%output = "stablehlo.iota"() {
  iota_dimension = 0 : i64
} : () -> tensor<4x5xi32>
// %output: [
//           [0, 0, 0, 0, 0],
//           [1, 1, 1, 1, 1],
//           [2, 2, 2, 2, 2],
//           [3, 3, 3, 3, 3]
//          ]

%output = "stablehlo.iota"() {
  iota_dimension = 1 : i64
} : () -> tensor<4x5xi32>
// %output: [
//           [0, 1, 2, 3, 4],
//           [0, 1, 2, 3, 4],
//           [0, 1, 2, 3, 4],
//           [0, 1, 2, 3, 4]
//          ]

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

is_finite

Семантика

Performs element-wise check whether the value in x is finite (ie is neither +Inf, -Inf, nor NaN) and produces a y tensor. Implements the isFinite operation from the IEEE-754 specification. For quantized types, the result is always true .

Входы

Выходы

Constraints

• (C1) shape(x) = shape(y) .

Примеры

// Logical values: -Inf, +Inf, NaN, ...
// %x: [0xFFF0000000000000, 0x7FF0000000000000, 0x7FF8000000000000, -10.0, -0.0, 0.0, 10.0]
%y = "stablehlo.is_finite"(%x) : (tensor<7xf64) -> tensor<7xi1>
// %y: [false, false, false, true, true, true, true]

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

бревно

Семантика

Performs element-wise logarithm operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: log from IEEE-754.
• For complex numbers: complex logarithm.
• For quantized types: dequantize_op_quantize(log, operand, type(result)) .

Входы

Выходы

Constraints

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

Примеры

// %operand: [[1.0, 2.0], [3.0, 4.0]]
%result = "stablehlo.log"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
// %result: [[0.0, 0.69314718055994529], [1.0986122886681098, 1.3862943611198906]]

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

log_plus_one

Семантика

Performs element-wise logarithm plus one operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: logp1 from IEEE-754.
• For complex numbers: complex logarithm plus one.
• For quantized types: dequantize_op_quantize(log_plus_one, operand, type(result)) .

Входы

Выходы

Constraints

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

Примеры

// %operand: [0.0, -0.999, 7.0, 6.38905621, 15.0]
%result = "stablehlo.log_plus_one"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// %result: [0.0, -6.90776825, 2.07944155, 2.0, 2.77258873]

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

логистический

Семантика

Performs element-wise logistic operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: division(1, addition(1, exp(-x))) from IEEE-754.
• For complex numbers: complex logistic.
• For quantized types: dequantize_op_quantize(logistic, operand, type(result)) .

Входы

Выходы

Constraints

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

Примеры

// %operand: [[0.0, 1.0], [2.0, 3.0]]
%result = "stablehlo.logistic"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
// %result: [[0.5, 0.73105858], [0.88079708, 0.95257413]]

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

карта

Семантика

Applies a map function computation to inputs along the dimensions and produces a result tensor.

More formally, result[result_index] = computation(inputs...[result_index]) . Note that dimensions are currently unused and will likely be removed in the future ( #487 ).

Входы

Выходы

Constraints

• (C1) shape(inputs...) = shape(result) .
• (C2) 0 < size(inputs) = N .
• (C3) dimensions = range(rank(inputs[0])) .
• (C4) computation has type (tensor<E0>, ..., tensor<EN-1>) -> tensor<E'> where Ei = element_type(inputs[i]) and E' = element_type(result) .

Примеры

// %input0: [[0, 1], [2, 3]]
// %input1: [[4, 5], [6, 7]]
%result = "stablehlo.map"(%input0, %input1) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = stablehlo.multiply %arg0, %arg1 : tensor<i64>
    stablehlo.return %0 : tensor<i64>
}) {
  dimensions = array<i64: 0, 1>
} : (tensor<2x2xi64>, tensor<2x2xi64>) -> tensor<2x2xi64>
// %result: [[0, 5], [12, 21]]

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

максимум

Семантика

Performs element-wise max operation on tensors lhs and rhs and produces a result tensor. Depending on the element type, does the following:

• For booleans: logical OR.
• For integers: integer maximum.
• For floats: maximum from IEEE-754.
• For complex numbers: lexicographic maximum for the (real, imaginary) pair. Imposing an ordering on complex numbers involves surprising semantics, so in the future we are planning to remove support for complex numbers for this operation ( #560 ).
• For quantized types:
  • dequantize_op_quantize(maximum, lhs, rhs, type(result)) .

Входы

Выходы

Constraints

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

Примеры

// %lhs: [[1, 2], [7, 8]]
// %rhs: [[5, 6], [3, 4]]
%result = "stablehlo.maximum"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[5, 6], [7, 8]]

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

минимум

Семантика

Performs element-wise min operation on tensors lhs and rhs and produces a result tensor. Depending on the element type, does the following:

• For booleans: logical AND.
• For integers: integer minimum.
• For floats: minimum from IEEE-754.
• For complex numbers: lexicographic minimum for the (real, imaginary) pair. Imposing an ordering on complex numbers involves surprising semantics, so in the future we are planning to remove support for complex numbers for this operation ( #560 ).
• For quantized types:
  • dequantize_op_quantize(minimum, lhs, rhs, type(result)) .

Входы

Выходы

Constraints

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

Примеры

// %lhs: [[1, 2], [7, 8]]
// %rhs: [[5, 6], [3, 4]]
%result = "stablehlo.minimum"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[1, 2], [3, 4]]

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

умножать

Семантика

Performs element-wise product of two tensors lhs and rhs and produces a result tensor. Depending on the element type, does the following:

• For booleans: logical AND.
• For integers: integer multiplication.
• For floats: multiplication from IEEE-754.
• For complex numbers: complex multiplication.
• For quantized types:
  • dequantize_op_quantize(multiply, lhs, rhs, type(result)) .

Входы

Выходы

Constraints

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

Примеры

// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.multiply"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[5, 12], [21, 32]]

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

отрицать

Семантика

Performs element-wise negation of operand tensor and produces a result tensor. Depending on the element type, does the following:

• For signed integers: integer negation.
• For unsigned integers: bitcast to signed integer, integer negation, bitcast back to unsigned integer.
• For floats: negate from IEEE-754.
• For complex numbers: complex negation.
• For quantized types: dequantize_op_quantize(negate, operand, type(result)) .

Входы

Выходы

Constraints

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

Примеры

// Negation operation with integer Tensors
// %operand: [0, -2]
%result = "stablehlo.negate"(%operand) : (tensor<2xi32>) -> tensor<2xi32>
// %result: [0, 2]

// Negation operation with with complex tensors
// %operand: (2.5, 0.0)
%result = "stablehlo.negate"(%operand) : (tensor<1xcomplex<f32>>) -> tensor<1xcomplex<f32>>
// %result: [-2.5, -0.0]

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

нет

Семантика

Performs element-wise NOT of tensor operand and produces a result tensor. Depending on the element type, does the following:

• For booleans: logical NOT.
• For integers: bitwise NOT.

Аргументы

Выходы

Constraints

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

Примеры

// Bitwise operation with with integer tensors
// %operand: [[1, 2], [3, 4]]
%result = "stablehlo.not"(%operand) : (tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[-2, -3], [-4, -5]]

// Bitwise operation with with boolean tensors
// %operand: [true, false]
%result = "stablehlo.not"(%operand) : (tensor<2xi1>) -> tensor<2xi1>
// %result: [false, true]

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

optimization_barrier

Семантика

Ensures that the operations that produce the operand are executed before any operations that depend on the result and prevents compiler transformations from moving operations across the barrier. Other than that, the operation is an identity, ie result = operand .

Аргументы

Выходы

Constraints

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

Примеры

// %operand0: 0.0
// %operand1: 1.0
%result0, %result1 = "stablehlo.optimization_barrier"(%operand0, %operand1) : (tensor<f32>, tensor<f32>) -> (tensor<f32>, tensor<f32>)
// %result0: 0.0
// %result1: 1.0

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

или

Семантика

Performs element-wise OR of two tensors lhs and rhs and produces a result tensor. Depending on the element type, does the following:

• For booleans: logical OR.
• For integers: bitwise OR.

Входы

Выходы

Constraints

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

Примеры

// Bitwise operation with with integer tensors
// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.or"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[5, 6], [7, 12]]

// Logical operation with with boolean tensors
// %lhs: [[false, false], [true, true]]
// %rhs: [[false, true], [false, true]]
%result = "stablehlo.or"(%lhs, %rhs) : (tensor<2x2xi1>, tensor<2x2xi1>) -> tensor<2x2xi1>
// %result: [[false, true], [true, true]]

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

outfeed

Семантика

Writes inputs to the outfeed and produces a result token.

Semantics of outfeed_config is implementation-defined.

Входы

Выходы

Примеры

%result = "stablehlo.outfeed"(%input0, %token) {
  outfeed_config = ""
} : (tensor<2x2x2xi64>, !stablehlo.token) -> !stablehlo.token

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

подушечка

Семантика

Expands operand by padding around the tensor as well as between the elements of the tensor with the given padding_value .

edge_padding_low and edge_padding_high specify the amount of padding added at the low-end (next to index 0) and the high-end (next to the highest index) of each dimension respectively. The amount of padding can be negative, where the absolute value of negative padding indicates the number of elements to remove from the specified dimension.

interior_padding specifies the amount of padding added between any two elements in each dimension which may not be negative. Interior padding occurs before edge padding such that negative edge padding will remove elements from the interior-padded operand.

More formally, result[result_index] is defined as:

• operand[operand_index] if result_index = edge_padding_low + operand_index * (interior_padding + 1) .
• padding_value otherwise.

Входы

Выходы

Constraints

• (C1) element_type(operand) = element_type(padding_value) = element_type(result) .
• (C2) size(edge_padding_low) = size(edge_padding_high) = size(interior_padding) = rank(operand) .
• (C3) 0 <= interior_padding .
• (C4) shape(result) = shape(operand) + edge_padding_low + max(shape(operand) - 1, 0) * interior_padding + edge_padding_high .

Примеры

// %operand: [
//            [1, 2, 3],
//            [4, 5, 6]
//           ]
// %padding_value: 0
%result = "stablehlo.pad"(%operand, %padding_value) {
  edge_padding_low = array<i64: 0, 1>,
  edge_padding_high = array<i64: 2, 1>,
  interior_padding = array<i64: 1, 2>
} : (tensor<2x3xi32>, tensor<i32>) -> tensor<5x9xi32>
// %result: [
//           [0, 1, 0, 0, 2, 0, 0, 3, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0],
//           [0, 4, 0, 0, 5, 0, 0, 6, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0]
//          ]

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

partition_id

Семантика

Produces partition_id of the current process.

Выходы

Примеры

%result = "stablehlo.partition_id"() : () -> tensor<ui32>

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

popcnt

Семантика

Performs element-wise count of the number of bits set in the operand tensor and produces a result tensor.

Входы

Выходы

Constraints

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

Примеры

// %operand: [0, 1, 2, 127]
%result = "stablehlo.popcnt"(%operand) : (tensor<4xi64>) -> tensor<4xi64>
// %result: [0, 1, 1, 7]

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

власть

Семантика

Performs element-wise exponentiation of lhs tensor by rhs tensor and produces a result tensor. Depending on the element type, does the following:

• For integers: integer exponentiation.
• For floats: pow from IEEE-754.
• For complex numbers: complex exponentiation.
• For quantized types: dequantize_op_quantize(power, lhs, rhs, type(result)) .

Входы

Выходы

Constraints

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

Примеры

// %lhs: [-2.0, -0.0, -36.0, 5.0, 3.0, 10000.0]
// %rhs: [2.0, 2.0, 1.1, 2.0, -1.0, 10.0]
%result = "stablehlo.power"(%lhs, %rhs) : (tensor<6xf64>, tensor<6xf64>) -> tensor<6xf64>
// %result: [4.0, 0.0, -nan, 25.0, 0.333333343, inf]

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

настоящий

Семантика

Extracts the real part, element-wise, from the operand and produces a result tensor. More formally, for each element x : real(x) = is_complex(x) ? real_part(x) : x .

Входы

Выходы

Constraints

• (C1) shape(result) = shape(operand) .
• (C2) element_type(result) is defined as:
  • complex_element_type(element_type(operand)) if is_complex(operand) .
  • element_type(operand) otherwise.

Примеры

// %operand: [(1.0, 2.0), (3.0, 4.0)]
%result = "stablehlo.real"(%operand) : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
// %result: [1.0, 3.0]

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

получение

Семантика

Receives data from a channel with channel_id and produces results .

If is_host_transfer is true , then the operation transfers data from the host. Otherwise, it transfers data from another device. What this means is implementation-defined. This flag duplicates the information provided in channel_type , so in the future we are planning to only keep one of them ( #666 ).

results consist of payload values which come first and a token which comes last. In the future, we are planning to split the payload and the token into two separate outputs to improve clarity ( #670 ).

Входы

Выходы

Constraints

• (C1) channel_type is defined as:
  • HOST_TO_DEVICE if is_host_transfer = true ,
  • DEVICE_TO_DEVICE otherwise.
• (C2) 0 < size(results) .
• (C3) is_empty(result[:-1]) or is_tensor(type(results[:-1])) .
• (C4) is_token(type(results[-1])) .

Примеры

%results0, %results1 = "stablehlo.recv"(%token) {
  channel_handle = #stablehlo.channel_handle<handle = 1, type = 3>,
  is_host_transfer = true
} : (!stablehlo.token) -> (tensor<2x2xi64>, !stablehlo.token)

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

уменьшать

Семантика

Applies a reduction function body to inputs and init_values along the dimensions and produces results tensors.

The order of reductions is implementation-defined, which means that body and init_values must form a monoid to guarantee that the operation produces the same results for all inputs on all implementations. However, this condition doesn't hold for many popular reductions. Eg floating-point addition for body and zero for init_values don't actually form a monoid because floating-point addition is not associative.

More formally, results...[j0, ..., jR-1] = reduce(input_slices_converted) where:

• input_slices = inputs...[j0, ..., :, ..., jR-1] , where : are inserted at dimensions .
• input_slices_converted = to_destination_type(input_slices..., type(func_inputs(body)[:len(func_inputs(body))//2])...) .
• init_values_converted = to_destination_type(init_values..., type(func_inputs(body)[len(func_inputs(body))//2:])...) .
• reduce(input_slices_converted) = exec(schedule) for some binary tree schedule where:
  • exec(node) = body(exec(node.left), exec(node.right)) .
  • exec(leaf) = leaf.value .
• schedule is an implementation-defined full binary tree whose in-order traversal consists of:
  • input_slices_converted...[index] values, for all index in index_space(input_slices_converted) in the ascending lexicographic order of index .
  • Interspersed with an implementation-defined amount of init_values_converted at implementation-defined positions.

Входы

Выходы

Constraints

• (C1) same(shape(inputs...)) .
• (C2) element_type(inputs...) = element_type(init_values...) .
• (C3) 0 < size(inputs) = size(init_values) = size(results) = N .
• (C4) 0 <= dimensions < rank(inputs[0]) .
• (C5) is_unique(dimensions) .
• (C6) body has type (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>) where is_promotable(element_type(inputs[i]), Ei) .
• (C7) shape(results...) = shape(inputs...) except that the dimension sizes of inputs... corresponding to dimensions are not included.
• (C8) element_type(results[i]) = Ei for all i in [0,N) .

Примеры

// %input = [[0, 1, 2, 3, 4, 5]]
// %init_value = 0
%result = "stablehlo.reduce"(%input, %init_value) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  dimensions = array<i64: 1>
} : (tensor<1x6xi64>, tensor<i64>) -> tensor<1xi64>
// %result = [15]

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

reduce_precision

Семантика

Performs element-wise conversion of operand to another floating-point type that uses exponent_bits and mantissa_bits and back to the original floating-point type and produces an output tensor.

More formally:

• The mantissa bits of the original value are updated to round the original value to the nearest value representable with mantissa_bits using roundToIntegralTiesToEven semantics.
• Then, if mantissa_bits are smaller than the number of mantissa bits of the original value, the mantissa bits are truncated to mantissa_bits .
• Then, if the exponent bits of the intermediate result don't fit into the range provided by exponent_bits , the intermediate result overflows to infinity using the original sign or underflows to zero using the original sign.
• For quantized types, performs dequantize_op_quantize( lambda operand: reduce_precision(operand, exponent_bits, mantissa_bits), operand, type(result)) .

Входы

Выходы

Constraints

• (C1) baseline_type(operand) = baseline_type(output) .
• (C2) 1 <= exponent_bits .
• (C3) 0 <= mantissa_bits .

Примеры

// Logical values: +Inf, NaN, +Denormal, 0.0, 65519.0, 65520.0
// %operand: [0x7FF0000000000000, 0x7FFFFFFFFFFFFFFF, 0x0000000000000001, 0.0, 65519.0, 65520.0]
%output = "stablehlo.reduce_precision"(%operand) {
  exponent_bits = 5 : i32,
  mantissa_bits = 10 : i32
} : (tensor<6xf64>) -> tensor<6xf64>
// Logical values: +Inf, NaN, 0.0, 0.0, 65504.0, +Inf
// %output: [0x7FF0000000000000, 0x7FFFFFFFFFFFFFFF, 0.0, 0.0, 65504.0, 0x7FF0000000000000]

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

reduce_scatter

Семантика

Within each process group in the StableHLO process grid, performs reduction, using computations , over the values of the operand tensor from each process, splits the reduction result along scatter_dimension into parts, and scatters the split parts between the processes to produce the result .

The operation splits the StableHLO process grid into process_groups which is defined as follows:

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

Afterwards, within each process_group :

• reduced_value = all_reduce(operand, replica_groups, channel_id, use_global_device_ids, computation) .
• parts@sender = split(reduced_value@sender, dim(process_groups, 1), scatter_dimension) .
• result@receiver = parts@sender[receiver_index] for all sender in process_group , where receiver_index = process_group.index(receiver) .

Входы

Выходы

Constraints

• (C1) dim(operand, scatter_dimension) % dim(process_groups, 1) = 0 .
• (C2) 0 <= scatter_dimension < rank(operand) .
• (C3) is_unique(replica_groups) .
• (C4) size(replica_groups) is defined as:
  • num_replicas if cross_replica is used.
  • num_replicas if cross_replica_and_partition is used.
  • num_processes if flattened_ids is used.
• (C5) 0 <= replica_groups < size(replica_groups) .
• (C6) If use_global_device_ids = true , then channel_id > 0 .
• (C7) computation has type (tensor<E>, tensor<E>) -> (tensor<E>) where is_promotable(element_type(operand), E) .
• (C8) shape(result) = shape(operand) except:
  • dim(result, scatter_dimension) = dim(operand, scatter_dimension) / dim(process_groups, 1) .
• (C9) element_type(result) = E .

Примеры

// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [[1, 2, 3, 4],
//                   [5, 6, 7, 8]]
// %operand@(1, 0): [[9, 10, 11, 12],
//                   [13, 14, 15, 16]]
%result = "stablehlo.reduce_scatter"(%operand) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
  "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  scatter_dimension = 1 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor<2x4xi64>) -> tensor<2x2xi64>
//
// %result@(0, 0): [[10, 12],
//                  [18, 20]]
// %result@(1, 0): [[14, 16],
//                  [22, 24]]

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

reduce_window

Семантика

Applies a reduction function body to windows of inputs and init_values and produces results .

The following diagram shows how elements in results... are computed from inputs... using a concrete example.

More formally, results...[result_index] = reduce(windows, init_values, axes(inputs...), body) (see reduce ) where:

• padded_inputs = pad(inputs..., init_values..., padding[:, 0], padding[:, 1], base_dilations - 1) .
• window_start = result_index * window_strides .
• window_end = window_start + (window_dimensions - 1) * window_dilations + 1 .
• windows = slice(padded_inputs..., window_start, window_end, window_dilations) .

Входы

Выходы

Constraints

• (C1) 0 < size(inputs) = size(init_values) = size(results) = N .
• (C2) same(shape(inputs...)) .
• (C3) element_type(inputs...) = element_type(init_values...) .
• (C4) size(window_dimensions) = rank(inputs[0]) .
• (C5) 0 < window_dimensions .
• (C6) size(window_strides) = rank(inputs[0]) .
• (C7) 0 < window_strides .
• (C8) size(base_dilations) = rank(inputs[0]) .
• (C9) 0 < base_dilations .
• (C10) size(window_dilations) = rank(inputs[0]) .
• (C11) 0 < window_dilations .
• (C12) shape(padding) = [rank(inputs[0]), 2] .
• (C13) body has type (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>) where is_promotable(element_type(inputs[i]), Ei) .
• (C14) same(shape(results...)) .
• (C15) shape(results[0]) = num_windows where:
  • dilated_input_shape = shape(inputs[0]) = 0 ? 0 : (shape(inputs[0]) - 1) * base_dilations + 1 .
  • padded_input_shape = padding[:, 0] + dilated_input_shape + padding[:, 1] .
  • dilated_window_shape = (window_dimensions - 1) * window_dilations + 1 .
  • is_empty_window = padded_input_shape = 0 || dilated_window_shape > padded_input_shape .
  • num_windows = is_empty_window ? 0 : floor((padded_input_shape - dilated_window_shape) / window_strides) + 1 .
• (C16) element_type(results[i]) = Ei for all i in [0,N) .

Примеры

// %input = [[1, 2], [3, 4], [5, 6]]
// %init_value = 0
%result = "stablehlo.reduce_window"(%input, %init_value) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  window_dimensions = array<i64: 2, 1>,
  window_strides = array<i64: 4, 1>,
  base_dilations = array<i64: 2, 1>,
  window_dilations = array<i64: 3, 1>,
  padding = dense<[[2, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<3x2xi64>, tensor<i64>) -> tensor<2x2xi64>
// %result = [[0, 0], [3, 4]]

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

остаток

Семантика

Performs element-wise remainder of dividend lhs and divisor rhs tensors and produces a result tensor.

More formally, the sign of the result is taken from the dividend, and the absolute value of the result is always less than the divisor's absolute value. The remainder is calculated as lhs - d * rhs , where d is given by:

• For integers: stablehlo.divide(lhs, rhs) .
• For floats: division(lhs, rhs) from IEEE-754 with rounding attribute roundTowardZero .
• For complex numbers: TBD ( #997 ).
• For quantized types:
  • dequantize_op_quantize(remainder, lhs, rhs, type(result)) .

For floating-point element types, this operation is in contrast with the remainder operation from IEEE-754 specification where d is an integral value nearest to the exact value of lhs/rhs with ties to even.

Входы

Выходы

Constraints

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

Примеры

// %lhs: [17, -17, 17, -17]
// %rhs: [3, 3, -3, -3]
%result = "stablehlo.remainder"(%lhs, %rhs) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi64>
// %result: [2, -2, 2, -2]

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

replica_id

Семантика

Produces replica_id of the current process.

Выходы

Примеры

%result = "stablehlo.replica_id"() : () -> tensor<ui32>

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

изменить форму

Семантика

Performs reshape of operand tensor to a result tensor. Conceptually, it amounts to keeping the same canonical representation but potentially changing the shape, eg from tensor<2x3xf32> to tensor<3x2xf32> or tensor<6xf32> .

More formally, result[result_index] = operand[operand_index] where result_index and operand_index have the same position in the lexicographic ordering of index_space(result) and index_space(operand) .

Входы

Выходы

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) and quantization_dimension(result) may differ, otherwise.
• (C2) size(operand) = size(result) .
• (C3) If is_per_axis_quantized(operand) :
  • reduce(dims(operand, [0, 1, ..., quantization_dimension(operand) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) = reduce(dims(result, [0, 1, ..., quantization_dimension(result) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) .
  • dim(operand, quantization_dimension(operand)) = dim(result, quantization_dimension(result)) .
  • reduce(dims(operand, [quantization_dimension(operand) + 1, ..., rank(operand) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) = reduce(dims(result, [quantization_dimension(result) + 1, ..., rank(result) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) .

Примеры

// %operand: [[1, 2, 3], [4, 5, 6]]
%result = "stablehlo.reshape"(%operand) : (tensor<2x3xi32>) -> tensor<3x2xi32>
// %result: [[1, 2], [3, 4], [5, 6]]

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

обеспечить регресс

Семантика

Reverses the order of elements in the operand along the specified dimensions and produces a result tensor. More formally, result[result_index] = operand[operand_index] where:

• operand_index[d] = dim(result, d) - result_index[d] - 1 if d in dimensions .
• operand_index[d] = result_index[d] otherwise.

Входы

Выходы

Constraints

• (C1) type(operand) = type(result) .
• (C2) is_unique(dimensions) .
• (C3) 0 <= dimensions < rank(result) .

Примеры

// %operand = [[1, 2], [3, 4], [5, 6]]
%result = "stablehlo.reverse"(%operand) {
  dimensions = array<i64: 1>
} : (tensor<3x2xi32>) -> tensor<3x2xi32>
// %result: [[2, 1], [4, 3], [6, 5]]

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

звонок

Семантика

Generates random numbers using the rng_distribution algorithm and produces a result tensor of a given shape shape .

If rng_distribution = UNIFORM , then the random numbers are generated following the uniform distribution over the interval [a, b) . If a >= b , the behavior is undefined.

If rng_distribution = NORMAL , then the random numbers are generated following the normal distribution with mean = a and standard deviation = b . If b < 0 , the behavior is undefined.

The exact way how random numbers are generated is implementation-defined. For example, they may or may not be deterministic, and they may or may not use hidden state.

In conversations with many stakeholders, this op has come up as effectively deprecated, so in the future we are planning to explore removing it ( #597 ).

Входы

Выходы

Constraints

• (C1) element_type(a) = element_type(b) = element_type(result) .
• (C2) If rng_distribution = NORMAL , then is_float(a) .
• (C3) shape(result) = shape .

Примеры

// %a = 0
// %b = 2
// %shape = [3, 3]
%result = "stablehlo.rng"(%a, %b, %shape) {
  rng_distribution = #stablehlo<rng_distribution UNIFORM>
} : (tensor<i32>, tensor<i32>, tensor<2xi64>) -> tensor<3x3xi32>
// %result: [
//           [1, 0, 1],
//           [1, 1, 1],
//           [0, 0, 0]
//          ]

rng_bit_generator

Семантика

Returns an output filled with uniform random bits and an updated output state output_state using the pseudorandom number generator algorithm rng_algorithm given an initial state initial_state . The output is guaranteed to be deterministic function of initial_state , but it is not guaranteed to be deterministic between implementations.

rng_algorithm is one of the following:

• DEFAULT : Implementation-defined algorithm.
• THREE_FRY : Implementation-defined variant of the Threefry algorithm.*
• PHILOX : Implementation-defined variant of the Philox algorithm.*

* See: Salmon et al. SC 2011. Parallel random numbers: as easy as 1, 2, 3.

Входы

Выходы

Constraints

• (C1) type(initial_state) = type(output_state) .
• (C2) size(initial_state) is defined as:
  • implementation-defined if rng_algorithm = DEFAULT .
  • 2 if rng_algorithm = THREE_FRY .
  • 2 or 3 if rng_algorithm = PHILOX .

Примеры

// %initial_state: [1, 2]
%output_state, %output = "stablehlo.rng_bit_generator"(%initial_state) {
  rng_algorithm = #stablehlo<rng_algorithm THREE_FRY>
} : (tensor<2xui64>) -> (tensor<2xui64>, tensor<2x2xui64>)
// %output_state: [1, 6]
// %output: [
//           [9236835810183407956, 16087790271692313299],
//           [18212823393184779219, 2658481902456610144]
//          ]

round_nearest_afz

Семантика

Performs element-wise rounding towards the nearest integer, breaking ties away from zero, on the operand tensor and produces a result tensor. Implements the roundToIntegralTiesToAway operation from the IEEE-754 specification. For quantized types, performs dequantize_op_quantize(round_nearest_afz, operand, type(result)) .

Входы

Выходы

Constraints

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

Примеры

// %operand = [-2.5, 0.4, 0.5, 0.6, 2.5]
%result = "stablehlo.round_nearest_afz"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// %result: [-3.0, 0.0, 1.0, 1.0, 3.0]

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

round_nearest_even

Семантика

Performs element-wise rounding towards the nearest integer, breaking ties towards the even integer, on the operand tensor and produces a result tensor. Implements the roundToIntegralTiesToEven operation from the IEEE-754 specification. For quantized types, performs dequantize_op_quantize(round_nearest_even, operand, type(result)) .

Входы

Выходы

Constraints

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

Примеры

// %operand = [-2.5, 0.4, 0.5, 0.6, 2.5]
%result = "stablehlo.round_nearest_even"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// %result: [-2.0, 0.0, 0.0, 1.0, 2.0]

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

rsqrt

Семантика

Performs element-wise reciprocal square root operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: rSqrt from IEEE-754.
• For complex numbers: complex reciprocal square root.
• For quantized types: dequantize_op_quantize(rsqrt, operand, type(result)) .

Входы

Выходы

Constraints

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

Примеры

// %operand: [[1.0, 4.0], [9.0, 25.0]]
%result = "stablehlo.rsqrt"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[1.0, 0.5], [0.33333343, 0.2]]

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

разбрасывать

Семантика

Produces results tensors which are equal to inputs tensors except that several slices specified by scatter_indices are updated with the values updates using update_computation .

The following diagram shows how elements in updates... map on elements in results... using a concrete example. The diagram picks a few example updates... indices and explains in detail which results... indices they correspond to.

More formally, for all update_index in index_space(updates[0]) :

• update_scatter_dims = [d for d in axes(updates[0]) and d not in update_window_dims] .
• update_scatter_index = update_index[update_scatter_dims...] .
• start_index is defined as:
  • scatter_indices[si0, ..., :, ..., siN] where si are individual elements in update_scatter_index and : is inserted at the index_vector_dim index, if index_vector_dim < rank(scatter_indices) .
  • [scatter_indices[update_scatter_index]] otherwise.
• For d_input in axes(inputs[0]) ,
  • full_start_index[d_input] = start_index[d_start] if d_input = scatter_dims_to_operand_dims[d_start] .
  • full_start_index[d_input] = 0 otherwise.
• update_window_index = update_index[update_window_dims...] .
• full_window_index = [wi0, ..., 0, ..., wiN] where wi are individual elements in update_window_index , and 0 is inserted at indices from inserted_window_dims .
• result_index = full_start_index + full_window_index .

Given that, results = exec(schedule, inputs) , where:

• schedule is an implementation-defined permutation of index_space(updates[0]) .
• exec([update_index, ...], results) = exec([...], updated_results) where:
  • If result_index is in bounds for shape(results...)
  • updates_converted = to_destination_type( updates...[update_index], type(func_inputs(update_computation) [len(func_inputs(update_computation))//2:])... )
  • updated_values = update_computation(results...[result_index], updates_converted)
  • updated_results is a copy of results with results...[result_index] set to updated_values... .
  • В противном случае
  • updated_results = results .
• exec([], results) = results .

If indices_are_sorted is true then the implementation can assume that scatter_indices are sorted with respect to scatter_dims_to_operand_dims , otherwise the behavior is undefined. More formally, for all i1 < i2 from indices(result) , full_start_index(i1) <= full_start_index(i2) .

If unique_indices is true then the implementation can assume that all result_index indices being scattered to are unique. If unique_indices is true but the indices being scattered to are not unique then the behavior is undefined.

Входы

Выходы

Constraints

• (C1) same(shape(inputs...)) .
• (C2) rank(inputs[0]) = size(update_window_dims) + size(inserted_window_dims) .
• (C3) same(shape(updates...)) .
• (C4) shape(updates[0]) = combine(update_scatter_dim_sizes, update_window_dim_sizes) where:
  • update_scatter_dim_sizes = shape(scatter_indices) except that the dimension size of scatter_indices corresponding to index_vector_dim is not included.
  • update_window_dim_sizes <= shape(inputs[0]) except that the dimension sizes in inputs[0] corresponding to inserted_window_dims are not included.
  • combine puts update_scatter_dim_sizes at axes corresponding to update_scatter_dims and update_window_dim_sizes at axes corresponding to update_window_dims .
• (C5) 0 < size(inputs) = size(updates) = N .
• (C6) element_type(updates...) = element_type(inputs...) .
• (C7) is_unique(update_window_dims) and is_sorted(update_window_dims) .
• (C8) 0 <= update_window_dims < rank(updates[0]) .
• (C9) is_unique(inserted_window_dims) and is_sorted(update_window_dims) .
• (C10) 0 <= inserted_window_dims < rank(inputs[0]) .
• (C11) size(scatter_dims_to_operand_dims) = index_vector_dim < rank(scatter_indices) ? dim(scatter_indices, index_vector_dim) : 1 .
• (C12) is_unique(scatter_dims_to_operand_dims) .
• (C13) 0 <= scatter_dims_to_operand_dims < rank(inputs[0]) .
• (C14) 0 <= index_vector_dim <= rank(scatter_indices) .
• (C15) update_computation has type (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>) , where is_promotable(element_type(inputs[i]), Ei) .
• (C16) shape(inputs...) = shape(results...) .
• (C17) element_type(results[i]) = Ei for all i in [0,N) .

Примеры

// %input: [
//          [[1, 2], [3, 4], [5, 6], [7, 8]],
//          [[9, 10], [11, 12], [13, 14], [15, 16]],
//          [[17, 18], [19, 20], [21, 22], [23, 24]]
//         ]
// %scatter_indices: [[[0, 2], [1, 0], [2, 1]], [[0, 1], [1, 0], [0, 9]]]
// %update: [
//           [[[1, 1], [1, 1]], [[1, 1], [1, 1]], [[1, 1], [1, 1]]],
//           [[[1, 1], [1, 1]], [[1, 1], [1, 1]], [[1, 1], [1, 1]]]
//          ]
%result = "stablehlo.scatter"(%input, %scatter_indices, %update) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  scatter_dimension_numbers = #stablehlo.scatter<
    update_window_dims = [2, 3],
    inserted_window_dims = [0],
    scatter_dims_to_operand_dims = [1, 0],
    index_vector_dim = 2>,
  indices_are_sorted = false,
  unique_indices = false
} : (tensor<3x4x2xi64>, tensor<2x3x2xi64>, tensor<2x3x2x2xi64>) -> tensor<3x4x2xi64>
// %result: [
//           [[1, 2], [5, 6], [7, 8], [7, 8]],
//           [[10, 11], [12, 13], [14, 15], [16, 17]],
//           [[18, 19], [20, 21], [21, 22], [23, 24]]
//          ]

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

выбирать

Семантика

Produces a result tensor where each element is selected from on_true or on_false tensor based on the value of the corresponding element of pred . More formally, result[result_index] = pred_element ? on_true[result_index] : on_false[result_index] , where pred_element = rank(pred) = 0 ? pred[] : pred[result_index] . For quantized types, performs dequantize_select_quantize(pred, on_true, on_false, type(result)) .

Входы

Выходы

Constraints

• (C1) rank(pred) = 0 or shape(pred) = shape(on_true) .
• (C2) baseline_type(on_true) = baseline_type(on_false) = baseline_type(result) .

Примеры

// %pred: [[false, true], [true, false]]
// %on_true: [[1, 2], [3, 4]]
// %on_false: [[5, 6], [7, 8]]
%result = "stablehlo.select"(%pred, %on_true, %on_false) : (tensor<2x2xi1>, tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[5, 2], [3, 8]]

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

select_and_scatter

Семантика

Scatters the values from the source tensor using scatter based on the outcome of reduce_window of the input tensor using select and produces a result tensor.

The following diagram shows how elements in result are computed from operand and source using a concrete example.

More formally:

• selected_values = reduce_window_without_init(...) with the following inputs:

  • inputs = [operand].
  • window_dimensions , window_strides , and padding which are used as is.
  • base_dilations = windows_dilations = 1 .
  • body is defined as:

  def body(arg0: tensor<E>, arg1: tensor<E>) -> tensor<E>:
    return select(arg0, arg1) ? arg0 : arg1;
  

  where E = element_type(operand) , and reduce_window_without_init works exactly like reduce_window , except that the schedule of the underlying reduce (see reduce ) doesn't include init values. It is currently unspecified what happens if the corresponding window doesn't have values ( #731 ).

• result[result_index] = reduce([source_values], [init_value], [0], scatter) where:

  • source_values = [source[source_index] for source_index in source_indices] .
  • selected_index(source_index) = operand_index if selected_values[source_index] has the operand element from operand_index .
  • source_indices = [source_index for source_index in indices(source) if selected_index(source_index) = result_index] .

Входы

Выходы

Constraints

• (C1) element_type(operand) = element_type(source) .
• (C2) shape(source) = num_windows where:
  • padded_operand_shape = padding[:, 0] + shape(operand) + padding[:, 1] .
  • is_empty_window = padded_operand_shape = 0 || window_dimensions > padded_operand_shape .
  • num_windows = is_empty_window ? 0 : floor((padded_operand_shape - window_dimensions) / window_strides) + 1 .
• (C3) element_type(init_value) = element_type(operand) .
• (C4) size(window_dimensions) = rank(operand) .
• (C5) 0 < window_dimensions .
• (C6) size(window_strides) = rank(operand) .
• (C7) 0 < window_strides .
• (C8) shape(padding) = [rank(operand), 2] .
• (C9) select has type (tensor<E>, tensor<E>) -> tensor<i1> where E = element_type(operand) .
• (C10) scatter has type (tensor<E>, tensor<E>) -> tensor<E> where is_promotable(element_type(operand), E) .
• (C11) shape(operand) = shape(result) .
• (C12) element_type(result) = E .

Примеры

// %operand: [[1, 5], [2, 5], [3, 6], [4, 4]]
// %source: [[5, 6], [7, 8]]
// %init_value: 0
%result = "stablehlo.select_and_scatter"(%operand, %source, %init_value) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.compare"(%arg0, %arg1) {
      comparison_direction = #stablehlo<comparison_direction GE>
    } : (tensor<i64>, tensor<i64>) -> tensor<i1>
    "stablehlo.return"(%0) : (tensor<i1>) -> ()
}, {
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  window_dimensions = array<i64: 3, 1>,
  window_strides = array<i64: 2, 1>,
  padding = dense<[[0, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<4x2xi64>, tensor<2x2xi64>, tensor<i64>) -> tensor<4x2xi64>
// %result: [[0, 0], [0, 0], [5, 14], [7, 0]]

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

отправлять

Семантика

Sends inputs to a channel channel_id and produces a result token.

If is_host_transfer is true , then the operation transfers data to the host. Otherwise, it transfers data to another device. What this means is implementation-defined. This flag duplicates the information provided in channel_type , so in the future we are planning to only keep one of them ( #666 ).

Входы

Выходы

Constraints

• (C1) channel_type is defined as:
  • DEVICE_TO_HOST if is_host_transfer = true ,
  • DEVICE_TO_DEVICE otherwise.

Примеры

%result = "stablehlo.send"(%operand, %token) {
  channel_handle = #stablehlo.channel_handle<handle = 1, type = 2>,
  is_host_transfer = true
} : (tensor<2x2xi64>, !stablehlo.token) -> !stablehlo.token

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

shift_left

Семантика

Performs element-wise left-shift operation on the lhs tensor by rhs number of bits and produces a result tensor.

Входы

Выходы

Constraints

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

Примеры

// %lhs: [-1, 0, 1]
// %rhs: [1, 2, 3]
%result = "stablehlo.shift_left"(%lhs, %rhs): (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
// %result: [-2, 0, 8]

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

shift_right_arithmetic

Семантика

Performs element-wise arithmetic right-shift operation on the lhs tensor by rhs number of bits and produces a result tensor.

Входы

Выходы

Constraints

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

Примеры

// %lhs: [-1, 0, 8]
// %rhs: [1, 2, 3]
%result = "stablehlo.shift_right_arithmetic"(%lhs, %rhs): (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
// %result: [-1, 0, 1]

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

shift_right_logical

Семантика

Performs element-wise logical right-shift operation on the lhs tensor by rhs number of bits and produces a result tensor.

Входы

Выходы

Constraints

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

Примеры

// %lhs: [-1, 0, 8]
// %rhs: [1, 2, 3]
%result = "stablehlo.shift_right_logical"(%lhs, %rhs): (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
// %result: [9223372036854775807, 0, 1]

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

знак

Семантика

Returns the sign of the operand element-wise and produces a result tensor. More formally, for each element x , the semantics can be expressed using Python syntax as follows:

def sign(x):
  if is_integer(x):
    if compare(x, 0, LT, SIGNED): return -1
    if compare(x, 0, EQ, SIGNED): return 0
    return 1
  elif is_float(x):
    if is_nan(x): return NaN
    if compare(x, -0.0, EQ, FLOAT): return -0.0
    if compare(x, +0.0, EQ, FLOAT): return +0.0
    if compare(x, 0.0, LT, FLOAT): return -1.0
    return 1.0
  elif is_complex(x):
    if is_nan(real(x)) or is_nan(imag(x)): return (NaN, NaN)
    if compare(x, (0.0, 0.0), EQ, FLOAT): return (0.0, 0.0)
    return divide(x, convert(abs(x), type(x)))

For quantized types, performs dequantize_op_quantize(sign, operand, type(result)) .

Входы

Выходы

Constraints

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

Примеры

// Logical values: +NaN, -1.0, -0.0, +0.0, 1.0
// operand: [0x7FFFFFFFFFFFFFFF, -1.0, -0.0, 0.0, 1.0]
%result = "stablehlo.sign"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// Logical values: +NaN, -1.0, -0.0, +0.0, 1.0
// %result: [0x7FFFFFFFFFFFFFFF, -1.0, -0.0, 0.0, 1.0]

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

синус

Семантика

Performs element-wise sine operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: sin from IEEE-754.
• For complex numbers: complex sine.
• For quantized types: dequantize_op_quantize(sine, operand, type(result)) .

Входы

Выходы

Constraints

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

Примеры

// %operand: [
//            [0.0, 1.57079632],       // [0, pi/2]
//            [3.14159265, 4.71238898] // [pi, 3pi/2]
//           ]
%result = "stablehlo.sine"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[0.0, 1.0], [0.0, -1.0]]

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

кусочек

Семантика

Extracts a slice from the operand using statically-computed starting indices and produces a result tensor. start_indices contain the starting indices of the slice for each dimension, limit_indices contain the ending indices (exclusive) for the slice for each dimension, and strides contain the strides for each dimension.

More formally, result[result_index] = operand[operand_index] where operand_index = start_indices + result_index * strides .

Входы

Выходы

Constraints

• (C1) element_type(operand) = element_type(result) .
• (C2) size(start_indices) = size(limit_indices) = size(strides) = rank(operand) .
• (C3) 0 <= start_indices <= limit_indices <= shape(operand) .
• (C4) 0 < strides .
• (C5) shape(result) = ceil((limit_indices - start_indices) / strides) .

Примеры

// %operand: [
//            [0, 0, 0, 0],
//            [0, 0, 1, 1],
//            [0, 0, 1, 1]
//           ]
%result = "stablehlo.slice"(%operand) {
  start_indices = array<i64: 1, 2>,
  limit_indices = array<i64: 3, 4>,
  strides = array<i64: 1, 1>
} : (tensor<3x4xi64>) -> tensor<2x2xi64>
// % result: [
//            [1, 1],
//            [1, 1]
//           ]

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

Сортировать

Семантика

Sorts 1-dimensional slices of inputs along the dimension dimension together, according to a comparator and produces results .

Unlike similar inputs in other operations, dimension allows negative values, with the semantics described below. In the future, this may be disallowed for consistency reasons ( #1377 ).

If is_stable is true, then the sorting is stable, that is, relative order of elements considered to be equal by the comparator is preserved. For the case where there is a single input, two elements e1 and e2 are considered to be equal by the comparator if and only if comparator(e1, e2) = comparator(e2, e1) = false . See the formalization below for how this generalizes to multiple inputs.

More formally, for all result_index in index_space(results[0]) :

• adjusted_dimension = dimension >= 0 ? dimension : rank(inputs[0]) + dimension .
• result_slice = [ri0, ..., :, ..., riR-1] where riN are individual elements in result_index , and : is inserted at adjusted_dimension .
• inputs_together = (inputs[0]..., ..., inputs[N-1]...) .
• results_together[result_slice] = sort(inputs_together[result_slice], comparator_together) .
• where sort sorts a 1-dimensional slice in non-descending order expecting that comparator_together returns true if the left-hand side argument is less than the right-hand second argument.

• def comparator_together(lhs_together, rhs_together):
    args = []
    for (lhs_el, rhs_el) in zip(lhs_together, rhs_together):
      args.append(lhs_el)
      args.append(rhs_el)
    return comparator(*args)
  

• (results[0]..., ..., results[N-1]...) = results_together .

Входы