StableHLO の仕様

StableHLO は、機械学習（ML）モデルの高レベル オペレーション（HLO）用のオペレーション セットです。StableHLO は、異なる ML フレームワークと ML コンパイラ間のポータビリティ レイヤとして機能します。StableHLO プログラムを生成する ML フレームワークは、StableHLO プログラムを使用する ML コンパイラと互換性があります。

Google の目標は、さまざまな ML フレームワーク（TensorFlow、JAX、PyTorch など）と ML コンパイラ（XLA、IREE など）間の相互運用性を高めることで、ML 開発を簡素化し、加速させることです。このドキュメントでは、その目的のために StableHLO プログラミング言語の仕様を説明します。

この仕様には、3 つの主要なセクションがあります。まず、プログラムのセクションでは、StableHLO 関数で構成される StableHLO プログラムの構造について説明します。StableHLO 関数自体は StableHLO オペレーションで構成されます。この構造内の [Ops] セクションには、個々のオペレーションのセマンティクスを指定します。[実行] セクションには、プログラム内で一緒に実行されるこれらのオペレーションのすべてのセマンティクスが示されています。最後に、表記セクションでは、仕様全体で使用される表記について説明します。

StableHLO の以前のリリースの仕様を表示するには、目的のタグ付きリリースでリポジトリを開きます。たとえば、StableHLO v0.19.0 仕様。StableHLO のマイナー バージョンの増加ごとに発生した変更を確認するには、VhloDialect.td のバージョンログをご覧ください。

プログラム

Program ::= {Func}

StableHLO プログラムは、任意の数の StableHLO 関数で構成されます。以下は、3 つの入力（%image、%weights、%bias）と 1 つの出力を持つ関数 @main を含むプログラムの例です。関数本体には 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 識別子は多くのプログラミング言語の識別子に似ていますが、2 つの特殊性があります。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 つのディメンション（つまり 2 つの軸）があり、サイズは 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 のセマンティクス（型、値、量子化されたテンソル型にゼロポイントが 1 つだけか複数ある可能性があるかなど）については、現在も議論が続いています。このディスカッションの結果に基づいて、ゼロポイントに関する仕様は今後変更される可能性があります（#1405）。

現在進行中のもう一つの議論では、QuantizationStorageMin と QuantizationStorageMax のセマンティクスについて、これらの値と量子化テンソルの値に制約を適用すべきかどうかを判断しています（#1406）。

最後に、未知のディメンション サイズの表現（#1407）と同様に、未知のスケールとゼロポイントの表現を検討する予定です。

量子化テンソル型は、量子化された要素を持つテンソルを表します。これらのテンソルは、通常の要素型ではなく、量子化された要素型を持つ点を除き、通常のテンソルとまったく同じです。

量子化されたテンソルでは、量子化はテンソル単位で行うことができます。つまり、テンソル全体に 1 つの scale と zero_point を設定します。または、軸単位で行うこともできます。つまり、特定のディメンション quantization_dimension のスライスごとに 1 組の scales と zero_points を設定します。より正式には、軸ごとの量子化を使用するテンソル t には、quantization_dimension の dim(t, 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 型のスカラー値を tensor<T> 型の 0 次元テンソル値で表現するのが一般的です）。

• ブール値型は、ブール値 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: IEEE-754 規則に従う 8 ビットの浮動小数点数。
  • f8E4M3FN 型と f8E5M2 型は、ディープラーニング用の FP8 形式で説明されている FP8 形式の E4M3 エンコードと E5M2 エンコードに対応しています。
  • f8E4M3FNUZ 型と f8E5M2FNUZ 型は、ディープ ニューラル ネットワーク用の 8 ビット数値形式で説明されている FP8 形式の E4M3 エンコードと E5M2 エンコードに対応しています。
  • ハイブリッド 8 ビット浮動小数点数（HFP8）トレーニングとディープ ニューラル ネットワークの推論で説明されている FP8 形式の E4M3 エンコードに対応する f8E4M3B11FNUZ タイプ。
  • BFloat16: Cloud TPU で高パフォーマンスを発揮させる秘訣で説明されている bfloat16 形式に対応する bf16 型。
  • f16、f32、f64 型は、IEEE 754 標準で説明されている binary16（「半精度」）、binary32（「単精度」）、binary64（「倍精度」）の各形式に対応しています。
  • tf32 型は TensorFloat32 形式に対応しており、StableHLO では限定的にサポートされています。
  • OCP マイクロスケーリング形式の仕様で説明されている f4E2M1FN、f6E2M3FN、f6E3M2FN、f8E8M0FNU MX（マイクロスケーリング）タイプ。
• 複素型は、同じ要素型の実部と虚部を持つ複素値を表します。サポートされている複合型は complex<f32>（どちらも f32 型）と complex<f64>（どちらも f64 型）です。

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

関数型は、名前付き関数と匿名関数の両方を表します。入力型（-> の左側の型のリスト）と出力の型（-> の右側の型のリスト）があります。多くのプログラミング言語では関数型がファースト クラスですが、StableHLO ではそうではありません。

StringType ::= 'string'

String 型はバイトのシーケンスを表します。多くのプログラミング言語とは異なり、StableHLO では文字列型はファーストクラスではなく、プログラム要素の静的メタデータを指定する場合にのみ使用されます。

運用

StableHLO オペレーション（オペレーションとも呼ばれる）は、ML モデルのクローズド セットの大まかなオペレーションを表します。前述のように、StableHLO 構文は MLIR に大きく影響を受けています。これは必ずしも最も人間工学的な代替手段ではありませんが、ML フレームワークと ML コンパイラ間の相互運用性を高めるという StableHLO の目標には最適です。

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

StableHLO オペレーション（オペレーションとも呼ばれる）には、名前、入出力、署名があります。この名前は、stablehlo. 接頭辞と、サポートされているいずれかのオペレーションを一意に識別するニーモニックで構成されます。サポートされているすべてのオペレーションの一覧については、以下をご覧ください。

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

オペレーションは入力を使用し、出力を生成します。入力は、入力値（実行時に計算される）、入力関数（StableHLO では関数がファーストクラス値ではないため、静的に提供される）、入力属性（これも静的に提供される）に分類されます。op によって消費、生成される入力と出力の種類は、そのニモニックによって異なります。たとえば、add オペレーションは 2 つの入力値を使用し、1 つの出力値を生成します。これに対し、select_and_scatter op は 3 つの入力値、2 つの入力関数、3 つの入力属性を使用します。

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

入力関数（別名: 匿名関数）は、名前付き関数と非常によく似ていますが、1）識別子がない（「匿名」という名前の由来）ことと、2）出力型を宣言しない（出力型は関数内の return オペレーションから推論される）点が異なります。

入力関数の構文には、MLIR との互換性を確保するために現在使用されていない部分（上記の Unused 生産を参照）が含まれています。MLIR には、ジャンプ演算子で接続された複数の演算子の「ブロック」を持つことができる、より一般的な「リージョン」のコンセプトがあります。これらのブロックには、Unused 本番環境に対応する ID が付いているため、ブロック同士を区別できます。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 シグネチャは、MLIR との互換性を確保するために、意図的に StableHLO 構文の一部になっています。

以下に、メモニクスが select_and_scatter の演算子の例を示します。3 つの入力値（%operand、%source、%init_value）、2 つの入力関数、3 つの入力属性（window_dimensions、window_strides、padding）を使用します。op のシグネチャには入力値の型のみが含まれます（インラインで提供される入力関数と属性の型は含まれません）。

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

整数定数は、10 進数または 16 進数表記の文字列を使用して整数値を表します。バイナリやオクタルなどの他の基数はサポートされていません。整数定数には次の制約があります。

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

浮動小数点定数は、10 進数または科学的記数法を使用する文字列で浮動小数点値を表します。また、16 進数表記を使用して、対応する型の浮動小数点形式で基盤となるビットを直接指定することもできます。浮動小数点定数には次の制約があります。

• （C1）16 進数以外の表記が使用されている場合、is_wellformed(float_literal, float_type)。
• （C2）16 進数表記を使用している場合、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 文字とエスケープ シーケンスを使用して指定されたバイトで構成されます。エンコードに依存しないため、これらのバイトの解釈は実装で定義されます。文字列リテラルの型は string です。

運用

abs

セマンティクス

operand テンソルに対して要素ごとの絶対値演算を実行し、result テンソルを生成します。要素のタイプに応じて、次の処理を行います。

• 符号付き整数の場合: 整数の剰余。
• 浮動小数点数: IEEE-754 の abs。
• 複素数の場合: 複素数モジュラス。
• 量子化された型の場合: dequantize_op_quantize(abs, operand, type(result))。

入力

出力

制約

• （C1）shape(result) = shape(operand)。
• （C2）baseline_element_type(result) は次のように定義されます。
  • is_complex(operand) の場合、complex_element_type(element_type(operand))。
  • そうでない場合は baseline_element_type(operand)。

例

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

その他の例

追加

セマンティクス

2 つのテンソル lhs と rhs の要素ごとの加算を行い、result テンソルを生成します。要素のタイプに応じて、次の処理を行います。

• ブール値の場合: 論理 OR。
• 整数の場合: 整数の加算。
• 浮動小数点数の場合: IEEE-754 の addition。
• 複素数の場合: 複素数加算。
• 量子化された型の場合: 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]]

その他の例

after_all

セマンティクス

inputs を生成するオペレーションが、result に依存するオペレーションの前に実行されるようにします。このオペレーションを実行しても何も行われません。このオペレーションは、result から inputs へのデータ依存関係を確立するためにのみ存在します。

入力

出力

例

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

その他の例

all_gather

セマンティクス

StableHLO プロセス グリッドの各プロセス グループ内で、各プロセスの operands テンソルの値を all_gather_dim に沿って連結し、results テンソルを生成します。

このオペレーションでは、StableHLO プロセス グリッドが、次のように定義された process_groups に分割されます。

• channel_id <= 0 and use_global_device_ids = false の場合、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 = true の場合、flattened_ids(replica_groups)

その後、各 process_group 内で次の操作を行います。

• process_group のすべての receiver に対して operands...@receiver = [operand@sender for sender in process_group]。
• process_group のすべての process に対して results...@process = concatenate(operands...@process, all_gather_dim)。

入力

出力

制約

• （C1）0 <= all_gather_dim < rank(operands...)。
• （C2）is_unique(replica_groups)。
• （C3）size(replica_groups) は次のように定義されます。
  • cross_replica を使用する場合は num_replicas。
  • cross_replica_and_partition を使用する場合は num_replicas。
  • flattened_ids を使用する場合は num_processes。
• （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 プロセス グリッドの各プロセス グループ内で、各プロセスの operands テンソルの値にリダクション関数 computation を適用し、results テンソルを生成します。

このオペレーションは、StableHLO プロセス グリッドを次のように定義された process_groups に分割します。

• channel_id <= 0 and use_global_device_ids = false の場合、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 = true の場合、flattened_ids(replica_groups)

その後、各 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) は次のように定義されます。
  • cross_replica を使用する場合は num_replicas。
  • cross_replica_and_partition を使用する場合は num_replicas。
  • flattened_ids を使用する場合は num_processes。
• （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 プロセス グリッド内の各プロセス グループ内で、split_dimension に沿って operands テンソルの値を分割し、分割された部分をプロセス間で分散し、分散された部分を concat_dimension に沿って連結して results テンソルを生成します。このオペレーションでは、StableHLO プロセス グリッドが、次のように定義された process_groups に分割されます。

• channel_id <= 0 の場合、cross_replica(replica_groups)。
• channel_id > 0 の場合は cross_partition(replica_groups)。

その後、各 process_group 内で次の操作を行います。

• process_group 内のすべての sender に対して split_parts...@sender = split(operands...@sender, split_count, split_dimension)。
• 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) は次のように定義されます。
  • cross_replica を使用する場合は num_replicas。
  • cross_partition を使用する場合は num_partitions。
• （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]]

その他の例

と

セマンティクス

2 つのテンソル lhs と rhs の要素ごとの AND を実行し、result テンソルを生成します。要素のタイプに応じて、次の処理を行います。

• ブール値の場合: 論理 AND。
• 整数の場合: ビット演算 AND。

入力

出力

制約

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

その他の例

atan2

セマンティクス

lhs テンソルと rhs テンソルの要素ごとの atan2 演算を実行し、result テンソルを生成します。要素のタイプに応じて、次の処理を行います。

• 浮動小数点数の場合: IEEE-754 の atan2。
• 複素数の場合: 複素 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_grad

セマンティクス

grad_output からバックプロパゲートされる batch_norm_training の複数の入力の勾配を計算し、grad_operand、grad_scale、grad_offset テンソルを生成します。より正式には、このオペレーションは、次のように Python 構文を使用して既存の StableHLO オペレーションへの分解として表現できます。

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]

batch_norm_inference

セマンティクス

feature_index 次元を除くすべての次元で operand テンソルを正規化し、result テンソルを生成します。より正式には、このオペレーションは、Python 構文を使用して既存の StableHLO オペレーションへの分解として次のように表現できます。

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 テンソルを生成します。より正式には、このオペレーションは、Python 構文を使用して既存の StableHLO オペレーションへの分解として次のように表現できます。

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] で、axes(operand) 内のすべての d について次のようにします。

• dim(operand, d) = 1 の場合、operand_index[d] = 0。
• それ以外の場合は 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）axes(operand) 内のすべての d について:
  • 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]
//            ]
//          ]

その他の例

ケース

セマンティクス

index の値に応じて、branches から関数を 1 つだけ実行することで、出力を生成します。より正式には、result = selected_branch() のようにします。

• 0 <= index < size(branches) の場合、selected_branch = branches[index]。
• それ以外の場合は 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 テンソルを生成します。要素のタイプに応じて、次の操作を行います。

• 浮動小数点数の場合: IEEE-754 の rootn(x, 3)。
• 複素数の場合: 複素立方根。
• 量子化型の場合: dequantize_op_quantize(cbrt, operand, type(result))

入力

出力

制約

• （C1）baseline_type(operand) = baseline_type(result)。

例

// %operand: [0.0, 1.0, 8.0, 27.0]
%result = "stablehlo.cbrt"(%operand) : (tensor<4xf64>) -> tensor<4xf64>
// %result: [0.0, 1.0, 2.0, 3.0]

その他の例

ceil

セマンティクス

operand テンソルの要素単位の ceil を実行し、result テンソルを生成します。IEEE-754 仕様の roundToIntegralTowardPositive オペレーションを実装します。量子化された型の場合は、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]

その他の例

cholesky

セマンティクス

行列のバッチの Cholesky 分解を計算します。

より正式には、index_space(result) のすべての i について、result[i0, ..., iR-3, :, :] は a[i0, ..., iR-3, :, :] の Cholesky 分解であり、下三角（lower が true の場合）または上三角（lower が false の場合）行列のいずれかになります。反対の三角形（厳密な上三角形または厳密な下三角形）の出力値は、実装で定義されます。

入力行列がエルミート正定値行列ではない i が存在する場合、動作は未定義です。

量子化された型の場合は、dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)) を実行します。

入力

出力

制約

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

例

// %a: [
//      [1.0, 2.0, 3.0],
//      [2.0, 20.0, 26.0],
//      [3.0, 26.0, 70.0]
//     ]
%result = "stablehlo.cholesky"(%a) {
  lower = true
} : (tensor<3x3xf32>) -> tensor<3x3xf64>
// %result: [
//           [1.0, 0.0, 0.0],
//           [2.0, 4.0, 0.0],
//           [3.0, 5.0, 6.0]
//          ]

クランプ

セマンティクス

operand テンソルのすべての要素を最小値と最大値の間にクランプし、result テンソルを生成します。より正式には、result[result_index] = minimum(maximum(operand[result_index], min_element), max_element)（min_element = rank(min) = 0 ? min[] : min[result_index]、max_element = rank(max) = 0 ? max[] : max[result_index]）です。量子化された型の場合は、dequantize_op_quantize(clamp, min, operand, max, type(result)) を実行します。

複素数の順序付けには意外なセマンティクスが伴うため、将来的には、この演算での複素数のサポートは終了する予定です（#560）。

入力

出力

制約

• （C1）rank(min) = 0 or shape(min) = shape(operand)。
• （C2）rank(max) = 0 or shape(max) = shape(operand)。
• （C3）baseline_element_type(min) = baseline_element_type(operand) = baseline_element_type(max)。
• （C4）baseline_type(operand) = baseline_type(result)。

例

// %min: [5, 10, 15]
// %operand: [3, 13, 23]
// %max: [10, 15, 20]
%result = "stablehlo.clamp"(%min, %operand, %max) : (tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<3xi32>
// %result: [5, 13, 20]

その他の例

collective_broadcast

セマンティクス

StableHLO プロセス グリッドの各プロセス グループ内で、operand テンソルの値をソースプロセスからターゲット プロセスに送信し、result テンソルを生成します。

このオペレーションでは、StableHLO プロセス グリッドが、次のように定義された process_groups に分割されます。

• channel_id <= 0 の場合は cross_replica(replica_groups)。
• channel_id > 0 の場合、cross_partition(replica_groups)。

その後、result@process は次のようになります。

• operand@process_groups[i, 0]: プロセスが process_groups[i] にあるような 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 は次のように定義されます。
  • cross_replica を使用する場合は num_replicas。
  • cross_partition を使用する場合は num_partitions。
• （C3）type(result) = type(operand)。

例

// num_replicas: 4
// num_partitions: 1
// %operand@(0, 0): [[1, 2]]
// %operand@(1, 0): [[3, 4]]
// %operand@(2, 0): [[5, 6]]
// %operand@(3, 0): [[7, 8]]
%result = "stablehlo.collective_broadcast"(%operand) {
  replica_groups = dense<[[2, 1]]> : tensor<1x2xi64>,
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor1x2xi64>) -> tensor<1x2xi64>
// %result@(0, 0): [[0, 0]]
// %result@(1, 0): [[5, 6]]
// %result@(2, 0): [[5, 6]]
// %result@(3, 0): [[0, 0]]

collective_permute

セマンティクス

StableHLO プロセス グリッド内の各プロセス グループ内で、operand テンソルの値をソースプロセスからターゲット プロセスに送信し、result テンソルを生成します。

このオペレーションは、StableHLO プロセス グリッドを次のように定義された process_groups に分割します。

• channel_id <= 0 の場合、cross_replica(source_target_pairs)。
• channel_id > 0 の場合は cross_partition(source_target_pairs)。

その後、result@process は次のようになります。

• operand@process_groups[i, 0]。process_groups[i, 1] = process のような i が存在する場合。
• それ以外の場合は 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 は次のように定義されます。
  • cross_replica を使用する場合は num_replicas。
  • cross_partition を使用する場合は num_partitions。
• （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]]

その他の例

比較

セマンティクス

comparison_direction と compare_type に従って lhs テンソルと rhs テンソルの要素ごとの比較を行い、result テンソルを生成します。

comparison_direction と compare_type の値のセマンティクスは次のとおりです。

ブール値と整数の要素型の場合:

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

compare_type = FLOAT の浮動小数点要素型の場合、op は次の IEEE-754 演算を実装します。

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

compare_type = TOTALORDER を含む浮動小数点要素型の場合、op は IEEE-754 の totalOrder オペレーションと compareQuietEqual オペレーションの組み合わせを使用します。

複雑な要素型の場合、指定された comparison_direction と compare_type を使用して、(real, imag) ペアの辞書順による比較が行われます。複素数に順序付けを適用すると、予期しないセマンティクスが発生するため、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 は次のように定義されます。
  • is_signed_integer(element_type(lhs)) の場合は SIGNED。
  • is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)) の場合は UNSIGNED。
  • is_float(element_type(lhs)) の場合は FLOAT または TOTALORDER。
  • is_complex(element_type(lhs)) の場合は FLOAT。

例

// %lhs: [1.0, 3.0]
// %rhs: [1.1, 2.9]
%result = "stablehlo.compare"(%lhs, %rhs) {
  comparison_direction = #stablehlo<comparison_direction LT>,
  compare_type = #stablehlo<comparison_type FLOAT>
} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xi1>
// %result: [true, false]

その他の例

複雑

セマンティクス

実数値と虚数値のペア（lhs と rhs）から複合値に要素単位で変換し、result テンソルを生成します。

入力

出力

制約

• （C1）type(lhs) = type(rhs)。
• （C2）shape(result) = shape(lhs)。
• （C3）element_type(result) の型は complex<E> で、E = element_type(lhs) です。

例

// %lhs: [1.0, 3.0]
// %rhs: [2.0, 4.0]
%result = "stablehlo.complex"(%lhs, %rhs) : (tensor<2xf64>, tensor<2xf64>) -> tensor<2xcomplex<f64>>
// %result: [(1.0, 2.0), (3.0, 4.0)]

その他の例

複合

セマンティクス

他の StableHLO オペレーションで構成された（コンポーズされた）オペレーションをカプセル化し、inputs と composite_attributes を受け取って results を生成します。op のセマンティクスは decomposition 属性によって実装されます。composite オペレーションは、プログラムのセマンティクスを変更せずに分解に置き換えることができます。分解をインライン化しても同じオペレーションのセマンティクスが得られない場合は、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>

その他の例

concatenate

セマンティクス

指定された引数と同じ順序で dimension 次元に沿って inputs を連結し、result テンソルを生成します。より正式には、result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1] です。ここで、

1. id = d0 + ... + dk-1 + kd。
2. d は dimension に等しく、d0 は inputs の d 番目のディメンション サイズです。

入力

出力

制約

• （C1）same(element_type(inputs...))。
• （C2）dim(inputs..., dimension) を除く same(shape(inputs...))。
• （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]]

その他の例

定数

セマンティクス

定数 value から output テンソルを生成します。

入力

出力

制約

• （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 テンソルを生成します。

boolean-to-any-supported-typeでは、値 false はゼロに変換され、値 true は 1 に変換されます。any-supported-type-to-booleanへの変換の場合、ゼロ値は false に変換され、ゼロ以外の値は true に変換されます。複雑な型での動作については以下をご覧ください。

integer-to-integer、integer-to-floating-point、floating-point-to-floating-point を含む変換では、ソース値が宛先の型で正確に表現できる場合、結果値はそのまま表現されます。それ以外の場合、動作は未定です（#180）。

浮動小数点数から整数への変換では、小数部分は切り捨てられます。切り捨てられた値を宛先型で表現できない場合、動作は未定です（#180）。

複素数から複素数への変換は、実部と虚部の変換における浮動小数点数から浮動小数点数への変換と同じ動作になります。

複素数から他の型への変換と他の型から複素数への変換では、それぞれソースの虚数値が無視されるか、宛先の虚数値がゼロになります。実部の変換は浮動小数点変換に従います。

原則として、このオペレーションはデクォンタイズ（量子化テンソルから通常のテンソルへの変換）、量子化（通常のテンソルから量子化テンソルへの変換）、再量子化（量子化テンソル間の変換）を表現できますが、現時点では専用のオペレーションがあります。最初のユースケースの場合は uniform_dequantize、2 つ目のユースケースと 3 つ目のユースケースの場合は uniform_quantize です。今後、これらの 2 つのオペレーションは 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 の場合、index_space(dim(result, output_spatial_dimensions...)) 内のすべての output_spatial_index について、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) は次のように定義されます。
  • result_dim = output_batch_dimension の場合、dim(lhs, input_batch_dimension) / batch_group_count。
  • result_dim = output_feature_dimension の場合、dim(rhs, kernel_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 テンソルを生成します。要素のタイプに応じて、次の操作を行います。

• 浮動小数点数の場合: IEEE-754 の cos。
• 複素数の場合: 複素余弦。
• 量子化型の場合: 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

セマンティクス

inputs と called_computations を受け取って results を生成する、実装定義のオペレーション call_target_name をカプセル化します。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 テンソルを生成します。要素のタイプに応じて、次の処理を行います。

• 整数の場合: 小数部を破棄して代数的な商を生成する整数除算。
• 浮動小数点数の場合: IEEE-754 の division。
• 複素数の場合: 複素除算。
• 量子化された型の場合:
  • 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: オペレーションの LHS と RHS が丸められる精度。精度タイプは、入力と出力のストレージ タイプとは独立しています。
• accumulation_type 累積に使用される精度。
• lhs_component_count、rhs_component_count、num_primitive_operations は、LHS や RHS を複数のコンポーネントに分解し、それらの値に対して複数の「プリミティブ」ドット演算を実行するアルゴリズムを実行する場合に適用されます。通常は、より高い精度をエミュレートします（例: 高精度の計算に bfloat16 AI データ型を活用: btf12_6x3x3）。分解のないアルゴリズムの場合、これらの値は 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）axes(operand) 内のすべての d について:
  • 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]
//            ]
//          ]

その他の例

dynamic_conv

セマンティクス

このオペレーションは、畳み込みオペレーションと機能的に同じですが、パディングは 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) は次のように定義されます。
  • result_dim = output_batch_dimension の場合、dim(lhs, input_batch_dimension) / batch_group_count。
  • result_dim = output_feature_dimension の場合、dim(rhs, kernel_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]]
//          ]]