StableHLO の仕様

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

Google の目標は、より多くのプロダクトを開発することで、ML 開発を簡素化、加速させることです。 さまざまな ML フレームワーク(TensorFlow、JAX、 PyTorch など)や ML コンパイラ(XLA や IREE など)を使用したデータ処理パイプラインの構築をサポートします。そのために、 に、StableHLO プログラミング言語の仕様が記載されています。

この仕様には、大きく分けて 3 つのセクションがあります。まず、 [プログラム] セクションでは、StableHLO プログラムの構造について説明します。 これは StableHLO オペレーションで構成され、それ自体が StableHLO オペレーションで構成されます。 その構造内の Ops セクションでは、 見ていきますExecution セクションでは、すべての 1 つのプログラム内で 一緒に実行されるからです最後に、 表記セクションでは、本コース全体で使用される表記法を 仕様。

StableHLO の以前のリリースの仕様を表示するには、 タグ付けされたリリースです。 たとえば、StableHLO v0.19.0 仕様などです。 StableHLO のマイナー バージョンのバンプごとに発生した変更を表示するには、以下をご覧ください。 VhloDialect.td にあるバージョン ログ。

プログラム

Program ::= {Func}

StableHLO プログラムは、任意の数の StableHLO 関数で構成されます。 以下は、3 つの入力を持つ関数 @main を使用したプログラムの例を示しています。 (%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 関数名前付き関数とも呼ばれます)には、 識別子、入出力、本文があります将来的には 互換性を高めるために関数のメタデータを追加導入する (#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} | '?'

テンソル型はテンソル、つまり多次元配列を表します。この会社には shapeelement type: シェイプは負でないか、 未知のディメンション サイズ(対応する要素の昇順) ディメンション(軸とも呼ばれます)は 0 から R-1 まで番号が付けられます。「 次元数 R はランクと呼ばれます。たとえば、tensor<2x3xf32> は次のようになります。 形状が 2x3、要素型が f32 のテンソル型。2 つの次元があります。 (つまり、2 つの軸)- 0 次元目と 1 次元目 - そのサイズは 2 と 3 です。順位は 2 です。

シェイプは、部分的または完全に未知(動的)にできます。たとえば、tensor<?x2xf64> 部分的に不明で、tensor<?x?xf64> は完全に不明です。動的 ディメンション サイズは ? を使用して表されます。シェイプのランク付けを解除することはできません。

将来的には、さまざまなテンソル型の拡張を 要素のタイプ(たとえば、レイアウトを含む) (#629)とスパース性 (#1078)。

QuantizedTensorType ::= 'tensor' '<' Shape QuantizedTensorElementType '>'
QuantizedTensorElementType ::= '!quant.uniform' '<'
                  QuantizationStorageType
                  ['<' QuantizationStorageMin ':' QuantizationStorageMax '>']
                  ':' QuantizationExpressedType
                  [':' QuantizationDimension]
                  ',' QuantizationParameters '>'
QuantizationStorageType ::= IntegerType
QuantizationStorageMin ::= IntegerConstant
QuantizationStorageMax ::= IntegerConstant
QuantizationExpressedType ::= FloatType
QuantizationDimension ::= IntegerConstant
QuantizationParameters ::= QuantizationParameter
                         | '{' QuantizationParameter {',' QuantizationParameter} '}'
QuantizationParameter ::= QuantizationScale ':' QuantizationZeroPoint
QuantizationScale ::= FloatConstant
QuantizationZeroPoint ::= IntegerConstant
名前 タイプ 制約
storage_type 整数型 (C1 ~ C3)、(C8)
storage_min 整数定数 (C1)、(C3)、(C7)
storage_max 整数定数 (C2)、(C3)、(C7)
expressed_type 浮動小数点型 (C4)
quantization_dimension オプションの整数定数 (C10-C12)
scales 浮動小数点定数の可変数 (C4 ~ C6)、(C9)、(C10)、(C13)
zero_points 整数定数の可変数 (C7 ~ C9)

量子化要素型は、ストレージ タイプの整数値を表します。 対応する storage_minstorage_max の範囲(両端を含む) 表現型の浮動小数点値。指定された整数値 i について、次の処理が行われます。 対応する浮動小数点値 f は次のように計算できます。 f = (i - zero_point) * scalescalezero_point が呼び出される) 量子化パラメータstorage_minstorage_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 つまたは 1 つの指標が ゼロ点が複数含まれることがあります。基づいて、 0 ポイント前後の仕様は変わる可能性があります。 (#1405)。

QuantizationStorageMin のセマンティクスに関して進行中の別のディスカッションがあります と QuantizationStorageMax を使用して、制約を適用するかどうかを決定します。 これらの値と量子化テンソルの値に適用される (#1406)。

最後に、未知のスケールとゼロの表現について 未知の表現をモデル化する方法について (#1407

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

量子化テンソルでは、量子化はテンソルごとに行うことができます。 テンソル全体で 1 つの scalezero_point、または軸ごとにすることもできる つまり、複数の scaleszero_points があり、スライスごとに 1 つのペアが quantization_dimension を返します。より正式には、テンソル t で表します。 軸ごとの量子化の場合、dim(t, quantization_dimension) 個のスライスがあります。 /quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :]i 番目のスライス内のすべての要素は、次のように scales[i]zero_points[i] を使用します。 生成します。量子化テンソル型は次のとおりです。 制約:

  • テンソルごとの量子化の場合: <ph type="x-smartling-placeholder">
      </ph>
    • 追加の制約はありません。
  • 軸ごとの量子化の場合: <ph type="x-smartling-placeholder">
      </ph>
    • (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 を包括的に表現するTtuple<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 ::= 'f8E4M3FN' | 'f8E5M2' | 'f8E4M3FNUZ' | 'f8E5M2FNUZ'
            | 'f8E4M3B11FNUZ' | 'bf16' | 'f16' | 'f32' | 'f64'
TensorFloat32 ::= 'tf32'
ComplexType ::= 'complex' '<' ComplexElementType '>'
ComplexElementType ::= 'f32' | 'f64'

要素型は、テンソル型の要素を表します。多くのプログラミングでは、 これらの型は StableHLO のファースト クラスではありません。つまり StableHLO プログラムは、これらの型の値を直接表すことはできません(そのため、 T 型のスカラー値を 0 次元のテンソルで表すのが慣用的です。 tensor<T> 型の値など)。

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

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

StringType ::= 'string'

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

運用

StableHLO オペレーションオペレーションとも呼ばれます)はクローズド セットを表す おおまかに説明しましたすでに説明したように StableHLO の構文は MLIR に大きく影響しますが、MLIR は 人間工学に基づいた代替手段ですが、StableHLO の目標である ML フレームワークと ML コンパイラ間の 相互運用性を向上できます

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

Ops は入力を消費し、出力を生成します。入力は以下のように分類されます。 入力値(実行中に計算される)、入力関数( StableHLO では関数はファースト クラス値ではないため、静的に保持されます。 入力属性(静的にも提供されます)入力と出力の種類 演算で消費、生成される情報はそのニモニックに左右されます。たとえば、add op は 2 つの入力値を消費し、1 つの出力値を生成します。これに対し、 select_and_scatter 演算は 3 つの入力値、2 つの入力関数、 3 つの入力属性。

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

入力関数匿名関数とも呼ばれます)は、 名前付き関数に似ていますが、次の点が異なります。1)識別子がない(つまり 「anonymous」という名前)、2)出力タイプを宣言しない(出力タイプは 関数内の return オペレーションから推論されます)。

入力関数の構文には、現在使用されていない部分が含まれます( 上記の Unused 本番環境など)は、MLIR との互換性を確保するために存在します。MLIR では 「リージョン」というより一般的なデータには複数の「ブロック」や ジャンプ オペレーションで接続されているオペレーションです。これらのブロックには、それぞれ対応する ID と Unused 本番環境に関連付けて、互いに区別できるようにします。 StableHLO にはジャンプ オペレーションがないため、MLIR 構文の対応する部分は次のとおりです。 未使用です(ただし、まだ存在します)。

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

入力属性には、名前と値があります。これは、サポートされている 使用します。これらは、プログラムの静的メタデータを指定する主要な方法です。 あります。たとえば、concatenate 演算では属性 dimension を使用して、 入力値を連結するディメンションを指定します。同様に slice 演算は start_indiceslimit_indices などの複数の属性を使用する を使用して、入力値をスライスする境界を指定します。

現時点では、StableHLO プログラムには属性が含まれていることがあり、 ファイアウォールルールがあります将来的には これらの属性を StableHLO opset に吸収するか、 StableHLO プログラムに 表示されますそれまでは 属性:

  • layout#629)。
  • mhlo.frontend_attributes#628)。
  • mhlo.sharding#619)。
  • output_operand_aliases#740)。
  • 位置情報のメタデータ(#594)。
OpSignature ::= '(' [ValueType {',' ValueType}] ')' '->' '(' [ValueType {',' ValueType}] ')'

演算シグネチャは、すべての入力値の型( -> の左側)とすべての出力値の型(すべての出力値の -> の右側にあります)。厳密に言えば、入力タイプは 出力型もほぼ常に冗長です( 出力の型は入力から推測できます)。それでも シグネチャは、MLIR との互換性を確保するために、StableHLO 構文に意図的に含まれています。

ニモニックが select_and_scatter の演算の例を次に示します。消費電力は 入力値(%operand%source%init_value)、2 つの入力関数 と 3 つの入力属性(window_dimensionswindow_stridespadding)。 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'

ブール定数は、ブール値 truefalse を表します。ブール値 定数の型は 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 進数表記で表します。その他のベース(例:2 進数または 8 進数はサポートされません。 整数定数には次の制約があります。

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

浮動小数点定数は、 小数または科学的表記を使用できますさらに、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

複素定数は実数部のリストを使用して複素数を表す (最初に来る)と虚数部(2 番目)があります。たとえば (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))。各要素の意味は次のとおりです。 <ph type="x-smartling-placeholder">
      </ph>
    • 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))。各要素の意味は次のとおりです。 <ph type="x-smartling-placeholder">
      </ph>
    • 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 テンソルに対して要素単位の abs 演算を実行し、result を生成します。 テンソルです。要素のタイプに応じて、次の操作を行います。

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

入力

ラベル 名前 タイプ 制約
(I1) operand 符号付き整数、浮動小数点数、複合型のテンソル、またはテンソルごとの量子化テンソル (C1 ~ C2)

出力

名前 タイプ 制約
result 符号付き整数型または浮動小数点型のテンソル、またはテンソルごとの量子化テンソル (C1 ~ C2)

制約

  • (C1)shape(result) = shape(operand)
  • (C2)baseline_element_type(result) は次のように定義されます。 <ph type="x-smartling-placeholder">
      </ph>
    • 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 つのテンソル lhsrhs を要素ごとに加算し、 result のテンソル。要素のタイプに応じて、次の操作を行います。

  • ブール値の場合: 論理 OR。
  • 整数の場合: 整数の加算。
  • 浮動小数点数: IEEE-754 の addition
  • 複素数の場合: 複素加算。
  • 量子化型の場合: dequantize_op_quantize(add, lhs, rhs, type(result))

入力

ラベル 名前 タイプ 制約
(I1) lhs テンソルまたは量子化テンソル (C1 ~ C6)
(I2) rhs テンソルまたは量子化テンソル (C1 ~ C5)、(C7)

出力

名前 タイプ 制約
result テンソルまたは量子化テンソル (C1 ~ C7)

制約

  • オペレーションで量子化されていないテンソルを使用する場合: <ph type="x-smartling-placeholder">
      </ph>
    • (C1)type(lhs) = type(rhs) = type(result)
  • オペレーションで量子化テンソルを使用する場合: <ph type="x-smartling-placeholder">
      </ph>
    • (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 へのデータの依存関係を確立するためにのみ存在します。

入力

ラベル 名前 タイプ
(I1) inputs token の可変数

出力

名前 タイプ
result token

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

その他の例

all_gather

セマンティクス

StableHLO プロセス グリッドの各プロセス グループ内で、値を連結する all_gather_dim の各プロセスから 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 内で以下を行います。

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

入力

ラベル 名前 タイプ 制約
(I1) operands テンソルの可変数またはテンソルごとの量子化テンソル (C1)、(C6)
(I2) all_gather_dim si64 型の定数 (C1)、(C6)
(I3) replica_groups si64 型の 2 次元のテンソル定数 (C2-C4)
(I4) channel_id si64 型の定数 (C5)
(I5) use_global_device_ids i1 型の定数 (C5)

出力

名前 タイプ 制約
results テンソルの可変数またはテンソルごとの量子化テンソル (C6)

制約

  • (C1)0 <= all_gather_dim < rank(operands...)
  • (C2)is_unique(replica_groups)
  • (C3)size(replica_groups) は次のように定義されます。 <ph type="x-smartling-placeholder">
      </ph>
    • 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...)、ただし: <ph type="x-smartling-placeholder">
      </ph>
    • 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 各要素の意味は次のとおりです。 <ph type="x-smartling-placeholder">
      </ph>
    • 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])) です。

入力

ラベル 名前 タイプ 制約
(I1) operands テンソルの可変数またはテンソルごとの量子化テンソル (C5)、(C6)
(I2) replica_groups si64 型の 1 次元テンソル定数の可変数 (C1 ~ C3)
(I3) channel_id si64 型の定数 (C4)
(I4) use_global_device_ids i1 型の定数 (C4)
(I5) computation 関数 (C5)

出力

名前 タイプ 制約
results テンソルの可変数またはテンソルごとの量子化テンソル (C6-C7)

制約

  • (C1)is_unique(replica_groups)
  • (C2)size(replica_groups) は次のように定義されます。 <ph type="x-smartling-placeholder">
      </ph>
    • 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

セマンティクス

all_to_all

StableHLO プロセス グリッドの各プロセス グループ内で、 operands テンソルを split_dimension に沿って分割し、分割を分散させる 分散された部分を連結して、プロセス間の concat_dimensionresults のテンソルを生成します。 このオペレーションにより、StableHLO プロセス グリッドが process_groups に分割されます。 次のように定義されます。

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

その後、各 process_group 内で以下を行います。

  • split_parts...@sender = split(operands...@sender, split_count, split_dimension) process_group のすべての sender に対して適用されました。
  • 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)

入力

ラベル 名前 タイプ 制約
(I1) operands テンソルの可変数またはテンソルごとの量子化テンソル (C1 ~ C3)、(C9)
(I2) split_dimension si64 型の定数 (C1)、(C2)、(C9)
(I3) concat_dimension si64 型の定数 (C3)、(C9)
(I4) split_count si64 型の定数 (C2)、(C4)、(C8)、(C9)
(I5) replica_groups si64 型の 2 次元のテンソル定数 (C5 ~ C8)
(I6) channel_id si64 型の定数

出力

名前 タイプ 制約
results テンソルの可変数またはテンソルごとの量子化テンソル (C9)

制約

  • (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) は次のように定義されます。 <ph type="x-smartling-placeholder">
      </ph>
    • 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 の場合を除く): <ph type="x-smartling-placeholder">
      </ph>
    • 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 つのテンソル lhsrhs の要素単位の AND を実行し、result を生成します。 テンソルです。要素のタイプに応じて、次の操作を行います。

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

入力

ラベル 名前 タイプ 制約
(I1) lhs ブール型または整数型のテンソル (C1)
(I2) rhs ブール型または整数型のテンソル (C1)

出力

名前 タイプ 制約
result ブール型または整数型のテンソル (C1)

制約

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

入力

ラベル 名前 タイプ 制約
(I1) lhs 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)
(I2) rhs 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

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

セマンティクス

batch_norm_training 誤差逆伝播の複数の入力の勾配を計算します。 grad_output から生成され、grad_operandgrad_scalegrad_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))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C1 ~ C3)、(C5)
(I2) scale 浮動小数点数またはテンソルごとの量子化型の 1 次元テンソル (C2)、(C4)、(C5)
(I3) mean 浮動小数点数またはテンソルごとの量子化型の 1 次元テンソル (C2)、(C4)
(I4) variance 浮動小数点数またはテンソルごとの量子化型の 1 次元テンソル (C2)、(C4)
(I5) grad_output 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C2)、(C3)
(I6) epsilon f32 型の定数
(I7) feature_index si64 型の定数 (C1)、(C5)

出力

名前 タイプ 制約
grad_operand 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C2)、(C3)
grad_scale 浮動小数点数またはテンソルごとの量子化型の 1 次元テンソル (C2)、(C4)
grad_offset 浮動小数点数またはテンソルごとの量子化型の 1 次元テンソル (C2)、(C4)

制約

  • (C1)0 <= feature_index < rank(operand)
  • (C2)operandscalemeanvariancegrad_outputgrad_operandgrad_scalegrad_offsetbaseline_element_type は同じです。
  • (C3)operandgrad_outputgrad_operand の形状が同じである。
  • (C4)scalemeanvariancegrad_scalegrad_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

セマンティクス

次を除くすべての次元で 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))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C1 ~ C7)
(I2) scale 浮動小数点数またはテンソルごとの量子化型の 1 次元テンソル (C2)、(C3)
(I3) offset 浮動小数点数またはテンソルごとの量子化型の 1 次元テンソル (C2)、(C4)
(I4) mean 浮動小数点数またはテンソルごとの量子化型の 1 次元テンソル (C5)
(I5) variance 浮動小数点数またはテンソルごとの量子化型の 1 次元テンソル (C2)、(C6)
(I6) epsilon f32 型の定数
(I7) feature_index si64 型の定数 (C1)、(C3 ~ C6)

出力

名前 タイプ 制約
result 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C2)、(C7)

制約

  • (C1)0 <= feature_index < rank(operand)
  • (C2)operandscaleoffsetmeanvarianceresult が 同じ 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 テンソルを正規化し、outputbatch_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))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C1)
(I2) scale 浮動小数点数またはテンソルごとに量子化された 1 次元テンソル (C2)、(C3)
(I3) offset 浮動小数点数またはテンソルごとに量子化された 1 次元テンソル (C2)、(C4)
(I4) epsilon f32 型の定数 (C1)、(C3 ~ C6)
(I5) feature_index si64 型の定数 (C1)、(C3 ~ C6)

出力

名前 タイプ 制約
output 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C7)
batch_mean 浮動小数点数またはテンソルごとに量子化された 1 次元テンソル (C2)、(C5)
batch_var 浮動小数点数またはテンソルごとに量子化された 1 次元テンソル (C2)、(C6)

制約

  • (C1)0 <= feature_index < rank(operand)
  • (C2)operandscaleoffsetbatch_meanbatch_varoutput が 同じ 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 は、指定された値のメモリ内表現とその動作を返します。 実装で定義されているのは、テンソルの正確な表現が 実装定義です。要素の型の正確な表現は、 使用されます。

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたは量子化テンソル (C1 ~ C2)

出力

名前 タイプ 制約
result テンソルまたは量子化テンソル (C1 ~ C2)

制約

  • (C1)E = is_quantized(operand) ? storage_type(operand) : element_type(operand)E' = is_quantized(result) ? storage_type(result) : element_type(result)R = rank(operand) がある場合: <ph type="x-smartling-placeholder">
      </ph>
    • num_bits(E') = num_bits(E) の場合は、shape(result) = shape(operand)
    • num_bits(E') < num_bits(E) の場合:
    • rank(result) = R + 1
    • すべての 0 <= i < R に対して dim(result, i) = dim(operand, i)
    • dim(result, R) * num_bits(E') = num_bits(E)
    • num_bits(E') > num_bits(E) の場合:
    • rank(result) = R - 1
    • すべての 0 <= i < R に対して dim(result, i) = dim(operand, i)
    • 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):

  • dim(operand, d) = 1 の場合は operand_index[d] = 0
  • そうでない場合は operand_index[d] = result_index[broadcast_dimensions[d]]

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたは量子化テンソル (C1-C2)、(C5-C6)
(I2) broadcast_dimensions si64 型の 1 次元のテンソル定数 (C2-C6)

出力

名前 タイプ 制約
result テンソルまたは量子化テンソル (C1)、(C3)、(C5 ~ C6)

制約

  • (C1)element_type(result) は次式で求められます。 <ph type="x-smartling-placeholder">
      </ph>
    • 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 について: <ph type="x-smartling-placeholder">
      </ph>
    • dim(operand, d) = 1 または
    • dim(operand, d) = dim(result, broadcast_dimensions[d])
  • (C6)is_per_axis_quantized(result) の場合: <ph type="x-smartling-placeholder">
      </ph>
    • 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 から関数を 1 つだけ実行すると、出力が生成されます index の値に応じて異なります。よりフォーマルな表現で、result = selected_branch() ここで

  • 0 <= index < size(branches) の場合は selected_branch = branches[index]
  • そうでない場合は selected_branch = branches[-1]

入力

ラベル 名前 タイプ 制約
(I1) index si32 型の 0 次元テンソル
(I2) branches 関数の可変数 (C1 ~ C4)

出力

名前 タイプ 制約
results テンソル、量子化テンソル、またはトークンの可変数 (C4)

制約

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

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

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

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C1)

制約

  • (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, :, :](下向きの三角形または (lowertrue の場合)または上三角形(lowerfalse の場合)の行列。 反対側の三角形(厳密な上三角形)の出力値 それに対応して厳密な下三角形が定義され、実装定義になります。

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

量子化型の場合 dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result))

入力

ラベル 名前 タイプ 制約
(I1) a 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1 ~ C3)
(I2) lower i1 型の 0 次元のテンソル定数

出力

名前 タイプ 制約
result 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

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

入力

ラベル 名前 タイプ 制約
(I1) min テンソルまたはテンソルごとの量子化テンソル (C1)、(C3)
(I2) operand テンソルまたはテンソルごとの量子化テンソル (C1 ~ C4)
(I3) max テンソルまたはテンソルごとの量子化テンソル (C2)、(C3)

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C4)

制約

  • (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](プロセスが次のようになっているような i が存在する場合) (process_groups[i]
  • broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) できません。

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたはテンソルごとの量子化テンソル (C3)
(I2) replica_groups si64 型の 1 次元テンソル定数の可変数 (C1)、(C2)
(I3) channel_id si64 型の定数

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C3)

制約

  • (C1)is_unique(replica_groups)
  • (C2)0 <= replica_groups < N。ここで、N は次のように定義されます。 <ph type="x-smartling-placeholder">
      </ph>
    • 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](該当する i が存在する場合) process_groups[i, 1] = process
  • broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) できません。

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたはテンソルごとの量子化テンソル (C5)
(I2) source_target_pairs si64 型の 2 次元のテンソル定数 (C1 ~ C4)
(I3) channel_id si64 型の定数

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C1)

制約

  • (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 は次のように定義されます。 <ph type="x-smartling-placeholder">
      </ph>
    • 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]]

その他の例

比較

セマンティクス

次に従って lhs テンソルと rhs テンソルの要素単位の比較を行います。 comparison_directioncompare_type があり、result テンソルを生成します。

comparison_directioncompare_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_directioncompare_type を使用して実行される) 複素数に順序を付けるには、意外な意味論が必要です。 将来的には複素数のサポートは終了し comparison_directionGEGTLE、または LT の場合 (#560)。

量子化型の場合。dequantize_compare(lhs, rhs, comparison_direction) を実行します。

入力

ラベル 名前 タイプ 制約
(I1) lhs テンソルまたはテンソルごとの量子化テンソル (C1 ~ C3)
(I2) rhs テンソルまたはテンソルごとの量子化テンソル (C1 ~ C2)
(I3) comparison_direction EQNEGEGTLELT の列挙型
(I4) compare_type FLOATTOTALORDERSIGNEDUNSIGNED の列挙型 (C3)

出力

名前 タイプ 制約
result ブール型のテンソル (C2)

制約

  • (C1)baseline_element_type(lhs) = baseline_element_type(rhs)
  • (C2)shape(lhs) = shape(rhs) = shape(result)
  • (C3)compare_type は次のように定義されます。 <ph type="x-smartling-placeholder">
      </ph>
    • 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]

その他の例

複雑

セマンティクス

実数列と実数行列のペアから複合値に要素単位で変換 架空の値(lhsrhs)から成り、result テンソルを生成します。

入力

ラベル 名前 タイプ 制約
(I1) lhs f32 型または f64 型のテンソル (C1 ~ C3)
(I2) rhs f32 型または f64 型のテンソル (C1)

出力

名前 タイプ 制約
result 複素型のテンソル (C2)、(C3)

制約

  • (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 オペレーションで構成された(構成)オペレーションをカプセル化します。 inputscomposite_attributes を取り、results を生成します。「 op のセマンティクスは、decomposition 属性によって実装されます。「 composite 演算は、プログラムを変更せずに分解に置き換えることができる 学びました。分解をインライン化しても同じ結果が得られない場合、 custom_call を使用することをおすすめします。

version フィールド(デフォルトは 0)は、複合レイヤがいつ 示されます。

入力

ラベル 名前 タイプ
(I1) inputs 値の可変数
(I2) name string 型の定数
(I3) composite_attributes 属性辞書
(I4) decomposition string 型の定数
(I5) version si32 型の定数

出力

名前 タイプ
results 値の可変数

制約

  • (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. ddimension と等しく、d0d 番目のディメンション サイズです /inputs

入力

ラベル 名前 タイプ 制約
(I1) inputs テンソルの可変数またはテンソルごとの量子化テンソル (C1 ~ C6)
(I2) dimension si64 型の定数 (C2)、(C4)、(C6)

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C5-C6)

制約

  • (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])、ただし以下を除く: <ph type="x-smartling-placeholder">
      </ph>
    • 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 テンソルを生成します。

入力

ラベル 名前 タイプ 制約
(I1) value 定数 (C1)

出力

名前 タイプ 制約
output テンソルまたは量子化テンソル (C1)

制約

  • (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 は次のようになります。 0 に変換され、値 true は 1 に変換されます。対象 any-supported-type-to-booleanが行われる場合、ゼロ値は false とゼロ以外の値は true に変換されます。この仕組みについては、下記をご覧ください。 複雑な型で機能します。

整数から整数または整数から浮動小数点数を含む変換の場合 または floating-point-to-floating-point(ソース値が が宛先タイプで表されている場合、結果の値は 必要があります。それ以外の場合の動作は未定 (#180)。

floating-point-to-integerへの変換の場合、小数部は次のようになります。 切り捨てられます。切り捨てられた値を宛先タイプで表現できない場合は、 動作は未定です(#180)。

複合型を含む変換は、 floating-point-to-floating-pointの変換(実数と浮動小数点数) 使用します。

complex-to-any-other-typeany-other-type-to-complex の変換の場合: ソースの虚偽値が無視されるか、宛先の虚数が それぞれゼロに設定されます。実数部分の変換は、 浮動小数点変換も実行できます

原理的には、この演算は逆量子化( 量子化テンソルから規則テンソルへの変換)、量子化(通常のテンソルから テンソルから量子化テンソルへの変換)と再量子化(量子化テンソルと テンソルがありますが、現時点では専用のオペレーションがあります。 最初のユースケースは uniform_dequantize、ユースケースは uniform_quantize です。 2 つ目と 3 つ目のユースケースです将来的には、この 2 つのオペレーションが統合される可能性があります。 convert に配置します(#1576)。

入力

ラベル 名前 タイプ 制約
(I1) operand テンソル (C1)

出力

名前 タイプ 制約
result テンソル (C1)

制約

  • (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 の要素がどのように計算されるかを示しています。 lhsrhs を具体例で説明します。

畳み込み

より正式には、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_indexindex_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) を実行します。

入力

ラベル 名前 タイプ 制約
(I1) lhs テンソルまたはテンソルごとの量子化テンソル (C1)、(C10 ~ C11)、(C14)(C25)、(C27 ~ C28)、(C31 ~ C32)、(C34)
(I2) rhs テンソルまたは量子化テンソル (C1)、(C14 ~ C16)、(C25)、(C27 ~ C29)、(C31 ~ C34)
(I3) window_strides si64 型の 1 次元のテンソル定数 (C2-C3)、(C25)
(I4) padding si64 型の 2 次元のテンソル定数 (C4)、(C25)
(I5) lhs_dilation si64 型の 1 次元のテンソル定数 (C5 ~ C6)、(C25)
(I6) rhs_dilation si64 型の 1 次元のテンソル定数 (C7 ~ C8)、(C25)
(I7) window_reversal i1 型の 1 次元のテンソル定数 (C9)
(I8) input_batch_dimension si64 型の定数 (C10)、(C13)、(C25)
(I9) input_feature_dimension si64 型の定数 (C11)、(C13 ~ C14)
(I10) input_spatial_dimensions si64 型の 1 次元のテンソル定数 (C12)、(C13)、(C25)
(I11) kernel_input_feature_dimension si64 型の定数 (C14)、(C18)
(I12) kernel_output_feature_dimension si64 型の定数 (C15 ~ C16)、(C18)、(C25)、(C29)
(I13) kernel_spatial_dimensions si64 型の 1 次元のテンソル定数 (C17 ~ C18)、(C25)
(I14) output_batch_dimension si64 型の定数 (C20)、(C25)
(I15) output_feature_dimension si64 型の定数 (C20)、(C25)、(C30)
(I16) output_spatial_dimensions si64 型の 1 次元のテンソル定数 (C19 ~ C20)、(C25)
(I17) feature_group_count si64 型の定数 (C11)、(C14)、(C16)、(C21)、(C23)
(I18) batch_group_count si64 型の定数 (C10)、(C15)、(C22)、(C23)、(C25)
(I19) precision_config DEFAULTHIGHHIGHEST の可変列挙型の数 (C24)

出力

名前 タイプ 制約
result テンソルまたは量子化テンソル (C25-C28)、(C30)、(C32-34)

制約

  • (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] の場合: <ph type="x-smartling-placeholder">
      </ph>
    • 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] の場合: <ph type="x-smartling-placeholder">
      </ph>
    • 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] の場合: <ph type="x-smartling-placeholder">
      </ph>
    • 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) は次のように定義されます。 <ph type="x-smartling-placeholder">
      </ph>
    • 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
  • オペレーションで量子化されていないテンソルを使用する場合: <ph type="x-smartling-placeholder">
      </ph>
    • (C27)element_type(lhs) = element_type(rhs) = element_type(result)
  • オペレーションで量子化テンソルを使用する場合: <ph type="x-smartling-placeholder">
      </ph>
    • (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))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

  • (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 内の先頭 0 ビットの数を要素ごとにカウントします result テンソルを生成します。

入力

ラベル 名前 タイプ 制約
(I1) operand 整数型のテンソル (C1)

出力

名前 タイプ 制約
result 整数型のテンソル (C1)

制約

  • (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 をカプセル化します。 inputscalled_computations を生成し、results を生成します。has_side_effect, backend_configapi_version は、 実装定義メタデータを使用します。

現在のところ、このオペレーションにはかなり整理されていないデータ コレクションが含まれています メタデータに、対応するオペレーションの有機的な進化を XLA コンパイラです将来的には、このメタデータの統合を予定しています。 (#741)。

入力

ラベル 名前 タイプ
(I1) inputs 値の可変数
(I2) call_target_name string 型の定数
(I3) has_side_effect i1 型の定数
(I4) backend_config string 型の定数または属性辞書
(I5) api_version si32 型の定数
(I6) called_computations string 型の定数の可変数

出力

名前 タイプ
results 値の可変数

%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
  • 複素数の場合: 複素除算。
  • 量子化型の場合: <ph type="x-smartling-placeholder">
      </ph>
    • dequantize_op_quantize(divide, lhs, rhs, type(result))

入力

ラベル 名前 タイプ 制約
(I1) lhs 整数型、浮動小数点数、複素数のテンソル、またはテンソルごとの量子化テンソル (C1)
(I2) rhs 整数型、浮動小数点数、複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 整数、浮動小数点、複合型のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

  • (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 は、次の要素の速度と精度のトレードオフを制御します。 コンピューティングの負荷を分散します以下のいずれかを指定できます( これらの列挙値のセマンティクスは指定されていませんが、 2024 年には #755):

  • DEFAULT: 計算は最も高速ですが、近似精度が最も低くなります。 あります。
  • HIGH: 計算は遅くなりますが、 あります。
  • HIGHEST: 計算が最も遅いものの、近似精度が最も高くなります。 あります。

DotAlgorithm は、実装に使用するアルゴリズムの主なプロパティを定義します。 精度も定義します。アルゴリズムの属性が フィールドが設定されている場合、precision_configDEFAULT にする必要があります。DotAlgorithms デフォルトのパラメータは実装であるため、デフォルト値はありません。 定義します。そのため、すべてのドット アルゴリズム フィールドを None に設定して、 空のドット アルゴリズムを使用します。このアルゴリズムでは代わりに precision_config 値が使用されます。

DotAlgorithm フィールドには次が含まれます。

  • lhs_precision_typerhs_precision_type は、LHS および 演算の RHS は四捨五入されます。精度の型は 入力と出力のストレージタイプを指定します
  • accumulation_type: 累積に使用する精度。
  • lhs_component_count さん、rhs_component_count さん、num_primitive_operations さん これは、LHS や RHS を 複数の「プリミティブ」を実行するドット演算 値 - 通常は高精度(例: bfloat16 AI データ型を活用して高精度の計算を行う: bf16_6x tf32_3x など)。分解を行わないアルゴリズムの場合、これらの値は 1 に設定する必要があります。
  • allow_imprecise_accumulation: 累積精度を低くするかどうかを指定します。 一部のステップで許可される(例: CUBLASLT_MATMUL_DESC_FAST_ACCUM)。

DotAlgorithm 属性の例:

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


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


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

どの組み合わせをサポートするかは、実装が決定します。イン 一般に、各アルゴリズムが各プラットフォームでサポートされていることは アクセラレータ タイプを指定します。特定のアルゴリズムが エラーがスローされるため、 あります。StableHLO の検証ではベスト エフォート型の検証、 どのハードウェアでもサポートされないことがわかっているアルゴリズムを防止できます。

xla_data.proto > Algorithm を参照してください。 をご覧ください。チケット #2483 は、サービス アカウントを作成する計画をキャプチャします。 バックエンドでサポートされるアルゴリズムに関する一元化されたドキュメントです。

入力

ラベル 名前 タイプ 制約
(I1) lhs テンソルまたはテンソルごとの量子化テンソル (C5 ~ C6)、(C9 ~ C10)、(C12 ~ C14)、(C17 ~ C18)、(C20)
(I2) rhs テンソルまたは量子化テンソル (C7-C10)、(C12-C20)
(I3) lhs_batching_dimensions si64 型の 1 次元のテンソル定数 (C1)、(C3)、(C5)、(C9)、(C12)
(I4) rhs_batching_dimensions si64 型の 1 次元のテンソル定数 (C1)、(C4)、(C7)、(C9)
(I5) lhs_contracting_dimensions si64 型の 1 次元のテンソル定数 (C2)、(C3)、(C6)、(C10)
(I6) rhs_contracting_dimensions si64 型の 1 次元のテンソル定数 (C2)、(C4)、(C8)、(C10)、(C16)
(I7) precision_config DEFAULTHIGHHIGHEST の可変列挙型の数 (C11)、(C21)
(I8) lhs_precision_type FloatType または TensorFloat32 (C21)
(I9) rhs_precision_type FloatType または TensorFloat32 (C21)
(I10) accumulation_type FloatType または TensorFloat32 (C21)
(I11) lhs_component_count si32 型の定数 (C21)、(C22)
(I12) rhs_component_count si32 型の定数 (C21)、(C23)
(I13) num_primitive_operations si32 型の定数 (C21)、(C24)
(I14) allow_imprecise_accumulation bool 型の定数 (C21)

出力

名前 タイプ 制約
result テンソルまたは量子化テンソル (C12)、(C14)、(C18 ~ C20)

制約

  • (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)
  • オペレーションで量子化されていないテンソルを使用する場合: <ph type="x-smartling-placeholder">
      </ph>
    • (C13)element_type(lhs) = element_type(rhs)
  • オペレーションで量子化テンソルを使用する場合: <ph type="x-smartling-placeholder">
      </ph>
    • (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) の場合: <ph type="x-smartling-placeholder">
      </ph>
    • (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_dimensionsknown_non_expanding_dimensions も使用できます。 ディメンションの展開動作に関する静的な知識を表すために使用します。 指定しない場合、すべてのディメンションが拡張される可能性があるものと見なされます。

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたは量子化テンソル (C1-C2)、(C5-C6)、(C9)
(I2) output_dimensions 整数型の 1 次元テンソル (C7)
(I3) broadcast_dimensions 整数型の 1 次元定数テンソル (C2-C6)
(I4) known_expanding_dimensions 整数型の 1 次元定数テンソル (C8-C9)
(I5) known_non_expanding_dimensions 整数型の 1 次元定数テンソル (C8-C9)

出力

名前 タイプ 制約
result テンソルまたは量子化テンソル (C1)、(C3)、(C5 ~ C7)

制約

  • (C1)element_type(result) は次式で求められます。 <ph type="x-smartling-placeholder">
      </ph>
    • 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 について: <ph type="x-smartling-placeholder">
      </ph>
    • dim(operand, d) = 1 または
    • dim(operand, d) = dim(result, broadcast_dimensions[d])
  • (C6)is_per_axis_quantized(result) の場合: <ph type="x-smartling-placeholder">
      </ph>
    • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)]
    • dim(operand, quantization_dimension(operand)) = 1 の場合: scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result)))
  • (C7)size(output_dimensions) = rank(result)
  • (C8)is_unique(known_expanding_dimensions + known_non_expanding_dimensions)
  • (C9)0 <= known_expanding_dimensions < rank(operand)
  • (C10)0 <= known_non_expanding_dimensions < rank(operand)

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

その他の例

dynamic_conv

セマンティクス

この処理は 畳み込み パディングは padding によって動的に指定されます。

入力

ラベル 名前 タイプ 制約
(I1) lhs テンソルまたはテンソルごとの量子化テンソル (C1)、(C10 ~ C11)、(C14)(C25)、(C26 ~ C27)、(C30 ~ C31)、(C33)
(I2) rhs テンソルまたは量子化テンソル (C1)、(C14 ~ C16)、(C26 ~ C28)、(C30 ~ C33)
(I3) padding 整数型の 2 次元テンソル (C4)
(I4) window_strides si64 型の 1 次元のテンソル定数 (C2 ~ C3)
(I5) lhs_dilation si64 型の 1 次元のテンソル定数 (C5-C6)
(I6) rhs_dilation si64 型の 1 次元のテンソル定数 (C7 ~ C8)
(I7) window_reversal i1 型の 1 次元のテンソル定数 (C9)
(I8) input_batch_dimension si64 型の定数 (C10)、(C13)
(I9) input_feature_dimension si64 型の定数 (C11)、(C13 ~ C14)
(I10) input_spatial_dimensions si64 型の 1 次元のテンソル定数 (C12)、(C13)
(I11) kernel_input_feature_dimension si64 型の定数 (C14)、(C18)
(I12) kernel_output_feature_dimension si64 型の定数 (C15 ~ C16)、(C18)、(C28)
(I13) kernel_spatial_dimensions si64 型の 1 次元のテンソル定数 (C17 ~ C18)
(I14) output_batch_dimension si64 型の定数 (C20)
(I15) output_feature_dimension si64 型の定数 (C20)、(C29)
(I16) output_spatial_dimensions si64 型の 1 次元のテンソル定数 (C19 ~ C20)
(I17) feature_group_count si64 型の定数 (C11)、(C14)、(C16)、(C21)、(C23)
(I18) batch_group_count si64 型の定数 (C10)、(C15)、(C22)、(C23)
(I19) precision_config DEFAULTHIGHHIGHEST の可変列挙型の数 (C24)

出力

名前 タイプ 制約
result テンソルまたは量子化テンソル (C25 ~ C27)、(C29)、(C31 ~ C33)

制約

  • (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] の場合: <ph type="x-smartling-placeholder">
      </ph>
    • 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] の場合: <ph type="x-smartling-placeholder">
      </ph>
    • 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] の場合: <ph type="x-smartling-placeholder">
      </ph>
    • 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) は次のように定義されます。 <ph type="x-smartling-placeholder">
      </ph>
    • 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
  • オペレーションで量子化されていないテンソルを使用する場合: <ph type="x-smartling-placeholder">
      </ph>
    • (C27)element_type(lhs) = element_type(rhs) = element_type(result)
  • オペレーションで量子化テンソルを使用する場合: <ph type="x-smartling-placeholder">
      </ph>
    • (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]]
//          ]]

その他の例

dynamic_gather

セマンティクス

この処理は ギャザー slice_sizes を値として動的に指定します。

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたはテンソルごとの量子化テンソル (C1)、(C7)、(C10 ~ C12)、(C14)
(I2) start_indices 整数型のテンソル (C2)、(C3)、(C13)
(I3) slice_sizes 整数型の 1 次元テンソル (C8)、(C11 ~ C13)
(I4) offset_dims si64 型の 1 次元のテンソル定数 (C1)、(C4 ~ C5)、(C13)
(I5) collapsed_slice_dims si64 型の 1 次元のテンソル定数 (C1)、(C6 ~ C8)、(C13)
(I6) start_index_map si64 型の 1 次元のテンソル定数 (C3)、(C9)、(C10)
(I7) index_vector_dim si64 型の定数 (C2)、(C3)、(C13)
(I8) indices_are_sorted i1 型の定数

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C5)、(C13 ~ C14)

制約

  • (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)。ここで: <ph type="x-smartling-placeholder">
      </ph>
    • batch_dim_sizes = shape(start_indices)(ディメンションのサイズを除く) index_vector_dim に対応する start_indices は含まれません。
    • offset_dim_sizes = shape(slice_sizes)(ディメンション サイズを除く) collapsed_slice_dims に対応する slice_sizes は含まれません。
    • combine は、batch_dims に対応する軸に batch_dim_sizes を配置します。 offset_dims に対応する軸の offset_dim_sizes
  • (C14)element_type(operand) = element_type(result)

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

その他の例

dynamic_iota

セマンティクス

この処理は iota ただし、結果の形状は output_shape によって動的に指定されます。

入力

ラベル 名前 タイプ 制約
(I1) output_shape 整数型の 1 次元テンソル (C1)、(C2)
(I2) iota_dimension si64 (C1)

出力

名前 タイプ 制約
result 整数、浮動小数点、複合型のテンソル、またはテンソルごとの量子化テンソル (C2)

制約

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

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

その他の例

dynamic_pad

セマンティクス

この処理は パッド 演算。ただし、edge_padding_lowedge_padding_highinterior_padding を使用 値として動的に指定されます。

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたはテンソルごとの量子化テンソル (C1)、(C2)、(C4)
(I2) padding_value 0 次元のテンソルまたはテンソルごとの量子化テンソル (C1)
(I3) edge_padding_low 整数型の 1 次元テンソル (C1)、(C4)
(I4) edge_padding_high 整数型の 1 次元テンソル (C1)、(C4)
(I5) interior_padding 整数型の 1 次元テンソル (C2-C4)

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C3-C6)

制約

  • (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 によって動的に指定されます。

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたは量子化テンソル (C1 ~ C3)
(I2) output_shape 整数型の 1 次元テンソル (C4)

出力

名前 タイプ 制約
result テンソルまたは量子化テンソル (C1 ~ C4)

制約

  • (C1)element_type(result) は次式で求められます。 <ph type="x-smartling-placeholder">
      </ph>
    • 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) の場合: <ph type="x-smartling-placeholder">
      </ph>
    • 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]]

その他の例

dynamic_slice

セマンティクス

動的に計算される開始インデックスを使用して 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

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたはテンソルごとの量子化テンソル (C1)、(C2)、(C4)
(I2) start_indices 整数型の 0 次元テンソルの可変数 (C2)、(C3)
(I3) slice_sizes si64 型の 1 次元のテンソル定数 (C2)、(C4)、(C5)

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C1)、(C5)

制約

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

セマンティクス

次のテンソルを除き、operand テンソルと等しい result テンソルを生成します。 start_indices から始まるスライスが、update の値で更新されます。 より正式には、result[result_index] は次のように定義されます。

  • 0 <= update_index < shape(update) の場合は update[update_index]。ここで: <ph type="x-smartling-placeholder">
      </ph>
    • adjusted_start_indices = clamp(0, start_indices, shape(operand) - shape(update))
    • update_index = result_index - adjusted_start_indices
  • そうでない場合は operand[result_index]

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたはテンソルごとの量子化テンソル (C1 ~ C4)、(C6)
(I2) update テンソルまたはテンソルごとの量子化テンソル (C2)、(C3)、(C6)
(I3) start_indices 整数型の 0 次元テンソルの可変数 (C4)、(C5)

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C1)

制約

  • (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 テンソル。要素のタイプに応じて、次の操作を行います。

  • 浮動小数点数: IEEE-754 の exp
  • 複素数の場合: 複素指数。
  • 量子化型の場合: dequantize_op_quantize(exponential, operand, type(result))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

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

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

その他の例

exponential_minus_one

セマンティクス

operand のテンソルに対して、要素単位の指数から 1 を引いた演算を実行します。 result テンソルを生成します。要素のタイプに応じて、次の操作を行います。

  • 浮動小数点数: IEEE-754 の expm1
  • 複素数の場合: 複素指数 - 1。
  • 量子化型の場合: dequantize_op_quantize(exponential_minus_one, operand, type(result))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

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

セマンティクス

実数と複素数に対する正方向フーリエ変換と逆フーリエ変換を実行 入出力です。

fft_type は、次のいずれかです。

  • FFT: 複雑な FFT を転送します。
  • IFFT: 逆複素数 FFT。
  • RFFT: 実数から複素数への転送 FFT。
  • IRFFT: 逆実数から複素数への変換 FFT(つまり、複素数を受け取り、実数を返します)。

より正式には、次の 1 次元のテンソルを取る関数 fft を考えます。 複合型を入力として、同じ型の 1 次元テンソルを生成 離散フーリエ変換を計算します。

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, ..., :])

さらに、次の 1 次元テンソルを取る関数 rfft があるとします。 浮動小数点数型が、入力変数の複雑な型の 1 次元テンソルを 同じ浮動小数点セマンティクスを持ち、次のように機能します。

  • 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, ..., :])

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点数または複素数のテンソル (C1)、(C2)、(C4)、(C5)
(I2) fft_type FFTIFFTRFFTIRFFT の列挙型 (C2)、(C5)
(I3) fft_length si64 型の 1 次元のテンソル定数 (C1)、(C3)、(C4)

出力

名前 タイプ 制約
result 浮動小数点数または複素数のテンソル (C2)、(C4)、(C5)

制約

  • (C1)size(fft_length) <= rank(operand)
  • (C2)operand 要素と result 要素型の関係はさまざま: <ph type="x-smartling-placeholder">
      </ph>
    • fft_type = FFTelement_type(operand)element_type(result) の場合 同じ複合型です。
    • fft_type = IFFTelement_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)operandresult の間にテンソル real があり、 浮動小数点型の場合は shape(real)[-size(fft_length):] = fft_length になります。
  • (C5)shape(result) = shape(operand)、ただし以下を除く: <ph type="x-smartling-placeholder">
      </ph>
    • 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)]

floor

セマンティクス

operand テンソルの要素単位の階数を実行し、result テンソルを生成します。 IEEE-754 の roundToIntegralTowardNegative オペレーションを実装します。 仕様。量子化型の場合 dequantize_op_quantize(floor, operand, type(result))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C1)

制約

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

その他の例

収集

セマンティクス

start_indices で指定されたオフセットの operand テンソルからスライスを収集 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 は次のように定義されます。 <ph type="x-smartling-placeholder">
      </ph>
    • start_indices[bi0, ..., :, ..., biN]bi は個々の要素) 次の場合、batch_index:index_vector_dim インデックスに挿入されます。 index_vector_dim <rank(start_indices)
    • そうでない場合は [start_indices[batch_index]]
  • axes(operand)d_operandの場合、 <ph type="x-smartling-placeholder">
      </ph>
    • 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
  • axes(operand)d_operandの場合、 <ph type="x-smartling-placeholder">
      </ph>
    • 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_dimsoperand_batching_dims
  • operand_index = full_start_index + full_batching_index + full_offset_index

indices_are_sortedtrue の場合、実装は以下を想定できます。 start_indicesstart_index_map を基準に並べ替えられます。それ以外の場合は、 動作が定義されていません。より正式には、indices(result) のすべての i1 < i2 について、 full_start_index(i1) <= full_start_index(i2)

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたはテンソルごとの量子化テンソル (C1)、(C8)、(C11)、(C17)、(C19 ~ C21)、(C23)
(I2) start_indices 整数型のテンソル (C2 ~ C3)、(C14)、(C17)、(C22)
(I3) offset_dims si64 型の 1 次元のテンソル定数 (C1)、(C4 ~ C5)、(C22)
(I4) collapsed_slice_dims si64 型の 1 次元のテンソル定数 (C1)、(C6 ~ C9)、(C22)
(I5) operand_batching_dims si64 型の 1 次元のテンソル定数 (C1)、(C6)、(C10 ~ C12)、(C16 ~ C18)、(C22)
(I6) start_indices_batching_dims si64 型の 1 次元のテンソル定数 (C13 ~ C17)
(I7) start_index_map si64 型の 1 次元のテンソル定数 (C3)、(C18 ~ C19)
(I8) index_vector_dim si64 型の定数 (C2 ~ C3)、(C15)、(C22)
(I9) slice_sizes si64 型の 1 次元のテンソル定数 (C9)、(C12)、(C20 ~ C22)
(I10) indices_are_sorted i1 型の定数

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C5)、(C22 ~ C23)

制約

  • (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)。ここで: <ph type="x-smartling-placeholder">
      </ph>
    • batch_dim_sizes = shape(start_indices)(ディメンションのサイズを除く) index_vector_dim に対応する start_indices は含まれません。
    • offset_dim_sizes = slice_sizes。ただし、ディメンションのサイズは slice_sizescollapsed_slice_dims に対応し、 operand_batching_dims は含まれません。
    • combine は、batch_dims に対応する軸に batch_dim_sizes を配置します。 offset_dims に対応する軸の offset_dim_sizes
  • (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

セマンティクス

指定された operanddimension のサイズを生成します。よりフォーマルな表現では result = dim(operand, dimension)。セマンティクスはシェイプのみに関係します。 型のコンポーネントです。要素タイプは何かです。

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたは量子化テンソル (C1)
(I2) dimension si64 型の定数 (C1)

出力

名前 タイプ
result si32 型の 0 次元テンソル

制約

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

<ph type="x-smartling-placeholder">
</ph>

セマンティクス

operand タプルの index の位置で要素を抽出し、 result。(より正式には、result = operand[index])。

入力

ラベル 名前 タイプ 制約
(I1) operand tuple (C1)、(C2)
(I2) index si32 型の定数 (C1)、(C2)

出力

名前 タイプ 制約
result サポートされている任意の型 (C2)

制約

  • (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 から関数を 1 つだけ実行して出力を生成する。 pred の値に応じて false_branch。(より正式には、result = pred ? true_branch() : false_branch())。

入力

ラベル 名前 タイプ 制約
(I1) pred i1 型の 0 次元テンソル
(I2) true_branch 関数 (C1 ~ C3)
(I3) false_branch 関数 (C1)、(C2)

出力

名前 タイプ 制約
results テンソル、量子化テンソル、またはトークンの可変数 (C3)

制約

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

その他の例

Image

セマンティクス

operand から虚数部を要素ごとに抽出し、 result テンソル。より正式には、各要素 x に対して次のように設定します。 imag(x) = is_complex(x) ? imaginary_part(x) : constant(0, element_type(result))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点数または複素数のテンソル (C1)、(C2)

出力

名前 タイプ 制約
result 浮動小数点型のテンソル (C1)、(C2)

制約

  • (C1)shape(result) = shape(operand)
  • (C2)element_type(result) は次のように定義されます。 <ph type="x-smartling-placeholder">
      </ph>
    • is_complex(operand) の場合は complex_element_type(element_type(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 は、最初に来るペイロード値と次のトークンで構成されます。 最後です。将来的には、ペイロードとトークンを 2 つに分割する予定です。 出力を分離してわかりやすくする (#670)。

入力

ラベル 名前 タイプ
(I1) token token
(I2) infeed_config string 型の定数

出力

名前 タイプ 制約
results テンソル、量子化テンソル、またはトークンの可変数 (C1 ~ C3)

制約

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

その他の例

Iota

セマンティクス

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

入力

ラベル 名前 タイプ 制約
(I1) iota_dimension si64 (C1)

出力

名前 タイプ 制約
output 整数、浮動小数点、複合型のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

  • (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 を実装する 準拠しています。量子化型の場合、結果は 常に true

入力

ラベル 名前 タイプ 制約
(I1) x 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
y ブール型のテンソル (C1)

制約

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

その他の例

log

セマンティクス

operand のテンソルに対して要素単位の対数演算を実行し、 result テンソル。要素のタイプに応じて、次の操作を行います。

  • 浮動小数点数: IEEE-754 の log
  • 複素数の場合: 複素対数
  • 量子化型の場合: dequantize_op_quantize(log, operand, type(result))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

  • (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 のテンソルに対して要素単位の対数と 1 演算を加え、 result テンソルを生成します。要素のタイプに応じて、次の操作を行います。

  • 浮動小数点数: IEEE-754 の logp1
  • 複素数の場合: 複素対数 + 1
  • 量子化型の場合: dequantize_op_quantize(log_plus_one, operand, type(result))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

  • (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 テンソル。要素のタイプに応じて、次の操作を行います。

  • 浮動小数点数: IEEE-754 の division(1, addition(1, exp(-x)))
  • 複素数の場合: 複素ロジスティック。
  • 量子化型の場合: dequantize_op_quantize(logistic, operand, type(result))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

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

その他の例

地図

<ph type="x-smartling-placeholder">
</ph>

セマンティクス

マップ関数 computationdimensions に沿って inputs に適用します。 result テンソルを生成します。

(より正式には、result[result_index] = computation(inputs...[result_index]))。

入力

ラベル 名前 タイプ 制約
(I1) inputs テンソルの可変数またはテンソルごとの量子化テンソル (C1 ~ C4)
(I2) dimensions si64 型の 1 次元のテンソル定数 (C3)
(I3) computation 関数 (C4)

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C1)、(C4)

制約

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

その他の例

最大

セマンティクス

テンソル lhsrhs に対して要素単位の最大演算を実行し、次を生成します。 result のテンソル。要素のタイプに応じて、次の操作を行います。

  • ブール値の場合: 論理 OR。
  • 整数の場合: 最大値の整数。
  • 浮動小数点数: IEEE-754 の maximum
  • 複素数の場合: (real, imaginary) ペアの辞書順の最大値。 複素数に順序を付けるには、意外な意味論が必要です。 将来的には複素数のサポートは終了し (#560)。
  • 量子化型の場合: <ph type="x-smartling-placeholder">
      </ph>
    • dequantize_op_quantize(maximum, lhs, rhs, type(result))

入力

ラベル 名前 タイプ 制約
(I1) lhs テンソルまたはテンソルごとの量子化テンソル (C1)
(I2) rhs テンソルまたはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C1)

制約

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

その他の例

最小

セマンティクス

テンソル lhsrhs に対して要素単位の min 演算を実行し、 result のテンソル。要素のタイプに応じて、次の操作を行います。

  • ブール値の場合: 論理 AND。
  • 整数の場合: 整数の最小値。
  • 浮動小数点数: IEEE-754 の minimum
  • 複素数の場合: (real, imaginary) ペアの辞書順の最小値。 複素数に順序を付けるには、意外な意味論が必要です。 将来的には複素数のサポートは終了し (#560)。
  • 量子化型の場合: <ph type="x-smartling-placeholder">
      </ph>
    • dequantize_op_quantize(minimum, lhs, rhs, type(result))

入力

ラベル 名前 タイプ 制約
(I1) lhs テンソルまたはテンソルごとの量子化テンソル (C1)
(I2) rhs テンソルまたはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C1)

制約

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

その他の例

乗算

セマンティクス

2 つのテンソル lhsrhs の要素単位の積を実行し、 result のテンソル。要素のタイプに応じて、次の操作を行います。

  • ブール値の場合: 論理 AND。
  • 整数の場合: 整数の乗算。
  • 浮動小数点数: IEEE-754 の multiplication
  • 複素数の場合: 複素乗算。
  • 量子化型の場合: <ph type="x-smartling-placeholder">
      </ph>
    • dequantize_op_quantize(multiply, lhs, rhs, type(result))

入力

ラベル 名前 タイプ 制約
(I1) lhs テンソルまたはテンソルごとの量子化テンソル (C1)
(I2) rhs テンソルまたはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C1)

制約

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

その他の例

negate

セマンティクス

operand テンソルの要素単位の否定を実行し、result を生成します。 テンソルです。要素のタイプに応じて、次の操作を行います。

  • 符号付き整数の場合: 整数の否定。
  • 符号なし整数の場合: 符号付き整数へのビットキャスト、整数否定、ビットキャスト 符号なし整数に戻します。
  • 浮動小数点数: IEEE-754 の negate
  • 複素数の場合: 複素否定。
  • 量子化型の場合: dequantize_op_quantize(negate, operand, type(result))

入力

ラベル 名前 タイプ 制約
(I1) operand 整数、浮動小数点、複合型のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 整数、浮動小数点、複合型のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

  • (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 の要素単位の NOT を実行し、result テンソルを生成します。 要素のタイプに応じて、次の操作を行います。

  • ブール値の場合: 論理 NOT。
  • 整数の場合: ビットごとの NOT。

引数

名前 タイプ 制約
operand ブール型または整数型のテンソル (C1)

出力

名前 タイプ 制約
result ブール型または整数型のテンソル (C1)

制約

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

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

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

その他の例

optimization_barrier

セマンティクス

operand を生成するオペレーションが result に依存し、コンパイラ変換を妨げる演算 障壁を越えたオペレーションの移行を防げます。それ以外は、オペレーションが ID(例: result = operand)。

引数

名前 タイプ 制約
operand テンソルの可変数、テンソルごとの量子化テンソルまたはトークン (C1)

出力

名前 タイプ 制約
result テンソルの可変数、テンソルごとの量子化テンソルまたはトークン (C1)

制約

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

その他の例

または

セマンティクス

2 つのテンソル lhsrhs の要素単位の OR を実行し、result を生成します。 テンソルです。要素のタイプに応じて、次の操作を行います。

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

入力

ラベル 名前 タイプ 制約
(I1) lhs 整数型またはブール型のテンソル (C1)
(I2) rhs 整数型またはブール型のテンソル (C1)

出力

名前 タイプ 制約
result 整数型またはブール型のテンソル (C1)

制約

  • (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 のセマンティクスは実装で定義されます。

入力

ラベル 名前 タイプ
(I1) inputs テンソルまたは量子化テンソルの可変数
(I2) token token
(I3) outfeed_config string 型の定数

出力

名前 タイプ
result token

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

その他の例

パッド

セマンティクス

テンソルの周囲や要素間にパディングを追加して operand を拡張します。 指定された padding_value を持つテンソルの約 1.5 倍です。

edge_padding_lowedge_padding_high は、追加するパディングの量を指定します インデックス値の低い値(インデックス 0 の次)と上限(インデックスの最大値の次)で示されます。 作成しますパディングの量は負にすることもできます。この場合、 負のパディングの絶対値は、削除する要素の数を示します。 すべての要素に適用されます。

interior_padding は、任意の 2 つの間に追加するパディングの量を指定します。 各次元の要素のうち負の値は使用できません。内部パディングあり 負のエッジ パディングによって要素が取り除かれるように 内部パディング オペランド。

より正式には、result[result_index] は次のように定義されます。

  • 次の場合は operand[operand_index] result_index = edge_padding_low + operand_index * (interior_padding + 1)
  • そうでない場合は padding_value

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたはテンソルごとの量子化テンソル (C1)、(C2)、(C4)
(I2) padding_value 0 次元のテンソルまたはテンソルごとの量子化テンソル (C1)
(I3) edge_padding_low si64 型の 1 次元のテンソル定数 (C1)、(C4)
(I4) edge_padding_high si64 型の 1 次元のテンソル定数 (C1)、(C4)
(I5) interior_padding si64 型の 1 次元のテンソル定数 (C2-C4)

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C3-C6)

制約

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

セマンティクス

現在のプロセスの partition_id を生成します。

出力

名前 タイプ
result ui32 型の 0 次元テンソル

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

その他の例

ポップ

セマンティクス

operand テンソルに設定されたビット数を要素ごとにカウントします result テンソルを生成します。

入力

ラベル 名前 タイプ 制約
(I1) operand 整数型のテンソル (C1)

出力

名前 タイプ 制約
result 整数型のテンソル (C1)

制約

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

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

その他の例

電力

セマンティクス

rhs テンソルで lhs テンソルの要素単位の指数化を実行します。 result テンソルが生成されます。要素のタイプに応じて、次の操作を行います。

  • 整数の場合: 整数のべき乗。
  • 浮動小数点数: IEEE-754 の pow
  • 複素数の場合: 複素べき乗。
  • 量子化型の場合: dequantize_op_quantize(power, lhs, rhs, type(result))

入力

ラベル 名前 タイプ 制約
(I1) lhs 整数、浮動小数点、複合型のテンソル、またはテンソルごとの量子化テンソル (C1)
(I2) rhs 整数、浮動小数点、複合型のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 整数、浮動小数点、複合型のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

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

その他の例

real

セマンティクス

operand から実数部分を要素ごとに抽出し、result を生成します。 テンソルです。より正式には、各要素 x に対して次のように設定します。 real(x) = is_complex(x) ? real_part(x) : x

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点数または複素数のテンソル (C1)、(C2)

出力

名前 タイプ 制約
result 浮動小数点型のテンソル (C1)、(C2)

制約

  • (C1)shape(result) = shape(operand)
  • (C2)element_type(result) は次のように定義されます。 <ph type="x-smartling-placeholder">
      </ph>
    • is_complex(operand) の場合は complex_element_type(element_type(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_transfertrue の場合、オペレーションはデータを ホストします。それ以外の場合は、別のデバイスからデータが転送されます。つまり 使用されます。このフラグは、 channel_type のため、今後はどちらか一方のみを残す予定です。 (#666)。

results は、最初に来るペイロード値と次のトークンで構成されます。 最後です。将来的には、ペイロードとトークンを 2 つに分割する予定です。 出力を分離してわかりやすくする (#670)。

入力

ラベル 名前 タイプ 制約
(I1) token token (C4)
(I2) channel_id si64 型の定数
(I3) channel_type DEVICE_TO_DEVICEHOST_TO_DEVICE の列挙型 (C1)
(I4) is_host_transfer i1 型の定数 (C1)

出力

名前 タイプ 制約
results テンソル、量子化テンソル、またはトークンの可変数 (C2-C4)

制約

  • (C1)channel_type は次のように定義されます。 <ph type="x-smartling-placeholder">
      </ph>
    • is_host_transfer = true の場合は HOST_TO_DEVICE
    • そうでない場合は 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)

その他の例

reduce

セマンティクス

リダクション関数 bodyinputsinit_values に沿って適用します。 dimensionsresults テンソルを生成します。

削減の順序は実装によって定義されます。つまり、bodyinit_values は、このオペレーションによって確実に すべての実装で、すべての入力に対して同じ結果が得られます。ただし この条件は 一般的な値下げ法の多くには当てはまりません。例:浮動小数点加算を init_valuesbody と 0 は実際にはモノイドを形成しません。理由は次のとおりです。 浮動小数点の加算は結合的ではありません。

より正式には、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 各要素の意味は次のとおりです。 <ph type="x-smartling-placeholder">
      </ph>
    • exec(node) = body(exec(node.left), exec(node.right))
    • exec(leaf) = leaf.value
  • schedule は実装定義の完全バイナリツリーであり、 トラバーサルは次のものから構成されます。 <ph type="x-smartling-placeholder">
      </ph>
    • input_slices_converted...[index] 個の値、すべての index 辞書順の昇順の index_space(input_slices_converted) /index
    • 実装定義の量と散在している 実装で定義されている位置の init_values_converted

入力

ラベル 名前 タイプ 制約
(I1) inputs テンソルの可変数またはテンソルごとの量子化テンソル (C1 ~ C4)、(C6)、(C7)
(I2) init_values 0 次元テンソルまたはテンソルごとの量子化テンソルの可変数 (C2)、(C3)
(I3) dimensions si64 型の 1 次元のテンソル定数 (C4)、(C5)、(C7)
(I4) body 関数 (C6)

出力

名前 タイプ 制約
results テンソルの可変数またはテンソルごとの量子化テンソル (C3)、(C7)、(C8)

制約

  • (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...)(ただし、ディメンションは dimensions に対応する inputs... のサイズは含まれません。
  • (C8)[0,N) のすべての i に対して element_type(results[i]) = Ei

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

その他の例

reduce_precision

セマンティクス

operand を別の浮動小数点型に要素単位で変換します exponent_bitsmantissa_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)) を実行します。

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C1)
(I2) exponent_bits si32 型の定数 (C2)
(I3) mantissa_bits si32 型の定数 (C3)

出力

名前 タイプ 制約
output 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C1)

制約

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

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

その他の例

reduce_scatter

セマンティクス

reduce_scatter

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_groupreceiver_index = process_group.index(receiver))。

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたはテンソルごとの量子化テンソル (C1)、(C2)、(C7)、(C8)
(I2) scatter_dimension si64 型の定数 (C1)、(C2)、(C8)
(I3) replica_groups si64 型の 2 次元のテンソル定数 (C3-C5)
(I4) channel_id si64 型の定数 (C6)
(I5) use_global_device_ids i1 型の定数 (C6)
(I6) computation 関数 (C7)

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C8-C9)

制約

  • (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) は次のように定義されます。 <ph type="x-smartling-placeholder">
      </ph>
    • cross_replica が使用されている場合は num_replicas
    • cross_replica_and_partition が使用されている場合は num_replicas
    • flattened_ids が使用されている場合は num_processes
  • (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)、ただし: <ph type="x-smartling-placeholder">
      </ph>
    • dim(result, scatter_dimension) = dim(operand, scatter_dimension) / dim(process_groups, 1)
  • (C9)element_type(result) = E

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

その他の例

reduce_window

セマンティクス

リダクション関数 bodyinputsinit_values のウィンドウに適用します。 results が生成されます。

次の図は、results... の要素がどのように計算されるかを示しています。 inputs... を使用します。

reduce_window

よりフォーマルな表現では results...[result_index] = reduce(windows, init_values, axes(inputs...), body)reduce を参照)ここで:

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

入力

ラベル 名前 タイプ 制約
(I1) inputs テンソルの可変数またはテンソルごとの量子化テンソル (C1 ~ C4)、(C6)、(C8)、(C10)、(C12)、(C13)、(C15)
(I2) init_values 0 次元テンソルまたはテンソルごとの量子化テンソルの可変数 (C1)、(C13)
(I3) window_dimensions si64 型の 1 次元のテンソル定数 (C4)、(C5)、(C15)
(I4) window_strides si64 型の 1 次元のテンソル定数 (C6)、(C7)、(C15)
(I5) base_dilations si64 型の 1 次元のテンソル定数 (C8)、(C9)、(C15)
(I6) window_dilations si64 型の 1 次元のテンソル定数 (C10)、(C11)、(C15)
(I7) padding si64 型の 2 次元のテンソル定数 (C12)、(C15)
(I8) body 関数 (C13)

出力

名前 タイプ 制約
results テンソルの可変数またはテンソルごとの量子化テンソル (C1)、(C14 ~ C16)

制約

  • (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。ここで: <ph type="x-smartling-placeholder">
      </ph>
    • 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)[0,N) のすべての i に対して element_type(results[i]) = Ei

// %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)
  • 浮動小数点数: IEEE-754 の division(lhs, rhs)(丸め属性あり) roundTowardZero
  • 複素数の場合: 未定 (#997)。
  • 量子化型の場合: <ph type="x-smartling-placeholder">
      </ph>
    • dequantize_op_quantize(remainder, lhs, rhs, type(result))

浮動小数点要素型の場合、この操作は IEEE-754 仕様の remainder 演算(d は整数値) lhs/rhs の正確な値に最も近い値に最も近い値を選択します。

入力

ラベル 名前 タイプ 制約
(I1) lhs 整数型、浮動小数点数、複素数のテンソル、またはテンソルごとの量子化テンソル (C1)
(I2) rhs 整数型、浮動小数点数、複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 整数型、浮動小数点数、複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

  • (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 ui32 型の 0 次元テンソル

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

その他の例

reshape

セマンティクス

operand テンソルを result テンソルに再形成します。概念的には、 同じ正規表現を維持できますが、変更される可能性がある 適用できます。例:tensor<2x3xf32> から tensor<3x2xf32> または tensor<6xf32> に変更する。

より正式には、result[result_index] = operand[operand_index] result_indexoperand_index は辞書順での位置が同じです index_space(result)index_space(operand) の順序付け。

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたは量子化テンソル (C1 ~ C3)

出力

名前 タイプ 制約
result テンソルまたは量子化テンソル (C1 ~ C3)

制約

  • (C1)element_type(result) は次式で求められます。 <ph type="x-smartling-placeholder">
      </ph>
    • 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) の場合: <ph type="x-smartling-placeholder">
      </ph>
    • 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]]

その他の例

reverse

セマンティクス

指定された dimensions に沿って operand 内の要素の順序を逆にします。 result のテンソルが生成されます。よりフォーマルな表現では result[result_index] = operand[operand_index] 各要素の意味は次のとおりです。

  • operand_index[d] = dim(result, d) - result_index[d] - 1 dimensionsd の場合。
  • そうでない場合は operand_index[d] = result_index[d]

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたはテンソルごとの量子化テンソル (C1)、(C3)
(I2) dimensions si64 型の 1 次元のテンソル定数 (C2)、(C3)

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C1)、(C3)

制約

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

<ph type="x-smartling-placeholder">
</ph>

セマンティクス

rng_distribution アルゴリズムを使用して乱数を生成し、 指定された形状 shaperesult テンソル。

rng_distribution = UNIFORM の場合、乱数が生成されます。 区間 [a, b) の均一分布に従います。a >= b の場合、 動作は未定義です。

rng_distribution = NORMAL の場合、乱数が生成されます。 平均 = a、標準偏差 = b の正規分布に従う b < 0 の場合、動作は未定義です。

乱数がどのように生成されるかは実装で定義します。対象 たとえば、決定論的場合もあれば、そうでない場合もあります。また、 隠されています。

多くの関係者との対話において、この商談は 将来的には削除を検討し (#597)。

入力

ラベル 名前 タイプ 制約
(I1) a 整数型、ブール値型、浮動小数点型の 0 次元テンソル (C1)、(C2)
(I2) b 整数型、ブール値型、浮動小数点型の 0 次元テンソル (C1)、(C2)
(I3) shape si64 型の 1 次元のテンソル定数 (C3)
(I4) rng_distribution UNIFORMNORMAL の列挙型 (C2)

出力

名前 タイプ 制約
result 整数型、ブール値型、浮動小数点型のテンソル (C1 ~ C3)

制約

  • (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出力は 0 ~ 100 ミリ秒で initial_state の決定的関数ですが、必ずしも 決定的です。

rng_algorithm は、次のいずれかです。

  • DEFAULT: 実装定義のアルゴリズム。
  • THREE_FRY: Threefry アルゴリズムの実装定義バリアント*。
  • PHILOX: 実装定義の Philox アルゴリズムのバリアント*。

* 参照: Salmon et al.SC 2011。並列乱数:1、2、3 と同じくらい簡単。 で確認できます。

入力

ラベル 名前 タイプ 制約
(I1) rng_algorithm DEFAULTTHREE_FRYPHILOX の列挙型 (C2)
(I2) initial_state ui64 型の 1 次元テンソル (C1)、(C2)

出力

名前 タイプ 制約
output_state ui64 型の 1 次元テンソル (C1)
output 整数型または浮動小数点型のテンソル

制約

  • (C1)type(initial_state) = type(output_state)
  • (C2)size(initial_state) は次のように定義されます。 <ph type="x-smartling-placeholder">
      </ph>
    • rng_algorithm = DEFAULT の場合は、実装で定義されます。
    • rng_algorithm = THREE_FRY の場合は 2
    • rng_algorithm = PHILOX の場合は 2 または 3

// %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 テンソルを生成します。実装 IEEE-754 仕様の roundToIntegralTiesToAway オペレーション。対象 処理されるため、 dequantize_op_quantize(round_nearest_afz, operand, type(result))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C1)

制約

  • (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 を生成する テンソルです。IEEE-754 の roundToIntegralTiesToEven オペレーションを実装します。 仕様。量子化型の場合 dequantize_op_quantize(round_nearest_even, operand, type(result))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 浮動小数点型のテンソルまたはテンソルごとの量子化テンソル (C1)

制約

  • (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 テンソルを生成します。要素のタイプに応じて、次の操作を行います。

  • 浮動小数点数: IEEE-754 の rSqrt
  • 複素数の場合: 複素数の逆平方根。
  • 量子化型の場合: dequantize_op_quantize(rsqrt, operand, type(result))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

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

その他の例

散布

セマンティクス

次を除く inputs テンソルと等しい results テンソルを生成します。 scatter_indices で指定された複数のスライスが、値で更新される update_computation を使用する updates

次の図は、updates... の要素が results... を使用します。この図はいくつかの例を示しています updates... は、どの results... インデックスを使用するかを詳しく説明しています。 表します。

散布

より正式には、index_space(updates[0]) のすべての update_index に対して以下を行います。

  • 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 は次のように定義されます。 <ph type="x-smartling-placeholder">
      </ph>
    • scatter_indices[si0, ..., :, ..., siN]si は個別) update_scatter_index: の要素が index_vector_dim インデックス(index_vector_dim が < rank(scatter_indices)
    • そうでない場合は [scatter_indices[update_scatter_index]]
  • d_inputaxes(inputs[0]))の場合、 <ph type="x-smartling-placeholder">
      </ph>
    • 次の場合は full_start_index[d_input] = start_index[d_start] d_input = scatter_dims_to_operand_dims[d_start]
    • そうでない場合は full_start_index[d_input] = 0
  • axes(inputs[0])d_inputの場合、 <ph type="x-smartling-placeholder">
      </ph>
    • 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_dimsinput_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) 各要素の意味は次のとおりです。 <ph type="x-smartling-placeholder">
      </ph>
    • result_indexshape(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_resultsresults のコピー(results...[result_index] を含む)です updated_values... に設定。
    • それ以外の場合
    • updated_results = results
  • exec([], results) = results

indices_are_sortedtrue の場合、実装は以下を想定できます。 scatter_indicesscatter_dims_to_operand_dims を基準として並べ替えられます。 それ以外の場合の動作は未定義です。より正式に、すべての i1 < i2indices(result)full_start_index(i1) <= full_start_index(i2)

unique_indicestrue の場合、実装はすべての 散布される result_index 個のインデックスは一意です。unique_indicestrue ですが、散らばっているインデックスが一意でない場合の動作は次のようになります。 未定義です。

入力

ラベル 名前 タイプ 制約
(I1) inputs テンソルの可変数またはテンソルごとの量子化テンソル (C1)、(C2)、(C4 ~ C6)、(C11)、(C13)、(C18)、(C21)、(C23 ~ C24)
(I2) scatter_indices 整数型のテンソル (C4)、(C15)、(C19)、(C22)
(I3) updates テンソルの可変数またはテンソルごとの量子化テンソル (C3-C6)、(C8)
(I4) update_window_dims si64 型の 1 次元のテンソル定数 (C2)、(C4)、(C7 ~ C8)
(I5) inserted_window_dims si64 型の 1 次元のテンソル定数 (C2)、(C4)、(C9 ~ C11)
(I6) input_batching_dims si64 型の 1 次元のテンソル定数 (C2)、(C4)、(C9)、(C12-13)、(C17-18)、(C20)
(I7) scatter_indices_batching_dims si64 型の 1 次元のテンソル定数 (C14 ~ C18)
(I8) scatter_dims_to_operand_dims si64 型の 1 次元のテンソル定数 (C19 ~ C21)
(I9) index_vector_dim si64 型の定数 (C4)、(C16)、(C19)、(C22)
(I10) indices_are_sorted i1 型の定数
(I11) unique_indices i1 型の定数
(I12) update_computation 関数 (C23)

出力

名前 タイプ 制約
results テンソルの可変数またはテンソルごとの量子化テンソル (C24-C25)

制約

  • (C1)same(shape(inputs...))
  • (C2)`rank(inputs[0]) = size(update_window_dims) + size(inserted_window_dims) <ph type="x-smartling-placeholder">
      </ph>
    • size(input_batching_dims)`.
  • (C3)same(shape(updates...))
  • (C4)shape(updates[0]) = combine(update_scatter_dim_sizes, update_window_dim_sizes)。ここで: <ph type="x-smartling-placeholder">
      </ph>
    • 次を除く update_scatter_dim_sizes = shape(scatter_indices) 対応する scatter_indices のディメンション サイズ index_vector_dim は含まれません。
    • 次を除く update_window_dim_sizes <= shape(inputs[0]) inserted_window_dims に対応する inputs[0] のディメンション サイズ と input_batching_dims は含まれません。
    • combine は、対応する軸に update_scatter_dim_sizes を配置します。 対応する軸の update_scatter_dimsupdate_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)[0,N) のすべての i に対して element_type(results[i]) = Ei

// %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 または pred の対応する要素の値に基づく on_false テンソル。 より正式には、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))

入力

ラベル 名前 タイプ 制約
(I1) pred i1 型のテンソル (C1)
(I2) on_true テンソルまたはテンソルごとの量子化テンソル (C1 ~ C2)
(I3) on_false テンソルまたはテンソルごとの量子化テンソル (C2)

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C2)

制約

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

セマンティクス

次に基づいて scatter を使用して source テンソルの値を分散します。 select を使用した input テンソルの reduce_window の結果。次の式が生成される: result テンソル。

次の図は、result の要素がどのように計算されるかを示しています。 operandsource を具体例で説明します。

select_and_scatter

よりフォーマルな表現:

  • selected_values = reduce_window_without_init(...) は、次の入力に置き換えます。

    • inputs = [operand].
    • window_dimensionswindow_stridespadding はそのまま使用されます。
    • 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 と完全に同じですが、基になるオブジェクトの schedulereducereduce を参照)には 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 selected_values[source_index] には operand 要素があります operand_index~。
    • source_indices = [source_index for source_index in indices(source) if selected_index(source_index) = result_index]

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたはテンソルごとの量子化テンソル (C1 ~ C4)、(C6)、(C8 ~ C11)
(I2) source テンソルまたはテンソルごとの量子化テンソル (C1)、(C2)
(I3) init_value 0 次元のテンソルまたはテンソルごとの量子化テンソル (C3)
(I4) window_dimensions si64 型の 1 次元のテンソル定数 (C2)、(C4)、(C5)
(I5) window_strides si64 型の 1 次元のテンソル定数 (C2)、(C6)、(C7)
(I6) padding si64 型の 2 次元のテンソル定数 (C2)、(C8)
(I7) select 関数 (C9)
(I8) scatter 関数 (C10)

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C11 ~ C12)

制約

  • (C1)element_type(operand) = element_type(source)
  • (C2)shape(source) = num_windows。ここで: <ph type="x-smartling-placeholder">
      </ph>
    • 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 の型は (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_transfertrue の場合、オペレーションによりデータが ホストします。それ以外の場合は、別のデバイスにデータが転送されます。つまり 使用されます。このフラグは、 channel_type のため、今後はどちらか一方のみを残す予定です。 (#666)。

入力

ラベル 名前 タイプ 制約
(I1) inputs テンソルまたは量子化テンソルの可変数
(I2) token token
(I3) channel_id si64 型の定数
(I4) channel_type DEVICE_TO_DEVICEDEVICE_TO_HOST の列挙型 (C1)
(I5) is_host_transfer i1 型の定数 (C1)

出力

名前 タイプ
result token

制約

  • (C1)channel_type は次のように定義されます。 <ph type="x-smartling-placeholder">
      </ph>
    • is_host_transfer = true の場合は DEVICE_TO_HOST
    • そうでない場合は 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 のテンソルが生成されます。

入力

ラベル 名前 タイプ 制約
(I1) lhs 整数型のテンソル (C1)
(I2) rhs 整数型のテンソル (C1)

出力

名前 タイプ 制約
result 整数型のテンソル (C1)

制約

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

入力

ラベル 名前 タイプ 制約
(I1) lhs 整数型のテンソル (C1)
(I2) rhs 整数型のテンソル (C1)

出力

名前 タイプ 制約
result 整数型のテンソル (C1)

制約

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

セマンティクス

rhslhs テンソルに対して要素単位の右シフト演算を実行します。 ビット数を計算し、result テンソルを生成します。

入力

ラベル 名前 タイプ 制約
(I1) lhs 整数型のテンソル (C1)
(I2) rhs 整数型のテンソル (C1)

出力

名前 タイプ 制約
result 整数型のテンソル (C1)

制約

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

入力

ラベル 名前 タイプ 制約
(I1) operand 符号付き整数、浮動小数点数、複合型のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 符号付き整数、浮動小数点数、複合型のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

  • (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 を生成します。 テンソルです。要素のタイプに応じて、次の操作を行います。

  • 浮動小数点数: IEEE-754 の sin
  • 複素数の場合: 複素正弦
  • 量子化型の場合: dequantize_op_quantize(sine, operand, type(result))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

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

その他の例

slice

セマンティクス

静的に計算された開始インデックスを使用して operand からスライスを抽出します result テンソルを生成します。start_indices には、次の開始インデックスが含まれます。 各ディメンションのスライス。limit_indices には終了インデックスが含まれます。 各ディメンションのスライスの(この値を含まない)、strides にはストライドが含まれます 作成します

より正式には、result[result_index] = operand[operand_index] operand_index = start_indices + result_index * strides

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたはテンソルごとの量子化テンソル (C1 ~ C3)、(C5)
(I2) start_indices si64 型の 1 次元のテンソル定数 (C2)、(C3)、(C5)
(I3) limit_indices si64 型の 1 次元のテンソル定数 (C2)、(C3)、(C5)
(I4) strides si64 型の 1 次元のテンソル定数 (C2)、(C4)

出力

名前 タイプ 制約
result テンソルまたはテンソルごとの量子化テンソル (C1)、(C5)

制約

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

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

その他の例

並べ替え

セマンティクス

inputs の 1 次元のスライスを、ディメンション dimension に沿ってまとめて並べ替えます。 comparator に従って results を生成します。

他の演算の類似入力とは異なり、dimension では負の値、 次のセマンティクスで示されます。将来的には、 パフォーマンス上の理由から (#1377)。

is_stable が true の場合、並べ替えは安定しています(相対順序)。 コンパレータで等しいと見なされた要素が保持されます。ケース 入力が 1 つの場合、e1e2 の 2 つの要素とみなされます。 次の場合にのみ、コンパレータによって comparator(e1, e2) = comparator(e2, e1) = false。以下の形式をご覧ください。 複数入力に一般化する方法を見てみましょう。

より正式には、index_space(results[0]) のすべての result_index に対して以下を行います。

  • adjusted_dimension = dimension >= 0 ? dimension : rank(inputs[0]) + dimension
  • result_slice = [ri0, ..., :, ..., riR-1]riN は個別) result_index: の要素は adjusted_dimension に挿入されます。
  • inputs_together = (inputs[0]..., ..., inputs[N-1]...)
  • results_together[result_slice] = sort(inputs_together[result_slice], comparator_together)
  • ここで、sort は 1 次元のスライスを非降順で並べ替えます。 左側の引数が一致していれば、comparator_togethertrue を返します。 右の秒引数より小さくなります。
  • def comparator_together(lhs_together, rhs_together):
      args = []
      for (lhs_el, rhs_el) in zip(lhs_together, rhs_together):
        args.append(lhs_el)
        args.append(rhs_el)
      return comparator(*args)
    
  • (results[0]..., ..., results[N-1]...) = results_together

入力

ラベル 名前 タイプ 制約
(I1) inputs テンソルの可変数またはテンソルごとの量子化テンソル (C1 ~ C5)
(I2) dimension si64 型の定数 (C4)
(I3) is_stable i1 型の定数
(I4) comparator 関数 (C5)

出力

名前 タイプ 制約
results テンソルの可変数またはテンソルごとの量子化テンソル (C2)、(C3)

制約

  • (C1)0 < size(inputs)
  • (C2)type(inputs...) = type(results...)
  • (C3)same(shape(inputs...) + shape(results...))
  • (C4)-R <= dimension < RR = rank(inputs[0]))。
  • (C5)comparator が type である (tensor<E1>, tensor<E1>, ..., tensor<EN-1>, tensor<EN-1>) -> tensor<i1>, ここで、Ei = element_type(inputs[i]) になります。

// %input0 = [[1, 2, 3], [3, 2, 1]]
// %input1 = [[3, 2, 1], [1, 2, 3]]
%result0, %result1 = "stablehlo.sort"(%input0, %input1) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>, %arg2: tensor<i64>, %arg3: tensor<i64>):
    %predicate = "stablehlo.compare"(%arg0, %arg1) {
      comparison_direction = #stablehlo<comparison_direction GT>
    } : (tensor<i64>, tensor<i64>) -> tensor<i1>
    "stablehlo.return"(%predicate) : (tensor<i1>) -> ()
}) {
  dimension = 0 : i64,
  is_stable = true
} : (tensor<2x3xi64>, tensor<2x3xi64>) -> (tensor<2x3xi64>, tensor<2x3xi64>)
// %result0 = [[3, 2, 3], [1, 2, 1]]
// %result1 = [[1, 2, 1], [3, 2, 3]]

その他の例

sqrt

セマンティクス

operand のテンソルに対して要素単位の平方根演算を実行し、 result テンソル。要素のタイプに応じて、次の操作を行います。

  • 浮動小数点数: IEEE-754 の squareRoot
  • 複素数の場合: 複素平方根。
  • 量子化型の場合: dequantize_op_quantize(sqrt, operand, type(result))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

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

// %operand: [[0.0, 1.0], [4.0, 9.0]]
%result = "stablehlo.sqrt"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[0.0, 1.0], [2.0, 3.0]]

その他の例

subtract

セマンティクス

2 つのテンソル lhsrhs の要素単位の減算を実行し、 result のテンソル。要素のタイプに応じて、次の操作を行います。

  • 整数の場合: 整数の減算。
  • 浮動小数点数: IEEE-754 の subtraction
  • 複素数の場合: 複素減算。
  • 量子化型の場合: <ph type="x-smartling-placeholder">
      </ph>
    • dequantize_op_quantize(subtract, lhs, rhs, type(result))

入力

ラベル 名前 タイプ 制約
(I1) lhs 整数、浮動小数点、複合型のテンソル、またはテンソルごとの量子化テンソル (C1)
(I2) rhs 整数、浮動小数点、複合型のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 整数、浮動小数点、複合型のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

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

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

その他の例

tan

セマンティクス

operand テンソルに対して要素単位のタンジェント演算を実行し、 result テンソル。要素のタイプに応じて、次の操作を行います。

  • 浮動小数点数: IEEE-754 の tan
  • 複素数の場合: 複素正接。
  • 量子化型の場合: dequantize_op_quantize(tan, operand, type(result))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

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

// %operand: [
//            [0.0, 1.57079632],       // [0, pi/2]
//            [3.14159265, 4.71238898] // [pi, 3pi/2]
//           ]
%result = "stablehlo.tan"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
// %result: [
//           [0.0, 1.63312e+16],
//           [0.0, 5.44375e+15]
//          ]

その他の例

tanh

セマンティクス

operand のテンソルに対して要素単位の双曲線正接演算を行います。 result テンソルを生成します。要素のタイプに応じて、次の操作を行います。

  • 浮動小数点数: IEEE-754 の tanh
  • 複素数の場合: 複素双曲線正接。
  • 量子化型の場合: <ph type="x-smartling-placeholder">
      </ph>
    • dequantize_op_quantize(tanh, operand, type(result))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

出力

名前 タイプ 制約
result 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

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

// %operand: [-1.0, 0.0, 1.0]
%result = "stablehlo.tanh"(%operand) : (tensor<3xf32>) -> tensor<3xf32>
// %result: [-0.76159416, 0.0, 0.76159416]

その他の例

行 / 列の入れ替え

セマンティクス

permutation を使用して operand テンソルの次元を並べ替え、 result のテンソル。よりフォーマルな表現で、result[result_index] = operand[operand_index] ここで result_index[d] = operand_index[permutation[d]] です。

入力

ラベル 名前 タイプ 制約
(I1) operand テンソルまたは量子化テンソル (C1 ~ C4)
(I2) permutation si64 型の 1 次元のテンソル定数 (C2-C4)

出力

名前 タイプ 制約
result テンソルまたは量子化テンソル (C1)、(C3 ~ C4)

制約

  • (C1)element_type(result) は次式で求められます。 <ph type="x-smartling-placeholder">
      </ph>
    • element_type(operand)!is_per_axis_quantized(operand) の場合)。
    • element_type(operand)(ただし、quantization_dimension(operand) と そうでない場合は、quantization_dimension(result) が異なる場合があります。
  • (C2)permutationrange(rank(operand)) の順列です。
  • (C3)shape(result) = dim(operand, permutation...)
  • (C4)is_per_axis_quantized(result) の場合、 quantization_dimension(operand) = permutation(quantization_dimension(result))

// %operand: [
//            [[1,2], [3,4], [5,6]],
//            [[7,8], [9,10], [11,12]]
//           ]
%result = "stablehlo.transpose"(%operand) {
  permutation = array<i64: 2, 1, 0>
} : (tensor<2x3x2xi32>) -> tensor<2x3x2xi32>
// %result: [
//           [[1,7], [3,9], [5,11]],
//           [[2,8], [4,10], [6,12]]
//          ]

その他の例

triangular_solve

セマンティクス

下三角形または上三角形を含む連立一次方程式をバッチで解く 係数行列です。

より正式には、ab を考慮して、result[i0, ..., iR-3, :, :] が解答となります。 left_side が次の場合に op(a[i0, ..., iR-3, :, :]) * x = b[i0, ..., iR-3, :, :] に変更: true または x * op(a[i0, ..., iR-3, :, :]) = b[i0, ..., iR-3, :, :] left_sidefalse であり、op(a) が決定される変数 x を解きます。 transpose_a で指定。次のいずれかになります。

  • NO_TRANSPOSE: a をそのまま使用してオペレーションを実行します。
  • TRANSPOSE: a の転置でオペレーションを実行します。
  • ADJOINT: a の共役転置でオペレーションを実行します。

lowertrue の場合、または入力データは a の下側の三角形からのみ読み取られます それ以外の場合は、a の上三角形。出力データは、同じ三角形内に返されます。 もう一方の三角形の値は実装定義です

unit_diagonal が true の場合、実装は対角線と a の要素が 1 と等しくない場合、動作は未定義です。

量子化型の場合 dequantize_op_quantize(lambda x, y: triangular_solve(x, y, left_side, lower, unit_diagonal, transpose_a), a, b, type(result))

入力

ラベル 名前 タイプ 制約
(I1) a 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1 ~ C3)
(I2) b 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1 ~ C4)
(I3) left_side i1 型の定数 (C3)
(I4) lower i1 型の定数
(I5) unit_diagonal i1 型の定数
(I6) transpose_a NO_TRANSPOSETRANSPOSEADJOINT の列挙型

出力

名前 タイプ 制約
result 浮動小数点数または複素数のテンソル、またはテンソルごとの量子化テンソル (C1)

制約

  • (C1)baseline_element_type(a) = baseline_element_type(b)
  • (C2)2 <= rank(a) = rank(b) = R
  • (C3)shape(a)shape(b) の関係は次のように定義される。 <ph type="x-smartling-placeholder">
      </ph>
    • shape(a)[:-3] = shape(b)[:-3]
    • dim(a, -2) = dim(a, -1) = dim(b, left_side ? -2 : -1)
  • (C4)baseline_type(b) = baseline_type(result)

// %a = [
//       [1.0, 0.0, 0.0],
//       [2.0, 4.0, 0.0],
//       [3.0, 5.0, 6.0]
//      ]
// %b = [
//       [2.0, 0.0, 0.0],
//       [4.0, 8.0, 0.0],
//       [6.0, 10.0, 12.0]
//      ]
%result = "stablehlo.triangular_solve"(%a, %b) {
  left_side = true,
  lower = true,
  unit_diagonal = false,
  transpose_a = #stablehlo<transpose NO_TRANSPOSE>
} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
// %result: [
//           [2.0, 0.0, 0.0],
//           [0.0, 2.0, 0.0],
//           [0.0, 0.0, 2.0]
//          ]

tuple

<ph type="x-smartling-placeholder">
</ph>

セマンティクス

val から result タプルを生成します。

入力

ラベル 名前 タイプ 制約
(I1) val 値の可変数 (C1)

出力

名前 タイプ 制約
result tuple (C1)

制約

  • (C1)result の型は tuple<E0, ..., EN-1>Ei = type(val[i]))です。

// %val0: [1.0, 2.0]
// %val1: (3)
%result = "stablehlo.tuple"(%val0, %val1) : (tensor<2xf32>, tuple<tensor<i32>>) -> tuple<tensor<2xf32>, tuple<tensor<i32>>>
// %result: ([1.0, 2.0], (3))

その他の例

uniform_dequantize

セマンティクス

量子化テンソル operand を要素ごとに変換します。 定義された量子化パラメータに基づく浮動小数点テンソル result operand 型。

(より正式には、result = dequantize(operand))。

入力

ラベル 名前 タイプ 制約
(I1) operand 量子化テンソル (C1)、(C2)

出力

名前 タイプ 制約
result 浮動小数点型のテンソル (C1)、(C2)

制約

  • (C1)shape(operand) = shape(result)
  • (C2)element_type(result) = expressed_type(operand)

// %operand: [10, 10]
%result = "stablehlo.uniform_dequantize"(%operand) : (tensor<2x!quant.uniform<i8:f32:0, {0.1:-30,0.5:-20}>>) -> tensor<2xf32>
// %result: [4.0, 15.0]

uniform_quantize

セマンティクス

浮動小数点テンソルまたは量子化テンソルを要素単位で変換します operand は、量子化に従って量子化テンソル result に変換します。 result 型で定義されたパラメータ。

よりフォーマルな表現では

  • is_float(operand) の場合: <ph type="x-smartling-placeholder">
      </ph>
    • result = quantize(operand, type(result))
  • is_quantized(operand) の場合: <ph type="x-smartling-placeholder">
      </ph>
    • float_result = dequantize(operand)
    • result = quantize(float_result, type(result))

入力

ラベル 名前 タイプ 制約
(I1) operand 浮動小数点型または量子化型のテンソル (C1)、(C2)

出力

名前 タイプ 制約
result 量子化テンソル (C1)、(C2)

制約

  • (C1)shape(operand) = shape(result)
  • (C2)expressed_type(result) = is_float(operand) ? element_type(operand) : expressed_type(operand)

// %operand: [4.0, 15.0]
%result = "stablehlo.uniform_quantize"(%operand) : (tensor<2xf32>) -> tensor<2x!quant.uniform<i8:f32:0, {0.1:-30,0.5:-20}>>
// %result: [10, 10]

// %operand: [10, 10]
%result = "stablehlo.uniform_quantize"(%operand) : (tensor<2x!quant.uniform<i8:f32:0, {0.1:-30,0.5:-20}>>) -> tensor<2x!quant.uniform<i8:f32:0, {0.1:-20,0.2:-30}>>
// %result: [20, 45]

しばらく

セマンティクス

body 関数を 0 回以上実行すると出力が生成されますが、 cond 関数は true を出力します。より形式的には、セマンティクスは、 次のように Python 構文を使用します。

internal_state = operand
while cond(*internal_state):
  internal_state = body(*internal_state)
results = internal_state

無限ループの動作は未定 (#383)。

入力

ラベル 名前 タイプ 制約
(I1) operand テンソル、量子化テンソル、またはトークンの可変数 (C1 ~ C3)
(I2) cond 関数 (C1)
(I3) body 関数 (C2)

出力

名前 タイプ 制約
results テンソル、量子化テンソル、またはトークンの可変数 (C3)

制約

  • (C1)cond の型は (T0, ..., TN-1) -> tensor<i1> であり、 Ti = type(operand[i])
  • (C2)body の型は (T0, ..., TN-1) -> (T0, ..., TN-1) であり、 Ti = type(operand[i])
  • (C3)type(results...) = type(operand...)

// %init_i: 1
// %init_sum: 0
// %one: 1
// %ten: 10
%results0, %results1 = "stablehlo.while"(%init_i, %init_sum) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %cond = "stablehlo.compare"(%arg0, %ten) {
      comparison_direction = #stablehlo<comparison_direction LT>
    } : (tensor<i64>, tensor<i64>) -> tensor<i1>
    stablehlo.return %cond : tensor<i1>
  }, {
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %new_sum = stablehlo.add %arg1, %one : tensor<i64>
    %new_i = stablehlo.add %arg0, %one : tensor<i64>
    stablehlo.return %new_i, %new_sum : tensor<i64>, tensor<i64>
}) : (tensor<i64>, tensor<i64>) -> (tensor<i64>, tensor<i64>)
// %results0: 10
// %results1: 10

その他の例

XOR

セマンティクス

2 つのテンソル lhsrhs の要素単位の XOR を実行し、result を生成します。 テンソルです。要素のタイプに応じて、次の操作を行います。

  • ブール値の場合: 論理 XOR。
  • 整数の場合: ビット単位の XOR。

入力

ラベル 名前 タイプ 制約
(I1) lhs ブール型または整数型のテンソル (C1)
(I2) rhs ブール型または整数型のテンソル (C1)

出力

名前 タイプ 制約
result ブール型または整数型のテンソル (C1)

制約

  • (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.xor"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[4, 4], [4, 12]]

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

その他の例

言語の相互運用

現時点では、StableHLO プログラムには、 StableHLO では定義されません。

モジュール、関数、呼び出し、リターン

StableHLO は、ModuleOp、FuncOp、CallOp、および ReturnOp。これは、既存の MLIR 機構との相互運用性を高めるために行われました。 便利なパスには、FuncOp と ModuleOp をターゲットとして記述します。 パイプラインはこれらのオペレーションの存在を想定しています完全な互換性の保証は 適用されます。テスト環境でこれらのオペレーションに変化が 互換性のない方法(削除など)が維持されるように、StableHLO の同等の機能が追加されます。 サポートしています。

CHLO

CHLO opset には、StableHLO に分解される上位レベルのオペレーションが含まれています。 現在のところ、CHLO の互換性は保証されていません。互換性 chlo-legalize-to-stablehlo パスは、 使用する必要があります。

シェイプの操作

中核的なプラットフォームから特定のオペレーションを 動的な StableHLO プログラムでの MLIR 言語による形状計算。 最も一般的な例としては、shape 言語があります。 shape_ofnum_elementstensor 言語などの演算 dimfrom_elements などの演算と、組み込みの index 型があります。

Dynamism RFC >O2 これらは対象範囲外であることを示しますが、index 型の一部のサポートは 相互運用性を確保するために 含まれていますこれらに対する互換性は保証されません。 型があります。shape-legalize-to-stablehlo では、 pass を使用して、これらのオペレーションを完全にサポートされている StableHLO オペレーションに変換できます。

非推奨のオペレーション

いくつかの StableHLO オペレーションが MHLO これらは非推奨で StableHLO の終了を予定しています。これらについて詳しくは 削除については、StableHLO v1.0 Cleanup #2283 をご覧ください。 これらのサポート終了に関するトラッカーの問題は #2340 です。

これらのオペレーションは、次のカテゴリに分類されます。

  • 「HLO にはありません」カテゴリであり、当初は StableHLO オペレーションに StableHLO opset を使用した場合が、後にあまり適していないと判断されました。 broadcast さん、create_token さん、cross-replica-sum さん、dot さん、einsum さん、 torch_index_selectunary_einsum#3)。
  • 未使用のオペレーション - ある時点で有用だったかもしれませんが、オペレーションは あるいはこれらの Ops を使用するパイプラインが 不要になったようにリファクタリングします。これには、maptuple#598)、 get_tuple_elementrngcomplex の比較 #560、 および畳み込み window_reversal#1181)。

これらのオペレーションの一部は、 既存のオペレーション(broadcastcreate_tokencross-replica-sumdotunary_einsum)であり、既存の互換性期間後に削除されます。 (6 か月間)。削除がまだ検討されているものもあります(einsumget_tuple_elementmaprng torch_index_selecttuplecomplex 比較、window_reversal)。コミュニティからのフィードバックは保留中ですが、 これらのオペレーションは削除されるか、完全にサポートされて仕様に追加されるかのいずれかです。終了 これらの op の将来性は既知であり、6 か月間の互換性が保証されています。

実行

シーケンシャル実行

StableHLO プログラムを実行するには、main 関数に入力値を指定します。 出力値の計算に使用されます関数の出力値は次の式で計算されます。 対応する return 演算をルートとして演算のグラフを実行します。

実行順序は、次のとおりに整合している限り、実装で定義されている すなわち op が使用前に実行される場合です。StableHLO では、 副作用 op は 1 つのトークンを使用し、1 つのトークンを生成します(複数のトークンが after_all で 1 つのトークンに多重化されるため)、サイド 効果も Dataflow と整合していますたとえば、次のプログラムでは、 実行順序は %0%1%2return の 2 つと、 %1%0%2return

func.func @main() -> tensor<f64> {
  %0 = stablehlo.constant dense<1.0> : tensor<f64>
  %1 = stablehlo.constant dense<2.0> : tensor<f64>
  %2 = stablehlo.add %0, %1 : tensor<f64>
  return %2 : tensor<f64>
}

より正式には、StableHLO プロセスとは以下の組み合わせを指します。 1)StableHLO プログラム、2)オペレーションのステータス(まだ実行されていない、 (すでに実行されている)、3)プロセスが取り組んでいる中間値などです。 このプロセスは、main 関数への入力値から始まり、 オペレーションのステータスと中間値を更新するオペレーションのグラフと、 出力値で終了します。形式化は未定 (#484)。

並列実行

StableHLO のプログラムは並行して実行でき、2 次元プロセス グリッドに編成される どちらも ui32 型である num_partitions による num_replicas です。

StableHLO プロセス グリッドで、StableHLO の num_replicas * num_partitions 同時に実行されます。各プロセスには固有の process_id = (replica_id, partition_id)、ここで replica_idreplica_ids = range(num_replicas))と partition_ids = range(num_partitions)partition_id: 両方とも タイプ ui32

プロセスグリッドのサイズはプログラムごとに静的に 今後、StableHLO プログラムに明示的に含める予定です。 #650)と、位置 すべてのプロセスについて静的に把握されます。各プロセスには replica_id によるプロセス グリッド内の位置へのアクセスと、 partition_id オペレーション。

プロセスグリッド内では、プログラムはすべて同じ プログラム、複数データ」すべて異なる場合があり(複数のプログラム内で、 複数のデータなどの中間要素です。将来的には 並行する StableHLO プログラムを定義する他のイディオムのサポートを導入します。 GSPMD を含む(#619)。

プロセスグリッド内では、プロセスはほとんど互いに独立しています。 個別のオペレーション ステータスがあり、入力/中間/出力の値は個別に ほとんどの op はプロセス間で別々に実行され、 ただし、以下で説明する少数のグループ演算を除きます。

ほとんどの演算の実行には、同じ ID の値のみが使用されるため、 通常、これらの値を名前で参照しても、あいまいさはありません。 しかし、グループ演算のセマンティクスを記述する場合には不十分であり、 name@process_id という表記が作られ、値 name を参照します。 追跡できます(この点から、不適格な name は、 name@(replica_id(), partition_id()) の省略形とみなされます)。

プロセス全体の実行順序は実装で定義されているが、 ポイントツーポイント通信とグループ演算によって導入される 以下で説明します。

ポイントツーポイント通信

StableHLO プロセスは、 StableHLO チャンネル。チャネルはタイプの正の ID で表されます。 si64。さまざまなオペレーションを通じて、チャネルに値を送信し、 チャンネルから受信します

さらなる形式化(例:これらのチャンネル ID の取得元、 プログラムが認識するプロセスや同期の種類を 未定です (#484)。

ストリーミング通信

すべての StableHLO プロセスは、次の 2 つのストリーミング インターフェースにアクセスできます。

  • 読み取り可能なインフィード
  • 書き込み可能なアウトフィード

チャネルはプロセス間での通信に使用するため、 双方にプロセスがあるため インフィードとアウトフィードには 実装後に定義できます。

さらなる形式化(例:ストリーミング通信が実行に与える影響 それによってどのような同期が行われるのかは 未定です (#484)。

集団オペレーション

StableHLO には、all_gatherall_reduceall_to_allcollective_broadcastcollective_permutereduce_scatter。これらのオペレーションはすべて、StableHLO プロセスでプロセスを分割します。 グリッドを StableHLO プロセス グループに分割し、その内部で共同計算を実行する 各プロセス グループを、他のプロセス グループから独立して作成します。

各プロセス グループ内で、グループ オペレーションによって同期が開始される場合がある 障壁となります。さらなる形式化(例:このプロセスがいつ実行されるのかを プロセスがどのようにこのバリアにたどり着いたか、 そうでない場合はどうなるかは 未定です (#484)。

プロセス グループにパーティション間通信が含まれている場合、つまり、 異なるプロセス グループ内のプロセスが呼び出されて、 チャネルが必要であり、グループオペレーションで si64 型の正の channel_id。レプリカ間の通信は必要ない です。

集合演算によって実行される計算は個々の演算に固有 各オペレーションのセクションで説明します。ただし、 プロセスグリッドがプロセスグループに分割され、 このセクションで説明しますより正式には StableHLO では 4 つの戦略があります

cross_replica

各プロセス グループ内では、レプリカ間の通信のみが発生します。この 戦略は replica_groups (レプリカ ID のリストのリスト)を受け取り、 partition_ids による replica_groups のデカルト積。replica_groups 固有の要素を持ち、すべての replica_ids を対象とする必要があります。よりフォーマルな表現では、 Python の構文:

def cross_replica(replica_groups: List[List[ReplicaId]]) -> List[List[ProcessId]]:
  for replica_group in replica_groups:
    for partition_id in partition_ids:
      process_group = []
      for replica_id in replica_group:
        process_group.append((replica_id, partition_id))
      yield process_group

たとえば、replica_groups = [[0, 1], [2, 3]]num_partitions = 2 の場合、次のようになります。 cross_replica で生成されるログ [[(0, 0), (1, 0)], [(0, 1), (1, 1)], [(2, 0), (3, 0)], [(2, 1), (3, 1)]]

cross_partition

各プロセス グループ内では、パーティション間通信のみが発生します。この 戦略は partition_groups (パーティション ID のリストのリスト)を受け取ります。また、 partition_groups のデカルト積を replica_ids で計算します。 partition_groups は一意の要素を持ち、すべての partition_ids をカバーする必要があります。 より正式には、Python 構文を使用します。

def cross_partition(partition_groups: List[List[PartitionId]]) -> List[List[ProcessId]]:
  for partition_group in partition_groups:
    for replica_id in replica_ids:
      process_group = []
      for partition_id in partition_group:
        process_group.append((replica_id, partition_id))
      yield process_group

たとえば、partition_groups = [[0, 1]]num_replicas = 4 の場合、次のようになります。 cross_partition で生成されるログ [[(0, 0), (0, 1)], [(1, 0), (1, 1)], [(2, 0), (2, 1)], [(3, 0), (3, 1)]]

cross_replica_and_partition

レプリカ間通信とパーティション間通信の両方が、 プロセスグループです。この戦略では replica_groups を使用します。これは、 レプリカ ID を計算し、各 replica_group のデカルト積を partition_idsreplica_groups には固有の要素があり、すべてを対象とする必要があります replica_ids。より正式には、Python 構文を使用します。

def cross_replica_and_partition(replica_groups: List[List[ReplicaId]]) -> List[List[ProcessId]]:
  for replica_group in replica_groups:
    process_group = []
    for partition_id in partition_ids:
      for replica_id in replica_group:
        process_group.append((replica_id, partition_id))
    yield process_group

たとえば、replica_groups = [[0, 1], [2, 3]]num_partitions = 2 の場合、次のようになります。 cross_replica_and_partition で生成されるログ [[(0, 0), (1, 0), (0, 1), (1, 1)], [(2, 0), (3, 0), (2, 1), (3, 1)]]

flattened_ids

この戦略では、flattened_id_groups を受け取ります。これは「フラット化された」リストのリストです。 replica_id * num_partitions + partition_id という形式のプロセス ID と プロセス ID に変換されます。flattened_id_groups には一意の要素が必要です process_ids のすべてを対象とします。より正式には、Python 構文を使用します。

def flattened_ids(flattened_id_groups: List[List[ui32]]) -> List[List[ProcessId]]:
  for flattened_id_group in flattened_id_groups:
    process_group = []
    for flattened_id in flattened_id_group:
      replica_id = flattened_id // num_partitions
      partition_id = flattened_id % num_partitions
      process_group.append((replica_id, partition_id))
    yield process_group

たとえば flattened_id_groups = [[0, 1, 2, 3], [4, 5, 6, 7]] の場合、次のようになります。 num_replicas = 4num_partitions = 2flattened_ids によって生成されるデータ [[(0, 0), (0, 1), (1, 0), (1, 1)], [(2, 0), (2, 1), (3, 0), (3, 1)]]

精度

現時点では、StableHLO では数値の精度は保証されません。 今後変更される可能性があります (#1156)。

量子化オペレーションの実行セマンティクス

量子化された StableHLO 演算の解釈は、 ハードウェアの要件と機能について説明します。たとえば、ハードウェアによっては、 「逆量子化、浮動小数点演算の実行」を使用して、量子化演算を解釈する 最後に量子化を行います。説明します。それ以外の方は、すべてのアクションを 整数演算による計算です。そのため、このモデルの 量子化された StableHLO オペレーションは、特定の 説明します。ハイブリッド量子化の解釈 (#1575)は、 仕様で規定されているセマンティクスに従います( 1792)。

エラー

StableHLO プログラムは、 実行前に多数のクラスのエラーが除外されます。 ただし、次のようなエラー条件は引き続き発生する可能性があります。整数オーバーフローを通じて 明示的に示されていない限り、これらのエラーはすべて 実装定義の動作になりますが、 (#1157)。

浮動小数点の例外

このルールの例外として、StableHLO プログラムの浮動小数点例外は 動作が明確に定義されています。例外が発生するオペレーションは、 IEEE-754 規格(無効なオペレーション、ゼロ除算、オーバーフロー、アンダーフロー、 (標準で定義されているとおり)デフォルトの結果を生成し、 対応するステータス フラグを生成せずに実行を続行する。 標準からの raiseNoFlag 例外処理。標準以外の 演算(複雑な算術や特定の超越関数など)は、 使用されます。

形状の不一致

StableHLO は、動的に形成されるテンソルをサポートしています。ただし、図形は 1 対 1 の それ以外の場合の動作は未定義です。StableHLO は明示的に 実行時にテンソルが特定の形状であることをアサートできる op を提供します。 正しいコードを生成するのは、プロデューサーの責任です。

具体的な例として、以下のプログラムは有効です。ただし実行時には %arg0%arg1 の正確な図形は同じである必要があります。そうでない場合、 プログラムの動作が定義されていません。

func.func @foo(%arg0: tensor<?xi32>, %arg1: tensor<?xi32>) -> tensor<?xi32> {
    %0 = stablehlo.add %arg0, %arg1 : tensor<?xi32>
    return %0 : tensor<?xi32>
}

Notation

構文の説明について、このドキュメントでは EBNF の修正された ISO フレーバーを使用しています 構文(ISO/IEC 14977:1996Wikipedia)、 1)ルールは = ではなく ::= を使用して定義される。

2)連結が , ではなく並置を使用して表現されている。

セマンティクスを記述する場合(「型」、「定数」、「Ops」のセクション内) サポートによって拡張された Python 構文に基づく数式を使用しています を使用して配列オペレーションを簡潔に表現します。うまくいく 小さなコード スニペットですが、まれにありますが、 常に明示的に導入される標準 Python 構文を使用します。

数式

dot_general の例を使用して、数式の仕組みを見てみましょう 仕様。このオペレーションの制約は次のとおりです。 dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...)

この数式で使用されている名前のソースは、1)グローバル関数、 例: dim、2)対応するプログラム要素のメンバー定義。例: lhslhs_batching_dimensionsrhsrhs_batching_dimensions の入力 「Inputs」で定義されたセクション(dot_general

前述のように、この数式の構文は Python ベースで、 簡潔さを重視した拡張機能です。数式を理解するために、データを 通常の Python 構文に変換します。

A)これらの数式では、= を使用して等式を表すため、最初のステップは Python 構文を取得する場合、次のように === に置き換えます。 dim(lhs, lhs_batching_dimensions...) == dim(rhs, rhs_batching_dimensions...)

B)また、これらの数式はスカラー式を変換する省略記号(...)もサポートしている 変換します。簡単に言えば、f(xs...) は テンソル xs 内のスカラー x、スカラー f(x) を計算して、すべての値を返す テンソルの結果として出力する」という意味です。標準の Python 構文では、 この例の数式は次のようになります。 [dim(lhs, dim1) for dim1 in lhs_batching_dimensions] == [dim(rhs, dim2) for dim2 in rhs_batching_dimensions]

省略記号のおかげで、多くの場合、同じレベルの スカラーがありますしかし、中には、下位レベルのセミ非フォーマル start_indices[bi0, ..., :, ..., biN] 式のように使用できます。 gather 仕様をご覧ください。簡潔さを示すために、簡潔にするために 標準的な Python 言語に翻訳するための厳密な形式が定められており、 状況に応じて直感的に理解できることを願っています。 不明瞭な数式がある場合はお知らせください。 改善します。

また、数式では、あらゆる種類のリストを展開するために省略記号が使用されています。 テンソル、テンソルのリスト( テンソル数など)が含まれます。この領域では、正確な 形式(例: リストは StableHLO 型システムの一部ではない)と 直感的な理解に頼りましょう

C)最後に取り上げる注目に値する車両は、 あります。StableHLO opset は暗黙的なブロードキャストをサポートしていませんが、 簡潔さも重視されます。簡単に言うと、スカラーが テンソルが想定されるコンテキストで使用される場合、スカラーは以下のようにブロードキャストされます。 予測された形状です。

dot_general の例を続けるために、別の制約もあります。 0 <= lhs_batching_dimensions < rank(lhs)dot_general での定義に基づく lhs_batching_dimensions はテンソルですが、0rank(lhs) はスカラーです。暗黙的ブロードキャストを適用すると、式は [0, ..., 0] <= lhs_batching_dimensions < [rank(lhs), ..., rank(lhs)] になります。

特定の dot_general 演算に適用すると、この数式は ブール値のテンソルに変換されます。数式を制約として使用すると、 数式が true またはテンソルと評価された場合、制約が成り立ちます。 true 要素のみ。

名前

数式における語彙のスコープには、1)グローバル関数、2)メンバー定義、

3)ローカル定義グローバル関数のリストを以下に示します。リスト 要素定義の記述はプログラム要素によって異なります。 適用先:

  • オペレーションの場合、メンバー定義には「入力」で導入された名前が含まれます。および 出力できます。
  • それ以外の場合、メンバーの定義には、インフラストラクチャの 対応する EBNF 非終端関数にちなんで名付けられたプログラム要素です。ほとんど これらの構造部分の名前は、タイムスタンプを変換して Snake の場合の非終端名(例: IntegerLiteral => integer_literal など)が使用されますが、このプロセスで名前が省略されることがあります(例: QuantizationStorageType =>storage_type)が使用され、この場合、名前は 「Inputs」と同様に明示的に導入され、/「出力」稼働中のセクション 仕様。
  • また、メンバー定義には、リソースを参照する self を 対応するプログラム要素です。

数式を評価する際に、次の型の値を使用できます。 1)Value(実際の値。例: dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>、 その種類を常に把握している場合) 2)Placeholderlhsrhsresult などの将来の値。実際の値) 値がまだわかっておらず、型のみがわかっている場合) 3)Type(「タイプ」セクションで定義したタイプ) 4)Function(「関数」セクションで定義されているグローバル関数)。

コンテキストに応じて、名前は異なる値を参照する場合があります。さらに表示 具体的には「Semantics」セクション(および他のプログラムの同等のセクション) 要素など)がランタイム ロジックを定義するため、すべての入力を Value として使用できます。 対照的に、「制約」はの Ops(および同等のもの)のセクション "compile-time"つまり 通常はランタイムの前に実行されるものなので そのため、定数入力のみが Value として使用でき、他の入力は Placeholder としてのみ利用可能です。

名前 「Semantics」 [制約] 内
グローバル機能 Function Function
定数入力 Value Value
非定数入力 Value Placeholder
出力 Value Placeholder
ローカル定義 定義によって異なる 定義によって異なる

transpose オペレーションの例について考えてみましょう。

%result = "stablehlo.transpose"(%operand) {
  permutation = dense<[2, 1, 0]> : tensor<3xi64>
} : (tensor<2x3x2xi32>) -> tensor<2x3x2xi32>

このオペレーションでは、permutation は定数であるため、Value として使用できます。 意味と制約の両方において重要です。一方、operandresult は セマンティクスでは Value として使用できますが、制約では Placeholder としてのみ使用できます。

関数

型の構成

型の構築に使用できる関数はありません。代わりに、 通常はより簡潔であるため、型構文を使用します。例: function_type( [tensor_type([], E), tensor_type([], E)], [tensor_type([], E)]) ではなく (tensor<E>, tensor<E>) -> (tensor<E>)

型の関数

  • element_type はテンソル型と量子化テンソル型に対して定義され、 それぞれ TensorElementType または QuantizedTensorElementType を返します。 (対応する TensorType または QuantizedTensorType の部分)。
def element_type(x: Value | Placeholder | Type):
 if type(x) == TensorType:
    return tensor_element_type(x)
  if type(x) == QuantizedTensorType:
    return quantized_tensor_element_type(x)
  if type(x) is not Type:
    return element_type(type(x))
  • is_per_axis_quantized(x: Value | Placeholder | Type) -> Value はショートカットです (is_quantized(x) and quantization_dimension(x) is not None

  • is_per_tensor_quantized(x: Value | Placeholder | Type) -> Valueis_quantized(x) and quantization_dimension(x) is None のショートカット。

  • is_promotable(x: Type, y: Type) -> bool は、タイプ x を昇格できるかどうかをチェックします y と入力します。xyQuantizedTensorElementType の場合、プロモーション storage_type にのみ適用されます。このバージョンのプロモーションは リダクション計算のコンテキストで現在使用されているもの( RFC をご覧ください)。

def is_promotable(x: Type, y: Type) -> Value:
  is_same_type = (is_bool(x) and is_bool(y)) or
    (is_integer(x) and is_integer(y)) or (is_float(x) and is_float(y)) or
    (is_complex(x) and is_complex(y)) or
    (is_quantized(x) and is_quantized(y) and expressed_type(x) = expressed_type(y))

  if is_same_type == False:
    return False

  if is_integer(x) or is_float(x):
    return bitwidth(x) <= bitwidth(y)

  if is_complex(x):
    return bitwidth(element_type(x)) <= bitwidth(element_type(y))

  if is_quantized(x):
    return bitwidth(storage_type(x)) <= bitwidth(storage_type(y))

  return false
  • is_quantized(x: Value | Placeholder | Type) -> Value は次のショートカットのショートカットです: is_quantized_tensor_element_type(x)

  • is_type_name(x: Value | Placeholder | Type) -> Value。全ユーザー できます。たとえば、xFloatType の場合、is_float(x)true を返します。 x が値またはプレースホルダの場合、この関数は次のショートカットを使用します。 is_type_name(type(x))

  • max_value(x: Type) -> Value は、各要素の最大値を返します。 TensorElementTypexTensorElementType でない場合、None を返します。

  • min_value(x: Type) -> Value は、指定可能な値の最小値を返します。 TensorElementTypexTensorElementType でない場合、None を返します。

  • member_name(x: Value | Placeholder | Type) -> Any。すべてのメンバーが利用可能 すべての型の member_name の定義。例: tensor_element_type(x) 対応する TensorTypeTensorElementType 部分を返します。 x が値またはプレースホルダの場合、この関数は次のショートカットを使用します。 member_name(type(x))x が適切なメンバーを持つ型でない場合、または そのような型の値またはプレースホルダは、None を返します。

  • is_empty_algorithm(*args: Type) は、すべてのドット アルゴリズム フィールドが設定されているかどうかを確認します。 宛先: Noneドット アルゴリズムでは実装が定義されているため、これが必要です。 デフォルト値を指定するのは誤りです。

値の構成

  • operation_name(*xs: Value | Type) -> Value。すべてのオペレーションで使用できます。 たとえば、add(lhs, rhs)lhsrhs の 2 つのテンソル値を取り、 これらの入力で add 演算を評価した結果の出力を返します。 オペレーションによっては、broadcast_in_dim。出力の型は次のとおりです。 「負荷が高い」、すなわちオペレーションの評価に必要。この例では これらの型を引数として受け取ります。

値に対する関数

  • Python のすべての演算子と関数を使用できます。例:両方 定期購入 およびスライス テンソル、量子化テンソルへのインデックス付けに Python の表記が タプルです。

  • to_destination_type(x: Value, destination_type: Type) -> Value の定義日 テンソルを変換し、type(x) に基づいて x の変換値を返します。 destination_type は次のようになります。

def to_destination_type(x: Value, destination_type: Type) -> Value:
  if type(x) == destination_type:
    return x

  if is_quantized(destination_type):
    if is_quantized(type(x)):
      return quantize(x, destination_type)
    assert is_float(type(x))
    return quantize(x, destination_type)

  if is_quantized(type(x)):
    assert destination_type = expressed_type(type(x))
    return dequantize(type(x))

  return convert(x, destination_type)

convertuniform_quantize、および uniform_dequantize オペレーション(#1576)。 マージ後は、上記の関数は不要になり、オペレーション名を使用できます。 convert を選択してください。

  • is_nan(x: Value) -> Value はテンソルに基づいて定義され、次の場合に true を返します。 x のすべての要素は NaN になり、それ以外の場合は false になります。x がテンソルでない場合、 None を返します。

  • is_sorted(x: Value) -> Value はテンソルに基づいて定義され、次の場合に true を返します。 x の要素が、昇順に対して昇順で並べ替えられます インデックスの辞書順。それ以外の場合は falsex が テンソルは None を返します。

  • is_unique(x: Value) -> Value はテンソルに基づいて定義され、x の場合は true を返します。 重複する要素がないか、それ以外の場合は false がない。x がテンソルでない場合、 None を返します。

  • member_name(x: Value) -> Any は、すべてのメンバー定義に対して定義されています。 すべての値の member_name。たとえば、real_part(x)RealPart を返します。 対応するComplexConstantの部分ですx が 適切なメンバーが None を返します。

  • same(x: Value) -> Value はテンソルに基づいて定義され、次の場合に true を返します。 x の要素はすべて互いに等しいか、そうでない場合は false である。テンソルが には要素がありません。つまり、「すべてが互いに等しい」とみなされます。つまり、 関数は true を返します。x がテンソルでない場合、None を返します。

  • split(x: Value, num_results: Value, axis: Value) -> Value の定義日 テンソルに変換し、axis 軸に沿って xnum_results スライスを返します。 x がテンソルでも dim(x, axis) % num_results != 0 でもない場合は、None を返します。

  • is_defined_in_parent_scope(x: Value) -> Value は文字列で定義される x が同じスコープで定義されている関数の名前である場合、true を返します。 親関数として使用できます。

  • is_namespaced_op_name(x: Value) -> Value は文字列に対して定義され、 x が有効な op 名である(つまり、次の規則を遵守している場合)は true 式: [a-zA-Z][a-zA-Z0-9_]*([.][a-zA-Z0-9_$]+)+

形状の計算

  • axes(x: Value | Placeholder | Type) -> Value は次のショートカットのショートカットです: range(rank(x))

  • dim(x: Value | Placeholder | Type, axis: Value) -> Value は次のショートカットのショートカットです: shape(x)[axis]

  • dims(x: Value | Placeholder | Type, axes: List) -> List は次のショートカットのショートカットです: list(map(lambda axis: dim(x, axis), axes))

  • index_space(x: Value | Placeholder | Type) -> Value はテンソルで定義される そして、並べ替えられた対応する TensorTypesize(x) インデックスを返します。 辞書順の昇順。例: [0, ..., 0][0, ..., 1]、...、 shape(x) - 1x がテンソル型、量子化テンソル型、値でない場合 またはこれらの型のプレースホルダの場合は、None を返します。

  • rank(x: Value | Placeholder | Type) -> Value は次のショートカットのショートカットです: size(shape(x))

  • shape(x: Value | Placeholder | Type) -> Value は、"Functions" 」セクション(member_name 経由)

  • size(x: Value | Placeholder | Type) -> Value は次のショートカットのショートカットです: reduce(lambda x, y: x * y, shape(x))

量子化計算

  • def baseline_element_type(x: Value | Placeholder | Type) -> Typeelement_type(baseline_type(x)) のショートカット。

  • baseline_type はテンソル型と量子化テンソル型に対して定義され、 「ベースライン」に変換します。つまり、同じ形状で、 要素タイプの量子化パラメータがデフォルト値にリセットされます。これは、 テンソル型と量子化テンソル型の両方を比較するための便利な手法として使用 これは頻繁に必要になります。量子化型では 量子化パラメータ(shape)を無視する型の比較 storage_typeexpressed_typestorage_minstorage_maxquantization_dimension(軸ごとの量子化型)はすべて一致する必要がありますが、 scaleszero points は異なる場合があります。

def baseline_type(x: Value | Placeholder | Type) -> Type:
  if type(x) == TensorType:
    return x
  if type(x) == QuantizedTensorType:
    element_type = quantized_tensor_element_type(x)
    baseline_element_type = QuantizedTensorElementType(
      storage_type = storage_type(element_type),
      storage_min = storage_min(element_type),
      storage_max = storage_max(element_type),
      expressed_type = expressed_type(element_type),
      quantization_dimension = quantization_dimension(element_type),
      scales = [constant(1.0, expressed_type(element_type))] * dim(x, quantization_dimension(element_type)),
      zero_points = [constant(0, storage_type(element_type))] * dim(x, quantization_dimension(element_type)))
    return QuantizedTensorType(shape(x), baseline_element_type)
  if type(x) is not Type:
    return baseline_element_type(type(x))
  • dequantize は量子化テンソル型に基づいて定義され、 浮動小数点テンソル型もサポートしていますこれは、量子化された要素を変換することで行われます。 これはストレージ タイプの整数値を対応する ゼロ点とスケールを使用して表現された型の浮動小数点値 値を返します。
def compute_zero_points(quantized_type, result_type):
  if is_per_tensor_quantized(quantized_type):
    return broadcast_in_dim(constant(zero_point(quantized_type), storage_type(quantized_type)), [], result_type)
  if is_per_axis_quantized(quantized_type):
    for i in index_space(result_type):
      d = quantization_dimension(quantized_type)
      zero_points[i] = zero_points(quantized_type)[i[d]]
    return zero_points

def compute_scales(quantized_type, result_type):
  if is_per_tensor_quantized(quantized_type):
    return broadcast_in_dim(constant(scale(quantized_type), expressed_type(quantized_type)), [],
            type(result_type))
  if is_per_axis_quantized(quantized_type):
    for i in index_space(result_type):
      d = quantization_dimension(quantized_type)
      scales[i] = scales(quantized_type)[i[d]]
    return scales

def dequantize(x: Value) -> Value:
  assert is_quantized(x)
  x_storage = bitcast_convert(x, storage_type(x))
  x_storage_sub = x_storage - compute_zero_points(type(x), type(x_storage))
  x_expressed_sub = convert(x_storage_sub, expressed_type(x))
  return x_expressed_sub * compute_scales(type(x), type(x_expressed_sub))
  • quantize は浮動小数点テンソル型で定義され、 数値化テンソル型です。これは、浮動小数点値を変換することで行われます。 表現された型の値をストレージ型の対応する整数値に変換する 量子化された要素型に関連付けられたゼロ点とスケールを使用します。
def quantize(x: Value, result_type: Type) -> Value:
  assert is_float(x) and is_quantized(result_type)
  zero_points = compute_zero_points(result_type, TensorType(shape(x), storage_type(result_type)))
  converted_zero_points = convert(zero_points, expressed_type(result_type))
  converted_min = convert(storage_min(result_type), expressed_type(result_type))
  converted_max = convert(storage_max(result_type), expressed_type(result_type))

  x_scaled = x / compute_scales(result_type, type(x))
  x_scaled_add_zp = x_scaled + converted_zero_points
  x_clamped = clamp(converted_min, x_scaled_add_zp, converted_max)
  x_rounded = round_nearest_even(x_clamped)
  return convert(x_rounded, result_type)
  • dequantize_op_quantize は、要素ごとの計算を指定するために使用されます。 数値化テンソルです逆量子化(量子化された要素を 演算を実行し、次に量子化します。つまり、 結果をストレージタイプに戻します現在のところ、この関数は テンソルごとの量子化に使用できます軸ごとの量子化は進行中です (#1574)。
def dequantize_op_quantize(op, *inputs_and_output_type):
  inputs = inputs_and_output_type[:-1]
  output_type = inputs_and_output_type[-1]

  float_inputs = map(dequantize, inputs)
  float_result = op(*float_inputs)
  return quantize(float_result, output_type)

def dequantize_batch_norm_grad_or_training_quantize(op, *inputs_and_output_types):
  inputs = inputs_and_output_type[:-3]
  float_inputs = map(dequantize, inputs)
  float_results = op(*float_inputs)
  return map(quantize, float_results, inputs_and_output_type[-3:])

def dequantize_compare(lhs, rhs, comparison_direction):
  float_lhs = dequantize(lhs)
  float_rhs = dequantize(rhs)
  return compare(float_lhs, float_rhs, comparison_direction, FLOAT)

def dequantize_select_quantize(pred, on_true, on_false, output_type):
  float_on_true = dequantize(on_true)
  float_on_false = dequantize(on_false)
  float_result = select(pred, float_on_true, float_on_false)
  return quantize(float_result, output_type)
  • hybrid_dequantize_then_op を使用して、ラベル値の重み付けのみの量子化を指定します。 浮動小数点の lhs と量子化された型の rhs を受け入れるハイブリッド演算。これは、 量子化された入力を表現された型に逆量子化し、計算を実行します 浮動小数点数で示されます。浮動小数点数 lhs テンソルの要素型と量子化 rhs の表現型 テンソルが同一でなければなりません。
def hybrid_dequantize_then_op(op, lhs, rhs):
  assert(is_float(lhs) and is_quantized(rhs) and element_type(lhs) == expressed_type(rhs))
  return op(lhs, dequantize(rhs))

グリッド計算

  • cross_partition(replica_groups: Value) -> Value。詳細については、「cross_replica」 ご覧ください。

  • cross_replica(replica_groups: Value) -> Value。詳細については、「cross_replica」 ご覧ください。

  • cross_replica_and_partition(replica_groups: Value) -> Value。詳しくは、 &quot;cross_replica_and_partition&quot;ご覧ください。

  • flattened_ids(replica_groups: Value) -> Value。「flattened_ids」については ご覧ください。

ダイナミズム

StableHLO 値には動的なディメンション サイズを設定できます。たとえば、tensor<?xi64>。 ただし、StableHLO 値に動的なディメンション数(ランクなし)を指定することはできません。 ダイナミズム(例:tensor<*xi64>)。オペランドと結果には動的な名前を使用できます。 サイズに制限がある場合でも同じです。制約は、 可能であれば静的に検証されます。それ以外の場合はランタイムに延期され、 不一致は未定義の動作になります例については以下をご覧ください。

単項要素演算の形状の不一致

次のおもちゃプログラムについて考えてみましょう。

func.func @foo(%arg0: tensor<?xf64>) {
  %0 = stablehlo.abs %arg0 : (tensor<?xf64>) -> tensor<2xf64>
  return
}

このようなプログラムは異常です。というのも、関数の形が 入力の形状ではありません。ただし、これは有効な StableHLO です。 。この構文では、abs オペレーションを静的に検証することはできません。 オペランドの正確な形状が不明であるためです。図形は 互換性があり、これは静的にチェックできます。? という結果が ランタイム時に 2 に設定されていても問題はありません。ただし、? は次の可能性がありました 他の整数になると、動作は未定義になります。

結果でディメンションのサイズが動的である場合、 未定義の動作です。確かに、「期待される」グループ化されたアセットが 不一致が生じることもあります

バイナリ要素ごとのオペレーションの形状の不一致

次のおもちゃプログラムについて考えてみましょう。

func.func @foo(%arg0: tensor<?xf64>, %arg1: tensor<?xf64>) {
  %0 = stablehlo.add %arg0, %arg0 : (tensor<?xf64>, tensor<?xf64>) -> tensor<?xf64>
  return
}

2 項要素演算に関しては、入力と出力の形状が 実行時に一致する必要があります。コンパイル時、静的ディメンションは同じである必要があります。 そうでない場合は互換性があれば十分です 入力に動的なディメンションがある場合、未定義のディメンションが存在する可能性がある サイズが一致しない場合があるためです。 サイズ(静的または動的)を指定します。すべての入力が 静的である場合、結果が動的かどうかは重要ではありません。つまり、 既知のディメンションは静的にチェックされ、動的ディメンションでは 制約が課せられます

出力の形状をオペランドとして受け取る op の形状の不一致

次のおもちゃプログラムについて考えてみましょう。

func.func @foo(%arg0: tensor<2xi32>) {
  %0 = stablehlo.dynamic_iota %arg0, dim = 0 : (tensor<2xi32>) -> tensor<3x4xi64>
  return
}

実行時のシェイプ オペランドの値は、結果のシェイプと一致する必要があります。 それ以外の場合の動作は未定義です。つまり、実行時に %arg0 には以下が必要です。 値 dense<[3, 4]> : tensor<2xi32> の値。シェイプ オペランドが定数の場合、 静的に検証できます。結果の形状が完全な動的な場合は、 不一致は指定できません