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

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

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

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

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

Программы

Program ::= {Func}

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

func.func @main(
  %image: tensor<28x28xf32>,
  %weights: tensor<784x10xf32>,
  %bias: tensor<1x10xf32>
) -> tensor<1x10xf32> {
  %0 = "stablehlo.reshape"(%image) : (tensor<28x28xf32>) -> tensor<1x784xf32>
  %1 = "stablehlo.dot"(%0, %weights) : (tensor<1x784xf32>, tensor<784x10xf32>) -> tensor<1x10xf32>
  %2 = "stablehlo.add"(%1, %bias) : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32>
  %3 = "stablehlo.constant"() {value = dense<0.0> : tensor<1x10xf32>} : () -> tensor<1x10xf32>
  %4 = "stablehlo.maximum"(%2, %3) : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32>
  "func.return"(%4): (tensor<1x10xf32>) -> ()
}

Функции

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

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

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

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

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

Типы

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

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

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

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

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

В будущем мы планируем изучить расширение типов тензоров за пределы размеров размеров и типов элементов, например, включив макеты ( #629 ) и разреженность ( #1078 ).

QuantizedTensorType ::= 'tensor' '<' Shape QuantizedTensorElementType '>'
QuantizedTensorElementType ::= '!quant.uniform' '<'
                  QuantizationStorageType
                  ['<' QuantizationStorageMin ':' QuantizationStorageMax '>']
                  ':' QuantizationExpressedType
                  [':' QuantizationDimension]
                  ',' QuantizationParameters '>'
QuantizationStorageType ::= IntegerType
QuantizationStorageMin ::= IntegerLiteral
QuantizationStorageMax ::= IntegerLiteral
QuantizationExpressedType ::= FloatType
QuantizationDimension ::= IntegerLiteral
QuantizationParameters ::= QuantizationParameter
                         | '{' QuantizationParameter {',' QuantizationParameter} '}'
QuantizationParameter ::= QuantizationScale [':' QuantizationZeroPoint]
QuantizationScale ::= FloatLiteral
QuantizationZeroPoint ::= IntegerLiteral

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

• (C1) type(storage_min) = storage_type .
• (C2) type(storage_max) = storage_type .
• (C3) min_value(storage_type) <= storage_min < storage_max <= max_value(storage_type) .
• (C4) type(scales...) = expressed_type .
• (C5) 0 < scales .
• (C6) is_finite(scales...) .
• (C7) storage_min <= zero_points <= storage_max .
• (C8) type(zero_points...) = storage_type .
• (C9) size(scales) = size(zero_points) .
• (C10) Если is_empty(quantization_dimension) , то size(scales) = 1 .
• (C11) 0 <= quantization_dimension .

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

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

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

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

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

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

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

TokenType ::= 'token'

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

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

Типы кортежей представляют собой кортежи, т. е. гетерогенные списки. Кортежи — это устаревшая функция, которая существует только для совместимости с HLO. В HLO кортежи используются для представления переменных входных и выходных данных. В StableHLO изначально поддерживаются переменные входные и выходные данные, и единственное использование кортежей в StableHLO — это всестороннее представление HLO ABI, где, например, T , tuple<T> и tuple<tuple<T>> могут существенно отличаться в зависимости от конкретной реализации. . В будущем мы планируем внести изменения в HLO ABI, которые могут позволить нам удалить типы кортежей из StableHLO ( #598 ).

TensorElementType ::= BooleanType | IntegerType | FloatType | ComplexType
BooleanType ::= 'i1'
IntegerType ::= SignedIntegerType | UnsignedIntegerType
SignedIntegerType ::= 'si2' | 'si4' | 'si8' | 'si16' | 'si32' | 'si64'
UnsignedIntegerType ::= 'ui2' | 'ui4' | 'ui8' | 'ui16' | 'ui32' | 'ui64'
FloatType ::= 'f4E2M1FN' | 'f6E2M3FN' | 'f6E3M2FN' | 'f8E3M4' | 'f8E4M3'
            | 'f8E4M3FN' | 'f8E4M3FNUZ' | 'f8E4M3B11FNUZ' | 'f8E5M2'
            | 'f8E5M2FNUZ' | 'f8E8M0FNU' | 'bf16' | 'f16' | 'f32' | 'f64'
TensorFloat32 ::= 'tf32'
ComplexType ::= 'complex' '<' ComplexElementType '>'
ComplexElementType ::= 'f32' | 'f64'

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

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

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

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

StringType ::= 'string'

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

Операции

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

%result = "stablehlo.select_and_scatter"(%operand, %source, %init_value) ({
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = "stablehlo.compare"(%arg0, %arg1) {
      comparison_direction = #stablehlo<comparison_direction GE>
    } : (tensor<i32>, tensor<i32>) -> tensor<i1>
    "stablehlo.return"(%0) : (tensor<i1>) -> ()
}, {
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
    "stablehlo.return"(%0) : (tensor<i32>) -> ()
}) {
  window_dimensions = dense<[3, 1]> : tensor<2xi64>,
  window_strides = dense<[2, 1]> : tensor<2xi64>,
  padding = dense<[[0, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<4x2xi32>, tensor<2x2xi32>, tensor<i32>) -> tensor<4x2xi32>

Константы

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

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

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

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

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

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

• (C1) is_wellformed(integer_literal, integer_type) .

FloatConstant  ::= FloatLiteral ':' FloatType
FloatLiteral   ::= SignPart IntegerPart FractionalPart ScientificPart
                 | '0x' [HexadecimalDigits]
SignPart       ::= ['-' | '+']
IntegerPart    ::= DecimalDigits
FractionalPart ::= ['.' [DecimalDigits]]
ScientificPart ::= [('e' | 'E') ['-' | '+'] DecimalDigits]

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

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

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

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

• (C1) is_wellformed(real_part, complex_element_type(complex_type)) .
• (C2) is_wellformed(imaginary_part, complex_element_type(complex_type)) .

TensorConstant ::= TensorLiteral ':' TensorType
TensorLiteral  ::= 'dense' '<' (DenseLiteral | ElementLiteral) '>'
DenseLiteral   ::= DenseDimension | DenseElements
DenseDimension ::= '[' [DenseLiteral {',' DenseLiteral}] ']'
DenseElements  ::= [ElementLiteral {',' ElementLiteral}]
ElementLiteral ::= BooleanLiteral | IntegerLiteral | FloatLiteral | ComplexLiteral

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

• (C1) has_syntax(tensor_literal, element_type(tensor_type)) , где:
  • has_syntax(element_literal: Syntax, element_type: Type) = is_wellformed(element_literal, type) .
  • has_syntax(tensor_literal: List, element_type: Type) = has_syntax(tensor_literal..., element_type) .
• (C2) has_shape(tensor_literal, shape(tensor_type)) , где:
  • has_shape(element_literal: Syntax, []) = true .
  • has_shape(tensor_literal: List, shape: List) = size(tensor_literal) = shape[0] and has_shape(tensor_literal..., shape[1:]) .
  • в противном случае false .

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

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

• (C1) has_syntax(quantized_tensor_literal, storage_type(quantized_tensor_type)) .
• (C2) has_shape(quantized_tensor_literal, shape(quantized_tensor_type)) .

StringConstant  ::= StringLiteral
StringLiteral   ::= '"' {stringCharacter | escapeSequence} '"'
stringCharacter ::= all ASCII characters except '\00', '\01', ... '\1f' and '"'
escapeSequence  ::= '\' ('"' | '\' | 'n' | 't' | (hexadecimalDigit hexadecimalDigit))

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

Операции

пресс

Семантика

Выполняет поэлементную операцию 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)) if is_complex(operand) .
  • baseline_element_type(operand) в противном случае.

Примеры

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

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

добавлять

Семантика

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

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

Входы

Выходы

Ограничения

• Если в операции используются неквантованные тензоры:
  • (C1) type(lhs) = type(rhs) = type(result) .
• Если в операции используются квантованные тензоры:
  • (C2) is_quantized(lhs) and is_quantized(rhs) and is_quantized(result) .
  • (C3) storage_type(lhs) = storage_type(rhs) = storage_type(result) .
  • (C4) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C5) (is_per_axis_quantized(lhs) or is_per_axis_quantized(rhs)) = is_per_axis_quantized(result) .
  • (C6) Если is_per_axis_quantized(lhs) , то quantization_dimension(lhs) = quantization_dimension(result) .
  • (C7) Если is_per_axis_quantized(rhs) , то quantization_dimension(rhs) = quantization_dimension(result) .

Примеры

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

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

после всего

Семантика

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

Входы

Выходы

Примеры

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

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

все_собрать

Семантика

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

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

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

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

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

Входы

Выходы

Ограничения

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

Примеры

// num_replicas: 2
// num_partitions: 1
// %operand0@(0, 0): [[1, 2], [3, 4]]
// %operand0@(1, 0): [[5, 6], [7, 8]]
// %operand1@(0, 0): [[11, 12], [13, 14]]
// %operand1@(1, 0): [[15, 16], [17, 18]]
%result:2 = "stablehlo.all_gather"(%operand0, %operand1) {
  all_gather_dim = 1 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  // channel_id = 0
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
  // use_global_device_ids = false
} : (tensor<2x2xi64>, tensor<2x2xi64>) -> (tensor<2x4xi64>, tensor<2x4xi64>)
// %result0@(0, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result0@(1, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result1@(0, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]
// %result1@(1, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]

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

all_reduce

Семантика

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

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

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

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

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

Входы

Выходы

Ограничения

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

Примеры

// num_replicas: 2
// num_partitions: 1
// %operand0@(0, 0): [1, 2, 3, 4]
// %operand0@(1, 0): [5, 6, 7, 8]
// %operand1@(0, 0): [9, 10, 11, 12]
// %operand1@(1, 0): [13, 14, 15, 16]
%result:2 = "stablehlo.all_reduce"(%operand0, %operand0) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  // channel_id = 0
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
  // use_global_device_ids = false
} : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>, tensor<4xi64>)
// %result0@(0, 0): [6, 8, 10, 12]
// %result0@(1, 0): [6, 8, 10, 12]
// %result1@(0, 0): [22, 24, 26, 28]
// %result1@(1, 0): [22, 24, 26, 28]

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

all_to_all

Семантика

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

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

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

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

Входы

Выходы

Ограничения

• (C1) 0 <= split_dimension < rank(operands...) .
• (C2) dim(operands..., split_dimension) % split_count = 0 .
• (C3) 0 <= concat_dimension < rank(operands...) .
• (C4) 0 < split_count .
• (C5) is_unique(replica_groups) .
• (C6) size(replica_groups) определяется как:
  • num_replicas если используется cross_replica .
  • num_partitions если используется cross_partition .
• (C7) 0 <= replica_groups < size(replica_groups) .
• (C8) dim(replica_groups, 1) = split_count .
• (C9) type(results...) = type(operands...) кроме случая, когда split_dimension != concat_dimension :
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count .
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count .

Примеры

// num_replicas: 2
// num_partitions: 1
// %operand1@(0, 0): [[1, 2, 3, 4],
//                    [5, 6, 7, 8]]
// %operand1@(1, 0): [[9, 10, 11, 12],
//                    [13, 14, 15, 16]]
// %operand2@(0, 0): [[17, 18, 19, 20],
//                    [21, 22, 23, 24]]
// %operand2@(1, 0): [[25, 26, 27, 28],
//                    [29, 30, 31, 32]]
%result:2 = "stablehlo.all_to_all"(%operand1, %operand2) {
  split_dimension = 1 : i64,
  concat_dimension = 0 : i64,
  split_count = 2 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
  // channel_id = 0
} : (tensor<2x4xi64>, tensor<2x4xi64>) -> (tensor<4x2xi64>, tensor<4x2xi64>)
// %result#0@(0, 0): [[1, 2], [5, 6], [9, 10], [13, 14]]
// %result#0@(1, 0): [[3, 4], [7, 8], [11, 12], [15, 16]]
// %result#1@(0, 0): [[17, 18], [21, 22], [25, 26], [29, 30]]
// %result#1@(1, 0): [[19, 20], [23, 24], [27, 28], [31, 32]]

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

и

Семантика

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

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

Входы

Выходы

Ограничения

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

Примеры

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

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

Атан2

Семантика

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

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

Входы

Выходы

Ограничения

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

Примеры

// %lhs: [0.0, 1.0, -1.0]
// %rhs: [0.0, 0.0, 0.0]
%result = "stablehlo.atan2"(%lhs, %rhs) : (tensor<3xf64>, tensor<3xf64>) -> tensor<3xf64>
// %result: [0.0, 1.57079637, -1.57079637] // [0.0, pi/2, -pi/2]

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

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

Семантика

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

def compute_sum(operand, feature_index):
  (sum,) = reduce(
      inputs=[operand],
      init_values=[constant(0, element_type(operand))],
      dimensions=[i for i in range(rank(operand)) if i != feature_index],
      body=lambda x, y: add(x, y))
  return sum

def compute_mean(operand, feature_index):
  sum = compute_sum(operand, feature_index)
  divisor = constant(size(operand) / dim(operand, feature_index),
                     element_type(operand))
  divisor_bcast = broadcast_in_dim(divisor, [], type(sum))
  return divide(sum, divisor_bcast)

def batch_norm_grad(operand, scale, mean, variance, grad_output, epsilon, feature_index):
  # Broadcast inputs to type(operand)
  scale_bcast = broadcast_in_dim(scale, [feature_index], type(operand))
  mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
  variance_bcast = broadcast_in_dim(variance, [feature_index], type(operand))
  epsilon_bcast = broadcast_in_dim(constant(epsilon, element_type(operand)), [],
                                   type(operand))

  # Perform normalization using the provided `mean` and `variance`
  # Intermediate values will be useful for computing gradients
  centered_operand = subtract(operand, mean_bcast)
  stddev = sqrt(add(variance_bcast, epsilon_bcast))
  normalized_operand = divide(centered_operand, stddev)

  # Use the implementation from batchnorm_expander.cc in XLA
  # Temporary variables have exactly the same names as in the C++ code
  elements_per_feature = broadcast_in_dim(
      constant(divide(size(operand), dim(operand, feature_index)),
               element_type(grad_output)),
      [], type(operand))
  i1 = multiply(grad_output, elements_per_feature)
  i2 = broadcast_in_dim(
      compute_sum(grad_output, feature_index), [feature_index], type(operand))
  i3 = broadcast_in_dim(
      compute_sum(multiply(grad_output, centered_operand), feature_index),
      [feature_index], type(operand))
  i4 = multiply(i3, centered_operand)
  i5 = divide(i4, add(variance_bcast, epsilon_bcast))
  i6 = subtract(subtract(i1, i2), i5)

  grad_operand =
      multiply(divide(divide(scale_bcast, stddev), elements_per_feature), i6)
  grad_scale =
      compute_sum(multiply(grad_output, normalized_operand), feature_index)
  grad_offset = compute_sum(grad_output, feature_index)

  return grad_operand, grad_scale, grad_offset

Для квантованных типов выполняет dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, mean, variance, grad_output: batch_norm_grad(operand, scale, mean, variance, grad_output, epsilon, feature_index), operand, scale, mean, variance, grad_output, type(grad_operand), type(grad_scale), type(feature_index)) .

Входы

Выходы

Ограничения

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

Примеры

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %mean: [2.0, 3.0]
// %variance: [1.0, 1.0]
// %grad_output: [
//                [[0.1, 0.1], [0.1, 0.1]],
//                [[0.1, 0.1], [0.1, 0.1]]
//               ]
%grad_operand, %grad_scale, %grad_offset =
"stablehlo.batch_norm_grad"(%operand, %scale, %mean, %variance, %grad_output) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>,
     tensor<2x2x2xf64>) -> (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>)
// %grad_operand: [
//                 [[0.0, 0.0], [0.0, 0.0]],
//                 [[0.0, 0.0], [0.0, 0.0]]
//                ]
// %grad_scale:  [0.0, 0.0]
// %grad_offset: [0.4, 0.4]

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

Семантика

Нормализует тензор operand по всем измерениям, кроме измерения feature_index , и создает result тензор. Более формально эту операцию можно выразить как декомпозицию существующих операций StableHLO с использованием синтаксиса Python следующим образом:

def batch_norm_inference(operand, scale, offset, mean, variance, epsilon, feature_index):
  # Broadcast inputs to shape(operand)
  scale_bcast = broadcast_in_dim(scale, [feature_index], type(operand))
  offset_bcast = broadcast_in_dim(offset, [feature_index], type(operand))
  mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
  variance_bcast = broadcast_in_dim(variance, [feature_index], type(operand))
  epsilon_bcast = broadcast_in_dim(constant(epsilon, element_type(operand)), [],
                                   type(operand))

  # Perform normalization using the provided `mean` and `variance` instead of
  # computing them like `batch_norm_training` does.
  centered_operand = subtract(operand, mean_bcast)
  stddev = sqrt(add(variance_bcast, epsilon_bcast))
  normalized_operand = divide(centered_operand, stddev)
  return add(multiply(scale_bcast, normalized_operand), offset_bcast)

Для квантованных типов выполняет dequantize_op_quantize(lambda operand, scale, offset, mean, variance: batch_norm_inference(operand, scale, offset, mean, variance, epsilon, feature_index), operand, scale, offset, mean, variance, type(result)) .

Входы

Выходы

Ограничения

• (C1) 0 <= feature_index < rank(operand) .
• (C2) operand , scale , offset , mean , variance и result имеют один и тот же baseline_element_type .
• (C3) size(scale) = dim(operand, feature_index) .
• (C4) size(offset) = dim(operand, feature_index) .
• (C5) size(mean) = dim(operand, feature_index) .
• (C6) size(variance) = dim(operand, feature_index) .
• (C7) baseline_type(operand) = baseline_type(result) .

Примеры

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %offset: [1.0, 1.0]
// %mean: [2.0, 3.0]
// %variance: [1.0, 1.0]
%result = "stablehlo.batch_norm_inference"(%operand, %scale, %offset, %mean, %variance) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>) -> tensor<2x2x2xf64>
// %result: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]

Batch_norm_training

Семантика

Вычисляет среднее значение и дисперсию по всем измерениям, за исключением измерения feature_index , и нормализует тензор operand , создавая output , тензоры batch_mean и batch_var . Более формально эту операцию можно выразить как декомпозицию существующих операций StableHLO с использованием синтаксиса Python следующим образом:

def compute_mean(operand, feature_index):
  (sum,) = reduce(
      inputs=[operand],
      init_values=[constant(0, element_type(operand))],
      dimensions=[i for i in range(rank(operand)) if i != feature_index],
      body=lambda x, y: add(x, y))
  divisor = constant(size(operand) / dim(operand, feature_index),
                     element_type(operand))
  divisor_bcast = broadcast_in_dim(divisor, [], type(sum))
  return divide(sum, divisor_bcast)

def compute_variance(operand, feature_index):
  mean = compute_mean(operand, feature_index)
  mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
  centered_operand = subtract(operand, mean_bcast)
  return compute_mean(mul(centered_operand, centered_operand), feature_index)

def batch_norm_training(operand, scale, offset, epsilon, feature_index):
  mean = compute_mean(operand, feature_index)
  variance = compute_variance(operand, feature_index)
  return batch_norm_inference(operand, scale, offset, mean, variance, epsilon,
                              feature_index),
         mean, variance

Для квантованных типов выполняет dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, offset: batch_norm_training(operand, scale, offset, epsilon, feature_index), operand, scale, offset, type(output), type(batch_mean), type(batch_var)) .

Входы

Выходы

Ограничения

• (C1) 0 <= feature_index < rank(operand) .
• (C2) operand , scale , offset , batch_mean , batch_var и output имеют один и тот же baseline_element_type .
• (C3) size(scale) = dim(operand, feature_index) .
• (C4) size(offset) = dim(operand, feature_index) .
• (C5) size(batch_mean) = dim(operand, feature_index) .
• (C6) size(batch_var) = dim(operand, feature_index) .
• (C7) baseline_type(output) = baseline_type(operand) .

Примеры

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %offset: [1.0, 1.0]
%output, %batch_mean, %batch_var = "stablehlo.batch_norm_training"(%operand, %scale, %offset) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>) ->
    (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>)
// %output: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]
// %batch_mean: [2.0, 3.0]
// %batch_var: [1.0, 1.0]

bitcast_convert

Семантика

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

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

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

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

Входы

Выходы

Ограничения

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

Примеры

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

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

Broadcast_in_dim

Семантика

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

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

Входы

Выходы

Ограничения

• (C1) element_type(result) определяется как:
  • element_type(operand) , если !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] if 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]

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

клетка

Семантика

Выполняет поэлементную ячейку тензора operand и создает тензор 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]

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

коллективное_вещание

Семантика

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

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

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

После этого result@process определяется следующим образом:

• operand@process_groups[i, 0] если существует i такой, что процесс находится вprocess_groups 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]]

коллективный_пермуте

Семантика

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

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

• cross_replica(source_target_pairs) , если channel_id <= 0 .
• cross_partition(source_target_pairs) , если 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 оператор реализует следующие операции IEEE-754:

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

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

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

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

Входы

Выходы

Ограничения

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs) .
• (C2) shape(lhs) = shape(rhs) = shape(result) .
• (C3) compare_type определяется как:
  • SIGNED , если is_signed_integer(element_type(lhs)) .
  • UNSIGNED , если is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)) .
  • FLOAT или TOTALORDER , если is_float(element_type(lhs)) .
  • FLOAT , если 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 . Семантика операции реализуется атрибутом decomposition . composite операцию можно заменить ее декомпозицией без изменения семантики программы. В тех случаях, когда встраивание декомпозиции не обеспечивает одинаковую семантику операций, предпочтительнее использовать 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 . Ниже показано, как это работает для сложных типов.

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

Для преобразований с плавающей запятой в целое число дробная часть усекается. Если усеченное значение не может быть представлено в целевом типе, поведение будет определено позднее ( #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]] .

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

Если 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 если 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) Если is_per_axis_quantized(rhs) , то quantization_dimension(rhs) = kernel_output_feature_dimension .
  • (C30) Если is_per_axis_quantized(result) , то quantization_dimension(result) = output_feature_dimension .
  • Если is_quantized(lhs) :
  • (C31) storage_type(lhs) = storage_type(rhs) .
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C33) Если is_per_tensor_quantized(rhs) , то is_per_tensor_quantized(result) .
  • Если !is_quantized(lhs) :
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result) .

Примеры

// %lhs: [[
//        [
//          [1], [2], [5], [6]
//        ],
//        [
//          [3], [4], [7], [8]
//        ],
//        [
//          [10], [11], [14], [15]
//        ],
//        [
//          [12], [13], [16], [17]
//        ]
//      ]]
//
// %rhs: [
//        [[[1]], [[1]], [[1]]],
//        [[[1]], [[1]], [[1]]],
//        [[[1]], [[1]], [[1]]]
//       ]
%result = "stablehlo.convolution"(%lhs, %rhs) {
  window_strides = array<i64: 4, 4>,
  padding = dense<0> : tensor<2x2xi64>,
  lhs_dilation = array<i64: 2, 2>,
  rhs_dilation = array<i64: 1, 1>,
  window_reversal = array<i1: false, false>,
  // In the StableHLO dialect, dimension numbers are encoded via:
  // `[<input dimensions>]x[<kernel dimensions>]->[output dimensions]`.
  // "b" is batch dimension, "f" is feature dimension,
  // "i" is input feature dimension, "o" is output feature dimension,
  // "0/1/etc" are spatial dimensions.
  dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
  batch_group_count = 1 : i64,
  feature_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>) -> tensor<1x2x2x1xi64>
// %result: [[
//            [[10], [26]],
//            [[46], [62]]
//          ]]

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

косинус

Семантика

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

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

Входы

Выходы

Ограничения

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

Примеры

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

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

count_leading_zeros

Семантика

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

Входы

Выходы

Ограничения

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

Примеры

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

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

custom_call

Семантика

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

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

Входы

Выходы

Примеры

%results = "stablehlo.custom_call"(%input0) {
  call_target_name = "foo",
  has_side_effect = false,
  backend_config = {bar = 42 : i32},
  api_version = 4 : i32,
  called_computations = [@foo]
} : (tensor<f64>) -> tensor<f64>

разделять

Семантика

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

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

Входы

Выходы

Ограничения

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

Примеры

// %lhs: [17.1, -17.1, 17.1, -17.1]
// %rhs: [3.0, 3.0, -3.0, -3.0]
%result = "stablehlo.divide"(%lhs, %rhs) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
// %result: [5.66666651, -5.66666651, -5.66666651, 5.66666651]

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

dot_general

Семантика

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

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

• lhs_result_dimensions = [d for d in axes(lhs) and d not in lhs_batching_dimensions and d not in lhs_contracting_dimensions] .
• rhs_result_dimensions = [d for d in axes(rhs) and d not in rhs_batching_dimensions and d not in rhs_contracting_dimensions] .
• result_batching_index + result_lhs_index + result_rhs_index = result_index где size(result_batching_index) = size(lhs_batching_dimensions) , size(result_lhs_index) = size(lhs_result_dimensions) и size(result_rhs_index) = size(rhs_result_dimensions) .
• transposed_lhs = transpose(lhs, lhs_batching_dimensions + lhs_result_dimensions + lhs_contracting_dimensions) .
• transposed_lhs_slice = slice(transposed_lhs, result_batching_index + result_lhs_index + [:, ..., :]) .
• reshaped_lhs_slice = reshape(transposed_lhs_slice, dims(lhs, lhs_contracting_dimensions)) .
• transposed_rhs = transpose(rhs, rhs_batching_dimensions + rhs_result_dimensions + rhs_contracting_dimensions) .
• transposed_rhs_slice = slice(transposed_rhs, result_batching_index + result_rhs_index + [:, ..., :]) .
• reshaped_rhs_slice = reshape(transposed_rhs_slice, dims(rhs, rhs_contracting_dimensions)) .
• dot_product = reduce( inputs=[multiply(reshaped_lhs_slice, reshaped_rhs_slice)], init_values=[constant(0, element_type(result))], dimensions=range(size(lhs_contracting_dimensions)), body=lambda x, y: add(x, y)) .

Для квантованных типов выполняет dequantize_op_quantize( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs, type(result)) .

Для гибридных квантованных типов выполняется hybrid_dequantize_then_op( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs) .

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

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

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

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

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

Пример атрибутов DotAlgorithm :

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


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


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

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

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

Входы

Выходы

Ограничения

• (C1) size(lhs_batching_dimensions) = size(rhs_batching_dimensions) .
• (C2) size(lhs_contracting_dimensions) = size(rhs_contracting_dimensions) .
• (C3) is_unique(lhs_batching_dimensions + lhs_contracting_dimensions) .
• (C4) is_unique(rhs_batching_dimensions + rhs_contracting_dimensions) .
• (C5) 0 <= lhs_batching_dimensions < rank(lhs) .
• (C6) 0 <= lhs_contracting_dimensions < rank(lhs) .
• (C7) 0 <= rhs_batching_dimensions < rank(rhs) .
• (C8) 0 <= rhs_contracting_dimensions < rank(rhs) .
• (C9) dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...) .
• (C10) dim(lhs, lhs_contracting_dimensions...) = dim(rhs, rhs_contracting_dimensions...) .
• (C11) size(precision_config) = 2 .
• (C12) shape(result) = dim(lhs, lhs_batching_dimensions) + dim(lhs, lhs_result_dimensions) + dim(rhs, rhs_result_dimensions) .
• Если в операции используются неквантованные тензоры:
  • (C13) element_type(lhs) = element_type(rhs) .
• Если в операции используются квантованные тензоры:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs) .
  • (C15) zero_points(rhs) = 0 .
  • (C16) Если 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) .
• Если !is_empty_algorithm(lhs_precision_type, rhs_precision_type, accumulation_type, lhs_component_count, rhs_component_count, num_primitive_operations allow_imprecise_accumulation) :
  • (C21) precision_config... = DEFAULT .
  • (C22) 0 < lhs_component_count .
  • (C23) 0 < rhs_component_count .
  • (C24) 0 < num_primitive_operations .

Примеры

// %lhs: [
//        [[1, 2],
//         [3, 4]],
//        [[5, 6],
//         [7, 8]]
//       ]
// %rhs: [
//        [[1, 0],
//         [0, 1]],
//        [[1, 0],
//         [0, 1]]
//       ]
%result = "stablehlo.dot_general"(%lhs, %rhs) {
  dot_dimension_numbers = #stablehlo.dot<
    lhs_batching_dimensions = [0],
    rhs_batching_dimensions = [0],
    lhs_contracting_dimensions = [2],
    rhs_contracting_dimensions = [1]
  >,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>],
  algorithm = #stablehlo.dot_algorithm<
    lhs_precision_type = tf32,
    rhs_precision_type = tf32,
    accumulation_type = f32,
    lhs_component_count = 1,
    rhs_component_count = 1,
    num_primitive_operations = 1,
    allow_imprecise_accumulation = false
  >
} : (tensor<2x2x2xi64>, tensor<2x2x2xi64>) -> tensor<2x2x2xi64>
// %result: [
//           [[1, 2],
//            [3, 4]],
//           [[5, 6],
//            [7, 8]]
//          ]

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

Dynamic_broadcast_in_dim

Семантика

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

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

Входы

Выходы

Ограничения

• (C1) element_type(result) определяется как:
  • element_type(operand) , если !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_nonexpanding_dimensions) .
• (C9) 0 <= known_expanding_dimensions < rank(operand) .
• (C10) 0 <= known_nonexpanding_dimensions < rank(operand) .

Примеры

// %operand: [
//            [1, 2, 3]
//           ]
%operand = stablehlo.constant dense<[[1, 2, 3]]> : tensor<1x3xi64>
%output_dimensions = stablehlo.constant dense<[2, 3, 2]> : tensor<3xi64>
%result = "stablehlo.dynamic_broadcast_in_dim"(%operand, %output_dimensions) {
  broadcast_dimensions = array<i64: 2, 1>,
  known_expanding_dimensions = array<i64: 0>,
  known_nonexpanding_dimensions = array<i64: 1>
} : (tensor<1x3xi64>, tensor<3xi64>) -> tensor<2x3x2xi64>
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

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

динамический_конв

Семантика

Эта операция функционально идентична операции свертки , но заполнение задается динамически с помощью padding .

Входы

Выходы

Ограничения

• (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 если 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) Если 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]]]
//        ]
// %padding: [[1, 1],
//            [1, 1]]
%result = "stablehlo.dynamic_conv"(%lhs, %rhs, %padding) {
  window_strides = array<i64: 4, 4>,
  lhs_dilation = array<i64: 2, 2>,
  rhs_dilation = array<i64: 1, 1>,
  window_reversal = array<i1: false, false>,
  dimension_numbers = #stablehlo.conv<raw
    input_batch_dimension = 0,
    input_feature_dimension = 3,
    input_spatial_dimensions = [0, 1],
    kernel_input_feature_dimension = 2,
    kernel_output_feature_dimension = 3,
    kernel_spatial_dimensions = [0, 1],
    output_batch_dimension = 0,
    output_feature_dimension = 3,
    output_spatial_dimensions = [1, 2]
  >,
  feature_group_count = 1 : i64,
  batch_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>, tensor<2x2xi64>) -> tensor<1x2x2x1xi64>
// %result: [[
//            [[1], [5]],
//            [[10], [14]]
//          ]]

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

динамический_сбор

Семантика

Эта операция функционально идентична операции сбора , при этом slice_sizes задаются динамически как значение.

Входы

Выходы

Ограничения

• (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) где:
  • batch_dim_sizes = shape(start_indices) за исключением того, что размер измерения start_indices соответствующий index_vector_dim не включается.
  • offset_dim_sizes = shape(slice_sizes) за исключением того, что размеры размеров в slice_sizes соответствующие collapsed_slice_dims , не включены.
  • combine помещает batch_dim_sizes в оси, соответствующие batch_dims , и offset_dim_sizes в оси, соответствующие 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]]
//            ]
//          ]

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

динамическая_йота

Семантика

Эта операция функционально идентична операции iota op, но форма результата задается динамически через 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]
//          ]

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

динамический_пад

Семантика

Эта операция функционально идентична операции Pad , но с edge_padding_low , edge_padding_high и interior_padding , заданными динамически как значения.

Входы

Выходы

Ограничения

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

Семантика

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

Входы

Выходы

Ограничения

• (C1) element_type(result) определяется как:
  • element_type(operand) , если !is_per_axis_quantized(operand) .
  • element_type(operand) за исключением того, что quantization_dimension(operand) и quantization_dimension(result) могут отличаться, в противном случае.
• (C2) size(operand) = size(result) .
• (C3) Если 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: [3, 2]
%result = "stablehlo.dynamic_reshape"(%operand, %output_shape) : (tensor<2x3xi64>, tensor<2xi64>) -> tensor<3x2xi64>
// %result: [[1, 2], [3, 4], [5, 6]]

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

динамический_срез

Семантика

Извлекает срез из operand используя динамически вычисляемые начальные индексы, и создает result тензор. start_indices содержит начальные индексы среза для каждого измерения, подлежащего потенциальной корректировке, а slice_sizes содержит размеры среза для каждого измерения. Более формально, result[result_index] = operand[operand_index] , где:

• 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

Семантика

Создает тензор result , который равен тензору operand , за исключением того, что срез, начинающийся с start_indices обновляется значениями в update . Более формально, result[result_index] определяется как:

• update[update_index] if 0 <= update_index < shape(update) где:
  • adjusted_start_indices = clamp(0, start_indices, shape(operand) - shape(update)) .
  • update_index = result_index - adjusted_start_indices .
• operand[result_index] в противном случае.

Входы

Выходы

Ограничения

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

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

экспоненциальный

Семантика

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

• Для чисел с плавающей запятой: exp IEEE-754.
• Для комплексных чисел: комплексная экспонента.
• Для квантованных типов: 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]]

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

экспоненциальный_минус_один

Семантика

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

• Для чисел с плавающей запятой: expm1 из IEEE-754.
• Для комплексных чисел: комплексная экспонента минус один.
• Для квантованных типов: 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_type — одно из следующих:

• FFT : прямое БПФ от комплекса к комплексу.
• IFFT : Обратное комплексное БПФ.
• RFFT : прямое БПФ от действительного к комплексному.
• IRFFT : обратное БПФ от действительного к комплексному (т.е. принимает комплексное значение, возвращает вещественное).

Более формально, учитывая функцию fft , которая принимает на вход одномерные тензоры комплексных типов, создает одномерные тензоры тех же типов, что и на выходе, и вычисляет дискретное преобразование Фурье:

Для fft_type = FFT result определяется как окончательный результат серии L вычислений, где L = size(fft_length) . Например, для 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]) .

Кроме того, учитывая функцию ifft , которая имеет ту же сигнатуру типа и вычисляет обратную функцию fft :

Для fft_type = IFFT result определяется как обратный вычислениям для fft_type = FFT . Например, для 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, ..., :]) .

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

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

(Когда дискретное преобразование Фурье вычисляется для вещественных операндов, первые N/2 + 1 элементов результата однозначно определяют остальную часть результата, поэтому результат rfft усекается, чтобы избежать вычисления избыточных элементов).

Для fft_type = RFFT result определяется как окончательный результат серии L вычислений, где L = size(fft_length) . Например, для 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]) .

Наконец, дана функция irfft , которая имеет ту же сигнатуру типа и вычисляет обратную функцию rfft :

Для fft_type = IRFFT result определяется как обратный вычислениям для fft_type = RFFT . Например, для 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) Отношения между типами operand и result элементов различаются:
  • Если fft_type = FFT , element_type(operand) и element_type(result) имеют один и тот же сложный тип.
  • Если fft_type = IFFT , element_type(operand) и element_type(result) имеют один и тот же сложный тип.
  • Если fft_type = RFFT , element_type(operand) — это тип с плавающей запятой, а element_type(result) — это сложный тип с той же семантикой с плавающей запятой.
  • Если fft_type = IRFFT , element_type(operand) — это сложный тип, а element_type(result) — это тип с плавающей запятой той же семантики с плавающей запятой.
• (C3) 1 <= size(fft_length) <= 3 .
• (C4) Если среди operand и result есть тензорное real с плавающей запятой, то shape(real)[-size(fft_length):] = fft_length .
• (C5) shape(result) = shape(operand) за исключением:
  • Если fft_type = RFFT , dim(result, -1) = dim(operand, -1) = 0 ? 0 : dim(operand, -1) / 2 + 1 .
  • Если 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)]

пол

Семантика

Выполняет поэлементное определение тензора operand и создает тензор result . Реализует операцию roundToIntegralTowardNegative из спецификации IEEE-754. Для квантованных типов выполняется 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]

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

собирать

Семантика

Собирает срезы тензора operand по смещениям, указанным в start_indices и создает result тензор.

На следующей диаграмме показано, как элементы result сопоставляются с элементами operand на конкретном примере. На диаграмме выбраны несколько примеров индексов result и подробно объяснено, каким индексам operand они соответствуют.

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

• batch_dims = [d for d in axes(result) and d not in offset_dims] .
• batch_index = result_index[batch_dims...] .
• start_index определяется как:
  • start_indices[bi0, ..., :, ..., biN] где bi — отдельные элементы в batch_index , а : вставляется в индекс index_vector_dim , если index_vector_dim < rank(start_indices) .
  • [start_indices[batch_index]] в противном случае.
• Для d_operand в axes(operand) ,
  • full_start_index[d_operand] = clamp(start_index[d_start], 0, dim(operand, d_operand) - slice_sizes[d_operand]) если d_operand = start_index_map[d_start] .
  • full_start_index[d_operand] = 0 в противном случае.
• Для d_operand в axes(operand) ,
  • full_batching_index[d_operand] = batch_index[d_start - (d_start < index_vector_dim ? 0 : 1)] , если d_operand = operand_batching_dims[i_batching] и d_start = start_indices_batching_dims[i_batching] .
  • full_batching_index[d_operand] = 0 в противном случае.
• offset_index = result_index[offset_dims...] .
• full_offset_index = [oi0, ..., 0, ..., oiN] , где oi — отдельные элементы в offset_index , а 0 вставляется в индексы из collapsed_slice_dims operand_batching_dims .
• operand_index = full_start_index + full_batching_index + full_offset_index .

Если indices_are_sorted имеет true , то реализация может предположить, что start_indices сортируются относительно start_index_map , в противном случае поведение не определено. Более формально, для всех i1 < i2 из indices(result) , full_start_index(i1) <= full_start_index(i2) .

Входы

Выходы

Ограничения

• (C1) rank(operand) = size(offset_dims) + size(collapsed_slice_dims) + size(operand_batching_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(concatenate(collapsed_slice_dims, operand_batching_dims))
• (C7) is_sorted(collapsed_slice_dims) .
• (C8) 0 <= collapsed_slice_dims < rank(operand) .
• (C9) slice_sizes[collapsed_slice_dims...] <= 1 .
• (C10) is_sorted(operand_batching_dims) .
• (C11) 0 <= operand_batching_dims < rank(operand) .
• (C12) slice_sizes[operand_batching_dims...] <= 1 .
• (C13) is_unique(start_indices_batching_dims) .
• (C14) 0 <= start_indices_batching_dims < rank(start_indices) .
• (C15) index_vector_dim not in start_indices_batching_dims .
• (C16) size(operand_batching_dims) == size(start_indices_batching_dims) .
• (C17) dim(operand, operand_batching_dims...) = dim(start_indices, start_indices_batching_dims...) .
• (C18) is_unique(concatenate(start_index_map, operand_batching_dims)) .
• (C19) 0 <= start_index_map < rank(operand) .
• (C20) size(slice_sizes) = rank(operand) .
• (C21) 0 <= slice_sizes <= shape(operand) .
• (C22) shape(result) = combine(batch_dim_sizes, offset_dim_sizes) где:
  • batch_dim_sizes = shape(start_indices) за исключением того, что размер измерения start_indices соответствующий index_vector_dim не включается.
  • offset_dim_sizes = slice_sizes за исключением того, что размеры измерений в slice_sizes соответствующие collapsed_slice_dims и operand_batching_dims не включены.
  • combine помещает batch_dim_sizes в оси, соответствующие batch_dims , и offset_dim_sizes в оси, соответствующие offset_dims .
• (C23) 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]]
//            ],
//            [
//             [[25, 26], [27, 28], [29, 30], [31, 32]],
//             [[33, 34], [35, 36], [37, 38], [39, 40]],
//             [[41, 42], [43, 44], [45, 46], [47, 48]]
//            ]
//           ]
// %start_indices: [
//                  [
//                   [[0, 0], [1, 0], [2, 1]],
//                   [[0, 1], [1, 1], [0, 9]]
//                  ],
//                  [
//                   [[0, 0], [2, 1], [2, 2]],
//                   [[1, 2], [0, 1], [1, 0]]
//                  ]
//                 ]
%result = "stablehlo.gather"(%operand, %start_indices) {
  dimension_numbers = #stablehlo.gather<
    offset_dims = [3, 4],
    collapsed_slice_dims = [1],
    operand_batching_dims = [0],
    start_indices_batching_dims = [1],
    start_index_map = [2, 1],
    index_vector_dim = 3>,
  slice_sizes = array<i64: 1, 1, 2, 2>,
  indices_are_sorted = false
} : (tensor<2x3x4x2xi32>, tensor<2x2x3x2xi64>) -> tensor<2x2x3x2x2xi32>
// %result: [
//           [
//            [
//             [[1, 2], [3, 4]],
//             [[3, 4], [5, 6]],
//             [[13, 14], [15, 16]]
//            ],
//            [
//             [[33, 34], [35, 36]],
//             [[35, 36], [37, 38]],
//             [[41, 42], [43, 44]]
//            ]
//           ],
//           [
//            [
//             [[1, 2], [3, 4]],
//             [[13, 14], [15, 16]],
//             [[21, 22], [23, 24]]
//            ],
//            [
//             [[43, 44], [45, 46]],
//             [[33, 34], [35, 36]],
//             [[27, 28], [29, 30]]
//            ]
//           ]
//          ]

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

get_dimension_size

Семантика

Возвращает размер заданного dimension operand . Более формально, result = dim(operand, dimension) . Семантика касается только компонента формы типа. Тип элемента может быть любым.

Входы

Выходы

Ограничения

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

Семантика

Извлекает элемент в index позиции кортежа operand и выдает result . Более формально, result = operand[index] .

Входы

Выходы

Ограничения

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

Примеры

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

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

если

Семантика

Создает результат выполнения ровно одной функции из true_branch или false_branch в зависимости от значения pred . Более формально, result = pred ? true_branch() : false_branch() .

Входы

Выходы

Ограничения

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

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

изображение

Семантика

Извлекает мнимую часть поэлементно из operand и создает result тензор. Более формально, для каждого элемента x : imag(x) = is_complex(x) ? imaginary_part(x) : constant(0, element_type(result)) .

Входы

Выходы

Ограничения

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

Примеры

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

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

подача

Семантика

Считывает данные с источника питания и выдает results .

Семантика infeed_config определяется реализацией.

results состоят из значений полезной нагрузки, которые идут первыми, и токена, который идет последним. В будущем мы планируем разделить полезную нагрузку и токен на два отдельных вывода для повышения ясности ( #670 ).

Входы

Выходы

Ограничения

• (C1) 0 < size(results) .
• (C2) is_empty(result[:-1]) или 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]]

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

йота

Семантика

Заполняет output тензор значениями в порядке возрастания, начиная с нуля, по измерению iota_dimension . Более формально,

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

Входы

Выходы

Ограничения

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

Семантика

Выполняет поэлементную проверку, является ли значение x конечным (т. е. не является ни +Inf, -Inf, ни NaN), и создает тензор y . Реализует операцию isFinite из спецификации IEEE-754. Для квантованных типов результат всегда true .

Входы

Выходы

Ограничения

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

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

бревно

Семантика

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

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

Входы

Выходы

Ограничения

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

Семантика

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

• Для чисел с плавающей запятой: logp1 из IEEE-754.
• Для комплексных чисел: комплексный логарифм плюс один.
• Для квантованных типов: dequantize_op_quantize(log_plus_one, operand, type(result)) .

Входы

Выходы

Ограничения

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

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

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

Семантика

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

• Для чисел с плавающей запятой: division(1, addition(1, exp(-x))) из IEEE-754.
• Для комплексных чисел: сложная логистика.
• Для квантованных типов: dequantize_op_quantize(logistic, operand, type(result)) .

Входы

Выходы

Ограничения

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

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

карта

Семантика

Применяет computation функции карты к inputs по dimensions и создает result тензор.

Более формально, result[result_index] = computation(inputs...[result_index]) .

Входы

Выходы

Ограничения

• (C1) shape(inputs...) = shape(result) .
• (C2) 0 < size(inputs) = N
• (C3) dimensions = range(rank(inputs[0])) .
• (C4) computation имеет тип (tensor<E0>, ..., tensor<EN-1>) -> tensor<E'> где Ei = element_type(inputs[i]) и 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]]

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

максимум

Семантика

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

• Для логических значений: логическое ИЛИ.
• Для целых чисел: целое число максимум.
• Для чисел с плавающей запятой: maximum из IEEE-754.
• Для комплексных чисел: лексикографический максимум для пары (real, imaginary) . Наложение порядка на комплексные числа требует удивительной семантики, поэтому в будущем мы планируем удалить поддержку комплексных чисел для этой операции ( #560 ).
• Для квантованных типов:
  • dequantize_op_quantize(maximum, lhs, rhs, type(result)) .

Входы

Выходы

Ограничения

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

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

минимум

Семантика

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

• Для логических значений: логическое И.
• Для целых чисел: целочисленный минимум.
• Для плавающих чисел: minimum IEEE-754.
• Для комплексных чисел: лексикографический минимум для пары (real, imaginary) . Наложение порядка на комплексные числа требует удивительной семантики, поэтому в будущем мы планируем удалить поддержку комплексных чисел для этой операции ( #560 ).
• Для квантованных типов:
  • dequantize_op_quantize(minimum, lhs, rhs, type(result)) .

Входы

Выходы

Ограничения

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

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

умножать

Семантика

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

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

Входы

Выходы

Ограничения

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

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

отрицать

Семантика

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

• Для целых чисел со знаком: целочисленное отрицание.
• Для целых чисел без знака: преобразование битов в целое число со знаком, отрицание целого числа, обратное преобразование целых чисел без знака.
• Для чисел с плавающей запятой: negate IEEE-754.
• Для комплексных чисел: комплексное отрицание.
• Для квантованных типов: dequantize_op_quantize(negate, operand, type(result)) .

Входы

Выходы

Ограничения

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

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

нет

Семантика

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

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

Аргументы

Выходы

Ограничения

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

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

оптимизация_барьер

Семантика

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

Аргументы

Выходы

Ограничения

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

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

или

Семантика

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

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

Входы

Выходы

Ограничения

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

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

перекормить

Семантика

Записывает inputs на выход и создает токен result .

Семантика outfeed_config определяется реализацией.

Входы

Выходы

Примеры

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

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

подушечка

Семантика

Расширяет operand путем заполнения вокруг тензора, а также между элементами тензора с заданным padding_value .

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

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

Более формально, result[result_index] определяется как:

• operand[operand_index] если result_index = edge_padding_low + operand_index * (interior_padding + 1) .
• padding_value в противном случае.

Входы

Выходы

Ограничения

• (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 текущего процесса.

Выходы

Примеры

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

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

попкнт

Семантика

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

Входы

Выходы

Ограничения

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

Примеры

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

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

власть

Семантика

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

• Для целых чисел: возведение в степень целого числа.
• Для поплавков: pow из IEEE-754.
• Для комплексных чисел: комплексное возведение в степень.
• Для квантованных типов: dequantize_op_quantize(power, lhs, rhs, type(result)) .

Входы

Выходы

Ограничения

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

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

настоящий

Семантика

Извлекает действительную часть поэлементно из operand и создает result тензор. Более формально, для каждого элемента x : real(x) = is_complex(x) ? real_part(x) : x .

Входы

Выходы

Ограничения

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

Примеры

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

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

получение

Семантика

Получает данные из канала с channel_id и выдает results .

Если is_host_transfer имеет true , операция передает данные с хоста. В противном случае он передает данные с другого устройства. Это означает, что это определяется реализацией. Этот флаг дублирует информацию, предоставленную в channel_type , поэтому в будущем мы планируем сохранить только один из них ( #666 ).

results состоят из значений полезной нагрузки, которые идут первыми, и токена, который идет последним. В будущем мы планируем разделить полезную нагрузку и токен на два отдельных вывода для повышения ясности ( #670 ).

Входы

Выходы

Ограничения

• (C1) channel_type определяется как:
  • HOST_TO_DEVICE если is_host_transfer = true ,
  • DEVICE_TO_DEVICE в противном случае.
• (C2) 0 < size(results) .
• (C3) is_empty(result[:-1]) или 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)

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

уменьшать

Семантика

Применяет body функции сокращения к inputs и init_values ​​вдоль dimensions и создает тензоры results .

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

Более формально, results...[j0, ..., jR-1] = reduce(input_slices_converted) , где:

• input_slices = inputs...[j0, ..., :, ..., jR-1] , где : вставляются в 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) для некоторого schedule двоичного дерева, где:
  • exec(node) = body(exec(node.left), exec(node.right)) .
  • exec(leaf) = leaf.value .
• schedule — это определяемое реализацией полное двоичное дерево, упорядоченный обход которого состоит из:
  • Значения input_slices_converted...[index] для всех index в index_space(input_slices_converted) в возрастающем лексикографическом порядке index .
  • Перемежается с определяемым реализацией количеством init_values_converted в позициях, определенных реализацией.

Входы

Выходы

Ограничения

• (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 имеет тип (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>) где is_promotable(element_type(inputs[i]), Ei) .
• (C7) shape(results...) = shape(inputs...) за исключением того, что размеры inputs... соответствующие dimensions не включены.
• (C8) element_type(results[i]) = Ei для всех i в [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]

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

уменьшить_точность

Семантика

Выполняет поэлементное преобразование operand в другой тип с плавающей запятой, который использует exponent_bits и mantissa_bits , и обратно в исходный тип с плавающей запятой и создает output тензор.

Более формально:

• Биты мантиссы исходного значения обновляются для округления исходного значения до ближайшего значения, которое можно представить с помощью mantissa_bits с использованием семантики roundToIntegralTiesToEven .
• Затем, если mantissa_bits меньше количества битов мантиссы исходного значения, биты мантиссы усекаются до mantissa_bits .
• Затем, если биты экспоненты промежуточного результата не помещаются в диапазон, предоставленный exponent_bits , промежуточный результат переполняется до бесконечности, используя исходный знак, или уменьшается до нуля, используя исходный знак.
• Для квантованных типов выполняет dequantize_op_quantize( lambda operand: reduce_precision(operand, exponent_bits, mantissa_bits), operand, type(result)) .

Входы

Выходы

Ограничения

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

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

уменьшить_разброс

Семантика

Внутри каждой группы процессов в сетке процессов StableHLO выполняет сокращение, используя computations , над значениями тензора operand каждого процесса, разбивает результат сокращения по scatter_dimension на части и распределяет разделенные части между процессами для получения 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 :

• 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] для всех sender в process_group , где receiver_index = process_group.index(receiver) .

Входы

Выходы

Ограничения

• (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) определяется как:
  • num_replicas если используется cross_replica .
  • num_replicas если используется cross_replica_and_partition .
  • num_processes если используется flattened_ids .
• (C5) 0 <= replica_groups < size(replica_groups) .
• (C6) Если use_global_device_ids = true , то channel_id > 0 .
• (C7) computation имеет тип (tensor<E>, tensor<E>) -> (tensor<E>) где is_promotable(element_type(operand), E) .
• (C8) shape(result) = shape(operand) за исключением:
  • 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]]

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

уменьшить_окно

Семантика

Применяет body функции сокращения к окнам inputs и init_values ​​и выдает results .

На следующей диаграмме показано, как элементы results... вычисляются на основе inputs... на конкретном примере.

Более формально, results...[result_index] = reduce(windows, init_values, axes(inputs...), body) (см. сокращение ), где:

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

Входы

Выходы

Ограничения

• (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 имеет тип (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>) где is_promotable(element_type(inputs[i]), Ei) .
• (C14) same(shape(results...)) .
• (C15) shape(results[0]) = num_windows где:
  • 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 для всех i в [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]]

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

остаток

Семантика

Вычисляет поэлементный остаток тензоров делимого lhs и rhs делителя и создает result тензор.

Более формально, знак результата берется из делимого, а абсолютное значение результата всегда меньше абсолютного значения делителя. Остаток рассчитывается как lhs - d * rhs , где d определяется выражением:

• Для целых чисел: stablehlo.divide(lhs, rhs) .
• Для чисел с плавающей запятой: division(lhs, rhs) из IEEE-754 с атрибутом округления roundTowardZero .
• Для комплексных чисел: подлежит уточнению ( #997 ).
• Для квантованных типов:
  • dequantize_op_quantize(remainder, lhs, rhs, type(result)) .

Для типов элементов с плавающей точкой эта операция в отличие от remainder операции из спецификации IEEE-754, где d является интегральным значением, ближайшим к точным значениям lhs/rhs со связями до ровности.

Входы

Выходы

Ограничения

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

Семантика

Создает replica_id текущего процесса.

Выходы

Примеры

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

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

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

Семантика

Выполняет изменение тензора operand до result . Концептуально, это равносильно тому же каноническому представлению, но потенциально изменяет форму, например, от tensor<2x3xf32> до tensor<3x2xf32> или tensor<6xf32> .

Более формально, result[result_index] = operand[operand_index] , где result_index и operand_index имеют одинаковую позицию в лексикографическом упорядочении index_space(result) и index_space(operand) .

Входы

Выходы

Ограничения

• (C1) element_type(result) определяется как:
  • element_type(operand) , если !is_per_axis_quantized(operand) .
  • element_type(operand) за исключением того, что quantization_dimension(operand) и quantization_dimension(result) могут отличаться, в противном случае.
• (C2) size(operand) = size(result) .
• (C3) Если 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]]

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

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

Семантика

Отправляет порядок элементов в operand вдоль указанных dimensions и дает тензор result . Более формально, result[result_index] = operand[operand_index] где:

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

Входы

Выходы

Ограничения

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

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

звонок

Семантика

Генерирует случайные числа, используя алгоритм rng_distribution и дает result тензор заданной shape .

Если rng_distribution = UNIFORM , то случайные числа генерируются после равномерного распределения по интервалу [a, b) . Если a >= b , поведение не определен.

Если rng_distribution = NORMAL , то случайные числа генерируются после нормального распределения со средним = a и стандартным отклонениями = b . Если b < 0 , поведение не определен.

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

В разговорах со многими заинтересованными сторонами этот OP появился как эффективно устаревший, поэтому в будущем мы планируем изучить его удаление ( #597 ).

Входы

Выходы

Ограничения

• (C1) element_type(a) = element_type(b) = element_type(result) .
• (C2) Если rng_distribution = NORMAL , то 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

Семантика

Возвращает output заполненный равномерными случайными битами и обновленным выходным состоянием output_state используя алгоритм генератора псевдордомов rng_algorithm с учетом начального состояния initial_state . Выход гарантированно является детерминированной функцией initial_state , но не гарантированно будет детерминированным между реализациями.

rng_algorithm является одним из следующих:

• DEFAULT : определяемый реализацией алгоритм.
• THREE_FRY : определяемый реализацией вариант алгоритма Threefry.*
• PHILOX : определяемый реализацией вариант алгоритма Philox.*

* См.: Salmon et al. SC 2011. Параллельные случайные числа: просто 1, 2, 3.

Входы

Выходы

Ограничения

• (C1) type(initial_state) = type(output_state) .
• (C2) size(initial_state) определяется как:
  • Реализация DEFATION IF rng_algorithm = DEFAULT .
  • 2 Если rng_algorithm = THREE_FRY .
  • 2 или 3 , если 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

Семантика

Выполняет элемент, округление к ближайшему целому числу, разрывает связи с нулем, на тензоре operand и дает тензор result . Реализует операцию roundToIntegralTiesToAway из спецификации IEEE-754. Для квантованных типов выполняет dequantize_op_quantize(round_nearest_afz, operand, type(result)) .

Входы

Выходы

Ограничения

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

Семантика

Выполняет элементное округление в сторону ближайшего целого числа, разбивая связи с ровным целым числом, на тензоре operand и дает тензор result . Реализует операцию roundToIntegralTiesToEven из спецификации IEEE-754. Для квантовых типов выполняет dequantize_op_quantize(round_nearest_even, operand, type(result)) .

Входы

Выходы

Ограничения

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

Семантика

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

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

Входы

Выходы

Ограничения

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

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

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

Семантика

Получает тензоры results , которые равны тенденциям inputs , за исключением того, что несколько срезов, указанных с помощью scatter_indices обновляются с updates значений с использованием update_computation .

На следующей диаграмме показано, как элементы в updates... Карта по элементам в results... используя конкретный пример. Диаграмма выбирает несколько примеров updates... Индексы и подробно объясняет, какие results... индексы, с которыми они соответствуют.

Более формально для всех update_index в 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 определяется как:
  • scatter_indices[si0, ..., :, ..., siN] , где si - это отдельные элементы в update_scatter_index и : вставлены в индексе index_vector_dim , если index_vector_dim < rank(scatter_indices) .
  • [scatter_indices[update_scatter_index]] в противном случае.
• Для d_input в axes(inputs[0]) ,
  • full_start_index[d_input] = start_index[d_start] если d_input = scatter_dims_to_operand_dims[d_start] .
  • full_start_index[d_input] = 0 в противном случае.
• Для d_input в axes(inputs[0]) ,
  • full_batching_index[d_input] = update_scatter_index[d_start - (d_start < index_vector_dim ? 0 : 1)] , если d_input = input_batching_dims[i_batching] и d_start = scatter_indices_batching_dims[i_batching] .
  • full_batching_index[d_input] = 0 в противном случае.
• update_window_index = update_index[update_window_dims...] .
• full_window_index = [wi0, ..., 0, ..., wiN] , где wi являются отдельными элементами в update_window_index , а 0 вставлено по индексам из inserted_window_dims и input_batching_dims .
• result_index = full_start_index + full_batching_index + full_window_index .

Учитывая это, results = exec(schedule, inputs) , где:

• schedule -это определяемая реализацией перестановка index_space(updates[0]) .
• exec([update_index, ...], results) = exec([...], updated_results) где:
  • Если result_index находится в границах для 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 - это копия results с results...[result_index] установлен на updated_values...
  • В противном случае
  • updated_results = results .
• exec([], results) = results .

Если indices_are_sorted является true , то реализация может предположить, что scatter_indices сортируется по отношению к scatter_dims_to_operand_dims , в противном случае поведение не определен. Более формально для всех i1 < i2 из indices(result) , full_start_index(i1) <= full_start_index(i2) .

Если unique_indices true , то реализация может предположить, что все индексы result_index разбросаны, являются уникальными. Если unique_indices true , но рассеянные показатели не являются уникальными, то поведение не определен.

Входы

Выходы

Ограничения

• (C1) same(shape(inputs...)) .
• (C2) `rank (inputs [0]) = size (update_window_dims) + size (inserted_window_dims)
  • размер (input_batching_dims) `.
• (C3) same(shape(updates...)) .
• (C4) shape(updates[0]) = combine(update_scatter_dim_sizes, update_window_dim_sizes) где:
  • update_scatter_dim_sizes = shape(scatter_indices) за исключением того, что размер размерности scatter_indices , соответствующий index_vector_dim , не включен.
  • update_window_dim_sizes <= shape(inputs[0]) за исключением того, что размеры измерений в inputs[0] соответствующие inserted_window_dims и input_batching_dims не включены.
  • combine puts update_scatter_dim_sizes по осям, соответствующим update_scatter_dims и update_window_dim_sizes по осям, соответствующим 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(concatenate(inserted_window_dims, input_batching_dims))
• (C10) is_sorted(inserted_window_dims) .
• (C11) 0 <= inserted_window_dims < rank(inputs[0]) .
• (C12) is_sorted(input_batching_dims) .
• (C13) 0 <= input_batching_dims < rank(inputs[0])) .
• (C14) is_unique(scatter_indices_batching_dims) .
• (C15) 0 <= scatter_indices_batching_dims < rank(scatter_indices) .
• (C16) index_vector_dim not in scatter_indices_batching_dims .
• (C17) size(input_batching_dims) == size(scatter_indices_batching_dims) .
• (C18) dim(inputs[0], input_batching_dims...) = dim(scatter_indices, scatter_indices_batching_dims...) .
• (C19) size(scatter_dims_to_operand_dims) = index_vector_dim < rank(scatter_indices) ? dim(scatter_indices, index_vector_dim) : 1 .
• (C20) is_unique(concatenate(scatter_dims_to_operand_dims, input_batching_dims)) .
• (C21) 0 <= scatter_dims_to_operand_dims < rank(inputs[0]) .
• (C22) 0 <= index_vector_dim <= rank(scatter_indices) .
• (C23) update_computation имеет тип (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>) , где is_promotable(element_type(inputs[i]), Ei) .
• (C24) shape(inputs...) = shape(results...) .
• (C25) element_type(results[i]) = Ei для всех i в [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]]
//          ],
//          [
//           [[25, 26], [27, 28], [29, 30], [31, 32]],
//           [[33, 34], [35, 36], [37, 38], [39, 40]],
//           [[41, 42], [43, 44], [45, 46], [47, 48]]
//          ]
//         ]
// %scatter_indices: [
//                    [
//                     [[0, 0], [1, 0], [2, 1]],
//                     [[0, 1], [1, 1], [0, 9]]
//                    ],
//                    [
//                     [[0, 0], [2, 1], [2, 2]],
//                     [[1, 2], [0, 1], [1, 0]]
//                    ]
//                   ]
// %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 = [3, 4],
    inserted_window_dims = [1],
    input_batching_dims = [0],
    scatter_indices_batching_dims = [1],
    scatter_dims_to_operand_dims = [2, 1],
    index_vector_dim = 3>,
  indices_are_sorted = false,
  unique_indices = false
} : (tensor<2x3x4x2xi64>, tensor<2x2x3x2xi64>, tensor<2x2x3x2x2xi64>) -> tensor<2x3x4x2xi64>
// %result: [
//           [
//            [[3, 4], [6, 7], [6, 7], [7, 8]],
//            [[9, 10],[11, 12], [15, 16], [17, 18]],
//            [[17, 18], [19, 20], [22, 23], [24, 25]]
//           ],
//           [
//            [[25, 26], [28, 29], [30, 31], [31, 32]],
//            [[35, 36], [38, 39], [38, 39], [39, 40]],
//            [[41, 42], [44, 45], [46, 47], [47, 48]]
//           ]
//          ]

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

выбирать

Семантика

Создает тензор result , в котором каждый элемент выбран из тензора on_true или on_false на основе значения соответствующего элемента pred . Более формально, result[result_index] = pred_element ? on_true[result_index] : on_false[result_index] , где pred_element = rank(pred) = 0 ? pred[] : pred[result_index] . Для квантовых типов выполняет dequantize_select_quantize(pred, on_true, on_false, type(result)) .

Входы

Выходы

Ограничения

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

Семантика

Разбросает значения из тензора source , используя scatter на основе результата reduce_window input тензора, используя select и дает тензор result .

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

Более формально:

• selected_values = reduce_window_without_init(...) со следующими входами:

  • inputs = [operand].
  • window_dimensions , window_strides и padding , которые используются как есть.
  • base_dilations = windows_dilations = 1 .
  • body определяется как:

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

  где E = element_type(operand) и reduce_window_without_init работает точно так же, как reduce_window , за исключением того, что schedule базового reduce (см. DEMUT ) не включает значения init. В настоящее время не указано, что происходит, если в соответствующем окне нет значений ( #731 ).

• result[result_index] = reduce([source_values], [init_value], [0], scatter) Где:

  • source_values = [source[source_index] for source_index in source_indices] .
  • selected_index(source_index) = operand_index if selected_values[source_index] имеет элемент operand от operand_index .
  • source_indices = [source_index for source_index in indices(source) if selected_index(source_index) = result_index] .

Входы

Выходы

Ограничения

• (C1) element_type(operand) = element_type(source) .
• (C2) shape(source) = num_windows , где:
  • 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 ISESS TYPE (tensor<E>, tensor<E>) -> tensor<i1> где E = element_type(operand) .
• (C10) scatter имеет тип (tensor<E>, tensor<E>) -> tensor<E> где 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]]

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

отправлять

Семантика

Отправляет inputs в канал channel_id и создает токен result .

Если is_host_transfer является true , то операция передает данные на хост. В противном случае он передает данные на другое устройство. Что это значит, определяется реализацией. Этот флаг дублирует информацию, представленную в channel_type , поэтому в будущем мы планируем сохранить только один из них ( #666 ).

Входы

Выходы

Ограничения

• (C1) channel_type определяется как:
  • DEVICE_TO_HOST if is_host_transfer = true ,
  • DEVICE_TO_DEVICE иначе.

Примеры

%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

Семантика

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

Входы

Выходы

Ограничения

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

Семантика

Выполняет элементную арифметическую операцию правого сдвига на тензоре lhs по количеству битов rhs и дает тензор result .

Входы

Выходы

Ограничения

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

Семантика

Выполняет элементную логическую операцию правого смены на тензоре lhs по количеству битов rhs и дает тензор result .

Входы

Выходы

Ограничения

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

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

знак

Семантика

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

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

Для квантовых типов выполняет dequantize_op_quantize(sign, operand, type(result)) .

Входы

Выходы

Ограничения

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

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

синус

Семантика

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

• Для поплавков: sin от IEEE-754.
• Для сложных чисел: сложный синус.
• Для квантовых типов: dequantize_op_quantize(sine, 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.sine"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[0.0, 1.0], [0.0, -1.0]]

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

кусочек

Семантика

Извлекает срез из operand используя статически выпускаемые начальные индексы и дает тензор result . start_indices содержат начальные индексы среза для каждого измерения, limit_indices содержат конечные индексы (исключительно) для среза для каждого измерения, а strides содержат шаги для каждого измерения.

Более формально, result[result_index] = operand[operand_index] , где operand_index = start_indices + result_index * strides .

Входы