StableHLO 規格

StableHLO 是機器內高階作業 (HLO) 的一組作業 學習 (ML) 模型StableHLO 是不同網路之間的可攜性層 機器學習架構和機器學習編譯器：產生 StableHLO 程式的機器學習架構 與使用 StableHLO 程式的機器學習編譯器相容。

我們的目標是打造更多 各種機器學習架構 (如 TensorFlow、JAX) 和 PyTorch) 和機器學習編譯器 (例如 XLA 和 IREE)。為達成這些目標 文件提供 StableHLO 程式設計語言規格。

這個規格包含三個主要部分首先， 程式一節說明 StableHLO 程式的結構 由 StableHLO 函式組成，本身由 StableHLO 運算組成。 在這個結構中，「Ops」區段會指定 個別作業「執行」部分提供所有物件的語意 這些運算會在程式內一起執行最後， 標記法一節討論在 規格。

如要查看先前 StableHLO 版本的規格，請在 例如 StableHLO v0.19.0 規格。 如要查看每個子版本增加 StableHLO 所帶來的變更，請參閱 VhloDialect.td 中的版本記錄，

程式

Program ::= {Func}

StableHLO 程式包含任意數量的 StableHLO 函式。 以下程式範例包含函式 @main，其中包含 3 個輸入內容 (%image、%weights 和 %bias) 和 1 項輸出內容。函式的主體 有 6 次操作

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

函式

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

StableHLO 函式 (也稱為「已命名函式」) 具有 ID、輸入內容/輸出內容和主體日後，我們將 為函式引入額外中繼資料，提升相容性 與 HLO 結合 (#425、 #626、 #740、 #744)。

ID

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

StableHLO ID 在許多程式設計中都類似於 ID 其中兩項特性：1) 所有 ID 都有線索 例如區分不同 ID 或 2) 值 ID ，藉此簡化 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} | '?'

張量類型表示張量，即多維陣列。他們的 shape 和元素類型：其中形狀代表非負數或 維度大小 (以相應值的遞增順序排列) 維度 (又稱為「axes」)，從 0 到 R-1。 維度數量「R」稱為「排名」。例如，tensor<2x3xf32> 是 形狀為 2x3 且元素類型為 f32 的張量類型。包含兩種維度 (也就是兩個軸) - 第 0 個維度和第 1 個維度，其大小 分別是 2 和 3。排名為 2。

形狀可以部分或完全不明 (動態)，例如：tensor<?x2xf64> 只有部分不明，tensor<?x?xf64> 完全不明。動態 尺寸會以 ? 表示。不得將形狀取消排名。

在未來，我們預計將更多的張量類型 維度大小和元素類型，例如： (#629) 及稀疏度 (#1078)。

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

量化元素類型代表儲存空間類型 從 storage_min 到 storage_max (包含)，對應到 表達類型的浮點值。若是指定整數值 i， 對應的浮點值 f 可以按 f = (i - zero_point) * scale，其中呼叫 scale 和 zero_point 量化參數。storage_min 和 storage_max 為選用項目 但該樣式的預設值是 min_value(storage_type) 和 max_value(storage_type)。量化元素類型具有 下列限制：

• (C1) type(storage_min) = storage_type。
• (C2) type(storage_max) = storage_type。
• (C3) min_value(storage_type) <= storage_min < storage_max <= max_value(storage_type)。
• (C4) type(scales...) = expressed_type。
• (C5) 0 < scales。
• (C6) is_finite(scales...)。
• (C7) storage_min <= zero_points <= storage_max。
• (C8) type(zero_points...) = storage_type。
• (C9) size(scales) = size(zero_points)。
• (C10) 如果值為 is_empty(quantization_dimension)，則值為 size(scales) = 1。
• (C11) 0 <= quantization_dimension。

目前 QuantizationScale 是浮點常數，但 對整數比例尺度表示強烈興趣，以倍數表示 轉變。我們預計於近期內探索這項功能 (#1404)。

關於 QuantizationZeroPoint 語意，我們目前正在進行討論， 包括類型、值，以及是否只能 量化張量類型中可能有多個零點根據 這次討論的結果，零點的規格可能會 (#1405)。

另一個進行中的討論涉及 QuantizationStorageMin 的語意 和 QuantizationStorageMax 判斷是否應套用任何限制條件 對這些值及量化張量值所施加的影響 (#1406)。

最後，我們打算探索代表未知規模和零的數值 類似於我們規劃如何探討未知 尺寸大小 (#1407)。

量化張量類型代表含量化元素的張量。這些 張量和一般張量相同，只不過其元素 有量化元素類型，而非一般元素類型

在量化張量中，量化可「每個張量」，也就是具有 整個張量一個 scale 和 zero_point，或者可以是每個軸。 意味著擁有多個 scales 和 zero_points，每個分層使用一組 特定維度 quantization_dimension。更正式，在張量 t 中 按軸量化時，會有 dim(t, quantization_dimension) 個配量 「quantization_dimension」的屬性值：t[:, ..., 0, ..., :], t[:, ..., 1, ..., :]， 依此類推。i 配量中的所有元素都會使用 scales[i] 和 zero_points[i] 做為 量化參數量化張量類型具有以下特性 限制：

• 根據張量量化：
  • 沒有其他限制。
• 按軸量化：
  • (C12) quantization_dimension < rank(self)。
  • (C13) dim(self, quantization_dimension) = size(scales)。

TokenType ::= 'token'

符記類型代表符記，即產生和使用的不透明值 來監控一些作業符記可用來設定執行作業的執行順序 ，相關說明請參閱「執行」一節。

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

元組類型代表元組，也就是異質清單。元組是舊版的 不只能與 HLO 相容。在 HLO 中，元組 來代表變異基因段輸入和輸出內容在 StableHLO 中，可變輸入和 StableHLO 預設支援輸出，而 StableHLO 中唯一的用途是 完全代表 HLO ABI，例如：T、tuple<T>和 tuple<tuple<T>> 可能會因特定因素而有重大差異 。我們預計日後會變更 HLO ABI 我們可以從 StableHLO 移除元組類型 (#598)。

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

元素類型代表張量類型的元素。與許多程式設計不同 這些類型不是 StableHLO 的第一類。也就是說 StableHLO 程式無法直接表示這些型別的值 (因此， 使用 0 維張量表示 T 類型的純量值 tensor<T> 類型的值)。

• 布林值類型代表布林值 true 和 false。
• 整數類型可以是帶正負號 (si) 或無正負號 (ui)，並具有 其中一個支援的位元寬度 (2、4、8、16、32 或 64)。 帶符號的 siN 類型代表從 -2^(N-1) 到 2^(N-1)-1 的整數值 包含不帶正負號的 uiN 類型代表 0 到 2^N-1 包含頭尾。
• 浮點類型可以是下列其中一項：
  • f8E4M3FN 和 f8E5M2 類型分別對應 FP8 格式的 E4M3 和 E5M2 編碼，如 適用於深度學習的 FP8 格式。
  • 與 E4M3 和 E5M2 對應的 f8E4M3FNUZ 和 f8E5M2FNUZ 類型 先前提過的 FP8 格式編碼 深層類神經網路的 8 位元數值格式。
  • 與 FP8 格式 E4M3 編碼對應的 f8E4M3B11FNUZ 類型 描述 深層類神經網路的混合 8 位元浮點 (HFP8) 訓練和推論。
  • 與上述 bfloat16 格式對應的 bf16 類型 BFloat16：在 Cloud TPU 上高效能的秘訣。
  • 分別對應的 f16、f32 和 f64 類型 binary16 (「半精度」)、binary32 (「單一精確度」) 和 binary64 (「雙精度」) 格式，如 IEEE 754 標準。
  • tf32 類型對應的是 TensorFloat32 格式 而且對 StableHLO 的支援有限。
• 複雜類型代表內含實際部分的複雜值 和相同元素類型的虛構部分。支援的複雜功能 類型為 complex<f32> (這兩個部分都是 f32 類型) 和 complex<f64> (這兩個部分都屬於 f64 類型)。

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

函式類型同時代表具名和匿名函式。他們擁有 輸入類型 (-> 左側的類型清單) 和輸出類型 (-> 右側的類型清單)。在許多程式設計中 函式類型為第一類別，但不在 StableHLO 中。

StringType ::= 'string'

字串類型代表位元組序列。與許多程式設計不同 字串型別不是 StableHLO 中的第一類別，只會用於 指定節目元素的靜態中繼資料。

作業

StableHLO 作業 (也稱為「作業」) 代表封閉組合 機器學習模型的高階作業如前所述 StableHLO 語法受到 MLIR 影響，不一定是 符合人體工學的替代方案，但最適合 StableHLO 達成 實現機器學習架構和機器學習編譯器之間的互通性。

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

StableHLO 作業 (也稱為「作業」) 具有名稱、 輸入/輸出及簽章名稱包含 stablehlo. 前置字串 「記憶」，用於識別其中一個支援的作業。請見下方說明 列出所有支援的作業

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

作業會耗用「輸入內容」並產生「輸出內容」。輸入內容會分類為 輸入值 (在執行期間計算)、輸入函數 (已提供) 靜態方式，因為在 StableHLO 函式中並不是一流的值) 以及 輸入屬性 (也以靜態方式提供)。輸入與輸出的種類 而相關運算的運用及產生方式就取決於其記憶。例如：add 運算會耗用 2 個輸入值，並產生 1 個輸出值。相較之下， select_and_scatter 運算會耗用 3 個輸入值、2 個輸入函式和 3 個輸入屬性。

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

輸入函式 (也稱為「匿名函式」) 非常 與已命名函式類似，但以下情況除外：1) 這些函式沒有 ID (因此 名為「anonymous」)，2) 都不會宣告輸出類型 (輸出類型是 透過函式中的 return 運算推論而得)。

輸入函式的語法包含目前未使用的部分 (請參閱 Unused 實際工作環境)，想要與 MLIR 相容。在 MLIR 中 還有較為籠統的「區域」概念可以有多個「區塊」 透過跳躍運算彼此串連在一起這些區塊的 ID 與 連結到 Unused 實際工作環境，以便彼此區別。 StableHLO 沒有跳躍運算，因此 MLIR 語法的對應部分為 未使用 (但仍然存在)。

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

輸入屬性具有名稱和一項支援的屬性值 常數。它們是指定節目靜態中繼資料的主要方式 元素。舉例來說，concatenate 運算使用 dimension 屬性來 也就是將輸入值串連起來的維度。同樣地 slice 運算使用多個屬性，例如 start_indices 和 limit_indices 可指定用來切割輸入值的邊界。

目前外部的 StableHLO 程式有時會包含屬性 而本文中則不會討論到這類主題。日後，我們將 無法將這些屬性吸收至 StableHLO 運算元，或禁止 出現在 StableHLO 程式中在此同時 屬性：

• layout (#629)。
• mhlo.frontend_attributes (#628)。
• mhlo.sharding (#619)。
• output_operand_aliases (#740)。
• 位置中繼資料 (#594)。

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

「運算簽名」包含所有輸入值的類型 ( -> 左側)，以及所有輸出值的類型 ( -> 右側的型別)。嚴格來說，輸入類型是 多餘和輸出類型則幾乎總是備援 大部分 StableHLO 運算 (輸出類型可從輸入推斷)。儘管如此 是專門為 StableHLO 語法的一部分，以便與 MLIR 相容。

以下是其記憶法為 select_and_scatter 的運算範例。會耗用 3 次 輸入值 (%operand、%source 和 %init_value)、2 個輸入函式 和 3 個輸入屬性 (window_dimensions、window_strides 和 padding)。 請注意，運算的簽名只包含輸入值的類型 (但並非以內嵌方式提供的輸入函式和屬性類型)。

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

常數

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

StableHLO 常數具有常值和類型，兩者共同代表 StableHLO 值。一般來說，類型是常數語法的一部分，除了 例如是不明確的時 (例如布林值常數一律為 i1 類型， 而整數常數則可有多種可能的類型)。

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

布林值常數代表布林值 true 和 false。布林值 常數的類型為 i1。

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

整數常數會透過使用小數或 十六進制標記法。其他二進位檔案或八進位數字。 整數常數具有下列限制：

• (C1) is_wellformed(integer_literal, integer_type)。

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

浮點常數會透過字串表示浮點值 使用小數或科學記號此外，十六進位標記法的 用來直接指定以 相應類型浮點常數有下列限制：

• (C1) 如果使用非十六進位標記法， is_wellformed(float_literal, float_type)。
• (C2) 如果使用十六進位標記法， size(hexadecimal_digits) = num_bits(float_type) / 4。

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

複雜常數會使用實際部分的清單表示複雜值 和虛部分 (第二次)。例如： (1.0, 0.0) : complex<f32> 代表 1.0 + 0.0i， (0.0, 1.0) : complex<f32> 代表 0.0 + 1.0i。這些元素的順序 部分資料儲存於記憶體中，則須由實作定義。複雜常數 具有下列限制：

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

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

Tensor 常數使用透過 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。這些元素在記憶體中的儲存順序為 您會瞭解自己的解決方案Tensor 常數有下列限制：

• (C1) has_syntax(tensor_literal, element_type(tensor_type))，其中：
  • has_syntax(element_literal: Syntax, element_type: Type) = is_wellformed(element_literal, type)。
  • has_syntax(tensor_literal: List, element_type: Type) = has_syntax(tensor_literal..., element_type)。
• (C2) has_shape(tensor_literal, shape(tensor_type))，其中：
  • has_shape(element_literal: Syntax, []) = true。
  • has_shape(tensor_literal: List, shape: List) = size(tensor_literal) = shape[0] and has_shape(tensor_literal..., shape[1:])。
  • 否則為 false。

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

量化張量常數會使用相同的 作為張量常數，且會將元素指定為 儲存空間類型量化張量常數有下列限制：

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

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

字串常值包含使用 ASCII 字元指定的位元組， 逸出序列。這些字元與編碼無關，因此 則是由實作定義字串常值屬於 string 類型。

作業數

腹部

語義學

對 operand 張執行元素的抽象運算，並產生 result 張量根據元素類型執行以下操作：

• 如果是帶正負號整數：整數模數。
• 浮點值：IEEE-754 的 abs。
• 針對複數：複數的模數。
• 量化類型：dequantize_op_quantize(abs, operand, type(result))。

輸入

輸出

限制

• (C1) shape(result) = shape(operand)。
• (C2) baseline_element_type(result) 定義為：
  • 如果 is_complex(operand)，則為 complex_element_type(element_type(operand))。
  • 否則為 baseline_element_type(operand)。

範例

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

更多範例

add

語義學

執行將兩個張量依元素加入 lhs 和 rhs 的結果，並產生 result 張量。根據元素類型執行以下操作：

• 布林值：邏輯 OR。
• 整數：加整數。
• 浮點值：IEEE-754 的 addition。
• 適用於複數：複雜加法。
• 量化類型：dequantize_op_quantize(add, lhs, rhs, type(result))。

輸入

輸出

限制

• 如果作業使用非量化張量：
  • (C1) type(lhs) = type(rhs) = type(result)。
• 如果作業使用量化張量：
  • (C2) is_quantized(lhs) and is_quantized(rhs) and is_quantized(result)。
  • (C3) storage_type(lhs) = storage_type(rhs) = storage_type(result)。
  • (C4) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)。
  • (C5) (is_per_axis_quantized(lhs) or is_per_axis_quantized(rhs)) = is_per_axis_quantized(result)。
  • (C6) 如果值為 is_per_axis_quantized(lhs)，則值為 quantization_dimension(lhs) = quantization_dimension(result)。
  • (C7) 如果值為 is_per_axis_quantized(rhs)，則值為 quantization_dimension(rhs) = quantization_dimension(result)。

範例

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

更多範例

after_all

語義學

確保產生 inputs 的作業會在任何要求前執行 處理仰賴 result 的作業。執行作業不會執行任何動作 它只是為了建立 result 到 inputs 的資料依附元件。

輸入

輸出

範例

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

更多範例

all_gather

語義學

在 StableHLO 程序網格中的每個程序群組中，將值串連起來 分別來自 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] 「process_group」的「receiver」。
• 所有產品的results...@process = concatenate(operands...@process, all_gather_dim) 「process_group」的「process」。

輸入

輸出

限制

• (C1) 0 <= all_gather_dim < rank(operands...)。
• (C2) is_unique(replica_groups)。
• (C3) size(replica_groups) 定義為：
  • 如果使用 cross_replica，則為 num_replicas。
  • 如果使用 cross_replica_and_partition，則為 num_replicas。
  • 如果使用 flattened_ids，則為 num_processes。
• (C4) 0 <= replica_groups < size(replica_groups)。
• (C5) 如果值為 use_global_device_ids = true，則值為 channel_id > 0。
• (C6) type(results...) = type(operands...)，但以下項目除外：
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1)。

範例

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

更多範例

all_reduce

語義學

在 StableHLO 程序網格中的每個處理程序群組中，套用縮減項目 函式 computation 設為每個程序的 operands 張量值 並產生 results 張量

這項作業會將 StableHLO 程序格線分割為 process_groups 定義如下：

• cross_replica(replica_groups) 如果 channel_id <= 0 and use_global_device_ids = false。
• cross_replica_and_partition(replica_groups) 如果 channel_id > 0 and use_global_device_ids = false。
• flattened_ids(replica_groups) 如果 channel_id > 0 and use_global_device_ids = true。

之後，在每個 process_group 內：

• results...@process[result_index] = exec(schedule) 代表一些二進位樹狀結構 schedule，其中：
  • exec(node) = computation(exec(node.left), exec(node.right))。
  • exec(leaf) = leaf.value。
• schedule 是實作定義的二進位樹狀結構，其順序為 週遊是 to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0]))。

輸入

輸出

限制

• (C1) is_unique(replica_groups)。
• (C2) size(replica_groups) 定義為：
  • 如果使用 cross_replica，則為 num_replicas。
  • 如果使用 cross_replica_and_partition，則為 num_replicas。
  • 如果使用 flattened_ids，則為 num_processes。
• (C3) 0 <= replica_groups < size(replica_groups)。
• (C4) 如果值為 use_global_device_ids = true，則值為 channel_id > 0。
• (C5) computation 採用 (tensor<E>, tensor<E>) -> (tensor<E>) 類型，其中 is_promotable(element_type(operand), E)。
• (C6) shape(results...) = shape(operands...)。
• (C7) element_type(results...) = E。

範例

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

更多範例

all_to_all

語義學

在 StableHLO 程序網格中的每個程序群組內， 沿著 split_dimension 的 operands 張量分為多個部分 將分散的元件串連起來 concat_dimension 並產生 results 張量。 這項作業會將 StableHLO 程序格線分割為 process_groups 定義如下：

• 如果 channel_id <= 0，則為 cross_replica(replica_groups)。
• 如果 channel_id > 0，則為 cross_partition(replica_groups)。

之後，在每個 process_group 內：

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

輸入

輸出

限制

• (C1) 0 <= split_dimension < rank(operands...)。
• (C2) dim(operands..., split_dimension) % split_count = 0。
• (C3) 0 <= concat_dimension < rank(operands...)。
• (C4) 0 < split_count。
• (C5) is_unique(replica_groups)。
• (C6) size(replica_groups) 的定義為：
  • 如果使用 cross_replica，則為 num_replicas。
  • 如果使用 cross_partition，則為 num_partitions。
• (C7) 0 <= replica_groups < size(replica_groups)。
• (C8) dim(replica_groups, 1) = split_count。
• (C9) type(results...) = type(operands...)，但 split_dimension != concat_dimension 除外：
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count。
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count。

範例

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

更多範例

和

語義學

執行 lhs 和 rhs 兩個張量的元素 AND 運算，然後產生 result 張量根據元素類型執行以下操作：

• 布林值：邏輯 AND。
• 整數：位元 AND。

輸入

輸出

限制

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

範例

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

更多範例

atan2

語義學

對 lhs 和 rhs 張執行元素的 atan2 運算，然後產生 result 張量。根據元素類型執行以下操作：

• 浮點值：IEEE-754 的 atan2。
• 複雜數字：複數的 atan2。
• 量化類型：dequantize_op_quantize(atan2, lhs, rhs, type(result))。

輸入

輸出

限制

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

範例

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

更多範例

batch_norm_grad

語義學

計算 batch_norm_training 反向傳播的多個輸入梯度 並產生 grad_operand、grad_scale 和 grad_offsetgrad_output 張量更正式的說法可以用分解的方式表示 使用 Python 語法的現有 StableHLO 作業：

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

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

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

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

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

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

  return grad_operand, grad_scale, grad_offset

如果是量化類型，請執行 dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, mean, variance, grad_output: batch_norm_grad(operand, scale, mean, variance, grad_output, epsilon, feature_index), operand, scale, mean, variance, grad_output, type(grad_operand), type(grad_scale), type(feature_index))。

輸入

輸出

限制

• (C1) 0 <= feature_index < rank(operand)。
• (C2) operand、scale、mean、variance、grad_output、grad_operand、 grad_scale 和 grad_offset 具有相同的 baseline_element_type。
• (C3) operand、grad_output 和 grad_operand 的形狀相同。
• (C4) scale、mean、variance、grad_scale 和 grad_offset 具有 我們可以看到
• (C5) size(scale) = dim(operand, feature_index)。

範例

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

batch_norm_inference

語義學

將 operand張量 (除了 feature_index 維度並產生 result 張量。更正式 作業可表示為現有 StableHLO 作業的分解 使用 Python 語法如下：

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

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

如果是量化類型，請執行 dequantize_op_quantize(lambda operand, scale, offset, mean, variance: batch_norm_inference(operand, scale, offset, mean, variance, epsilon, feature_index), operand, scale, offset, mean, variance, type(result))。

輸入

輸出

限制

• (C1) 0 <= feature_index < rank(operand)。
• (C2) operand、scale、offset、mean、variance 和 result 有 相同的 baseline_element_type。
• (C3) size(scale) = dim(operand, feature_index)。
• (C4) size(offset) = dim(operand, feature_index)。
• (C5) size(mean) = dim(operand, feature_index)。
• (C6) size(variance) = dim(operand, feature_index)。
• (C7) baseline_type(operand) = baseline_type(result)。

範例

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

batch_norm_training

語義學

計算所有維度 (feature_index 除外) 的平均值和變異數 維度，並將產生 output、batch_mean 的 operand 張量正規化 和 batch_var 張張量。更正式的說法是以 使用 Python 語法來分解現有 StableHLO 作業 如下：

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

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

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

如果是量化類型，請執行 dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, offset: batch_norm_training(operand, scale, offset, epsilon, feature_index), operand, scale, offset, type(output), type(batch_mean), type(batch_var))。

輸入

輸出

限制

• (C1) 0 <= feature_index < rank(operand)。
• (C2) operand、scale、offset、batch_mean、batch_var 和 output 有 相同的 baseline_element_type。
• (C3) size(scale) = dim(operand, feature_index)。
• (C4) size(offset) = dim(operand, feature_index)。
• (C5) size(batch_mean) = dim(operand, feature_index)。
• (C6) size(batch_var) = dim(operand, feature_index)。
• (C7) baseline_type(output) = baseline_type(operand)。

範例

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

bitcast_convert

語義學

對 operand 張量執行點陣圖作業，並產生 result 張量 其中部分 operand 的位元會以 result 張類型。

更正式，給予E = element_type(operand)、E' = element_type(result)， 和 R = rank(operand)：

• 如果是 num_bits(E') < num_bits(E)， bits(result[i0, ..., iR-1, :]) = bits(operand[i0, ..., iR-1])。
• 如果是 num_bits(E') > num_bits(E)， bits(result[i0, ..., iR-2]) = bits(operand[i0, ..., iR-2, :])。
• 如果是 num_bits(E') = num_bits(E)， bits(result[i0, ..., iR-1]) = bits(operand[i0, ..., iR-1])。

bits 會傳回指定值及其行為的記憶體內表示法 是實作的定義，因為張量的確切表示方式為 而元素類型的具體呈現方式 也會定義相同的實作方式

輸入

輸出

限制

• (C1) 指定 E = is_quantized(operand) ? storage_type(operand) : element_type(operand)、E' = is_quantized(result) ? storage_type(result) : element_type(result) 和 R = rank(operand)：
  • 如果是 num_bits(E') = num_bits(E)，shape(result) = shape(operand)。
  • 如果為 num_bits(E') < num_bits(E)：
  • rank(result) = R + 1。
  • 所有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]]。

輸入

輸出

限制

• (C1) element_type(result) 的提供者：
  • 如果 !is_per_axis_quantized(operand)，則為 element_type(operand)。
  • element_type(operand)，但quantization_dimension(operand)， scales(operand)，zero_points(operand)可能與 quantization_dimension(result)、scales(result) 和 zero_points(result) 否則會怎麼樣
• (C2) size(broadcast_dimensions) = rank(operand)。
• (C3) 0 <= broadcast_dimensions < rank(result)。
• (C4) is_unique(broadcast_dimensions)。
• (C5) 對於 axes(operand) 中的所有 d：
  • dim(operand, d) = 1或
  • dim(operand, d) = dim(result, broadcast_dimensions[d])。
• (C6) 如果 is_per_axis_quantized(result)：
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)]。
  • 如果值為 dim(operand, quantization_dimension(operand)) = 1，則 scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result)))。

範例

// %operand: [
//            [1, 2, 3]
//           ]
%result = "stablehlo.broadcast_in_dim"(%operand) {
  broadcast_dimensions = array<i64: 2, 1>
} : (tensor<1x3xi32>) -> tensor<2x3x2xi32>
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

更多範例

保護殼

語義學

這個輸出內容只會從 branches 執行一個函式，進而產生輸出內容 視 index 的值而定更正式，result = selected_branch() 其中：

• 如果 0 <= index < size(branches)，則為 selected_branch = branches[index]。
• 否則為 selected_branch = branches[-1]。

輸入

輸出

限制

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

範例

// %index: -1
// %result_branch0: [0, 0]
// %result_branch1: [1, 1]
%result0, %result1 = "stablehlo.case"(%index) ({
  "stablehlo.return"(%result_branch0, %result_branch0) : (tensor<2xi64>, tensor<2xi64>) -> ()
}, {
  "stablehlo.return"(%result_branch1, %result_branch1) : (tensor<2xi64>, tensor<2xi64>) -> ()
}) : (tensor<i32>) -> (tensor<2xi64>, tensor<2xi64>)
// %result0: [1, 1]
// %result1: [1, 1]

更多範例

cbrt

語義學

對 operand 張執行元素依據元素的立方根運算，然後產生 result 張量。根據元素類型執行以下操作：

• 浮點值：IEEE-754 的 rootn(x, 3)。
• 針對複數：複雜的立方根。
• 量化類型：dequantize_op_quantize(cbrt, operand, type(result))

輸入

輸出

限制

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

範例

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

更多範例

錫爾

語義學

執行 operand 張元素依據元素，產生 result 張量。 實作 IEEE-754 中的 roundToIntegralTowardPositive 作業 規格。如果是量化類型，請執行 dequantize_op_quantize(ceil, operand, type(result))。

輸入

輸出

限制

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

範例

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

更多範例

Cholesky

語義學

計算一批矩陣的 Cholesky 分解。

更正式，針對「index_space(result)」的所有i， result[i0, ..., iR-3, :, :] 是 Cholesky 分解法 a[i0, ..., iR-3, :, :]，可以是下三角函數 (如果 lower 是 true) 或上方三角形 (lower 是 false) 矩陣。 對角三角形的輸出值，即嚴格上三角形或 相對來說，嚴格的下限三角形則為實作定義。

如果 i 存在，且輸入矩陣不是 Hermitian 正確定值 ，則此行為為未定義。

如果是量化類型，請執行 dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result))。

輸入

輸出

限制

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

範例

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

限制取值範圍

語義學

將 operand 張量的每個元素限制在最小值和最大值之間 值並產生 result 張量。更正式，result[result_index] = minimum(maximum(operand[result_index], min_element), max_element)， 其中 min_element = rank(min) = 0 ? min[] : min[result_index]、 max_element = rank(max) = 0 ? max[] : max[result_index]。以量化類型來說 執行dequantize_op_quantize(clamp, min, operand, max, type(result))。

對複雜數字表示排序涉及出乎意料的語意 因此，我們預計日後將停止支援複數數字 用於這項作業 (#560)。

輸入

輸出

限制

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

範例

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

更多範例

collective_broadcast

語義學

在 StableHLO 程序網格中的每個程序群組內，將 從來源程序到目標程序的 operand 張量，並產生 result 張量。

這項作業會將 StableHLO 程序格線分割為 process_groups 定義如下：

• 如果 channel_id <= 0，則為 cross_replica(replica_groups)。
• 如果 channel_id > 0，則為 cross_partition(replica_groups)。

之後，result@process 是由以下公式表示：

• operand@process_groups[i, 0] (如果有 i 來進行程序 位置：process_groups[i]。
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) 反之。

輸入

輸出

限制

• (C1) is_unique(replica_groups)。
• (C2) 0 <= replica_groups < N，其中 N 定義為：
  • 如果使用 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)) 反之。

輸入

輸出

限制

• (C1) dim(source_target_pairs, 1) = 2。
• (C2) is_unique(source_target_pairs[:, 0])。
• (C3) is_unique(source_target_pairs[:, 1])。
• (C4) 0 <= source_target_pairs < N，其中 N 定義為：
  • 如果使用 cross_replica，則為 num_replicas。
  • 如果使用 cross_partition，則為 num_partitions。
• (C5) type(result) = type(operand)。

範例

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

更多範例

比較

語義學

執行依據lhsrhs comparison_direction 和 compare_type，並會產生 result 張量。

comparison_direction 和 compare_type 的值如下 語意：

布林值和整數元素類型：

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

針對具有 compare_type = FLOAT 的浮點元素類型，運算會實作 下列 IEEE-754 作業：

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

針對包含 compare_type = TOTALORDER 的浮點元素類型，運算 會結合以下項目的 totalOrder 和 compareQuietEqual 作業： IEEE-754。

如果是複雜的元素類型，(real, imag) 組合的字典式比較如下： 並根據所提供的 comparison_direction 和 compare_type 執行。 對複雜數字表示排序涉及出乎意料的語意 因此，我們預計日後將停止支援複數數字 當 comparison_direction 為 GE、GT、LE 或 LT 時 (#560)。

適用於量化類型。執行dequantize_compare(lhs, rhs, comparison_direction)。

輸入

輸出

限制

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs)。
• (C2) shape(lhs) = shape(rhs) = shape(result)。
• (C3) compare_type 定義為：
  • 如果 is_signed_integer(element_type(lhs))，則為 SIGNED。
  • 如果 is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs))，則為 UNSIGNED。
  • FLOAT 或 TOTALORDER (如果 is_float(element_type(lhs)))。
  • 如果 is_complex(element_type(lhs))，則為 FLOAT。

範例

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

更多範例

複雜

語義學

這個外掛程式能從一對實數，執行元素轉換到複雜值， 虛值 (lhs 和 rhs)，並產生 result 張量。

輸入

輸出

限制

• (C1) type(lhs) = type(rhs)。
• (C2) shape(result) = shape(lhs)。
• (C3) element_type(result) 採用 complex<E> 類型，其中 E = element_type(lhs)。

範例

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

更多範例

複合

語義學

封裝其他 StableHLO 作業 (由其他 StableHLO 作業組成) 的作業。 搭乘 inputs 和 composite_attributes，產生 results。 運算的語意是透過 decomposition 屬性實作。 composite 運算可以在不變更程式的情況下替換為分解 語意內嵌分解則未提供 運算語意，建議使用 custom_call。

version 欄位 (預設為 0) 用來表示複合的 語意異動

輸入

輸出

限制

• (C1) is_namespaced_op_name(name)
• (C2) is_defined_in_parent_scope(decomposition)
• (C3) types(inputs...) == input_types(decomposition)
• (C4) types(results...) == output_types(decomposition)

範例

%results = "stablehlo.composite"(%input0, %input1) {
  name = "my_namespace.my_op",
  composite_attributes = {
    my_attribute = "my_value"
  },
  decomposition = @my_op,
  version = 1 : i32
} : (tensor<f32>, tensor<f32>) -> tensor<f32>

更多範例

串連

語義學

按照指定的順序串接與 dimension 維度的 inputs 並產生 result 張量更正式 result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1]，其中：

1. id = d0 + ... + dk-1 + kd。
2. d等於 dimension，且 d0 是第 d 個維度大小 (共 inputs 個)。

輸入

輸出

限制

• (C1) same(element_type(inputs...))。
• (C2) same(shape(inputs...))，但 dim(inputs..., dimension) 除外。
• (C3) 0 < size(inputs)。
• (C4) 0 <= dimension < rank(inputs[0])。
• (C5) element_type(result) = element_type(inputs[0])。
• (C6) shape(result) = shape(inputs[0])，但以下情況除外：
  • dim(result, dimension) = dim(inputs[0], dimension) + ...。

範例

// %input0: [[1, 2], [3, 4], [5, 6]]
// %input1: [[7, 8]]
%result = "stablehlo.concatenate"(%input0, %input1) {
  dimension = 0 : i64
} : (tensor<3x2xi64>, tensor<1x2xi64>) -> tensor<4x2xi64>
// %result: [[1, 2], [3, 4], [5, 6], [7, 8]]

更多範例

常數

語義學

從常數 value 產生 output 張量。

輸入

輸出

限制

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

範例

%output = "stablehlo.constant"() {
  value = dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32>
} : () -> tensor<2x2xf32>
// %output: [[0.0, 1.0], [2.0, 3.0]]

更多範例

完成轉換

語義學

這個外掛程式能執行從某個元素類型到另一個元素類型的轉換 operand 張量並產生 result 張量。

針對 boolean-to-any-supported-type 的轉換，false 的值是 轉換成零，而 true 值會轉換成 1。適用對象 any-supported-type-to-boolean 的轉換，沒有值會轉換成 false，而非零值會轉換為 true。請參閱下文，瞭解操作方式 並處理較複雜的類型

轉換包含整數至整數、整數至浮點值 或 floating-point-to-floating-point (如果來源值無法確切) 目的地類型所代表的結果值 這種表示法否則，行為會待定。 (#180)。

針對涉及floating-point-to-integer的轉換，分數部分為 遭到截斷。如果無法以截斷值表示目的地類型， 行為是待定 (#180)。

如果轉換涉及複雜至複雜，請遵循 floating-point-to-floating-point的轉換次數，可用於轉換真實和 虛構部分

對於complex-to-any-other-type的轉換，以及complex-to-any-other-type轉換， 系統會忽略來源虛設值，或目的地假值為 。實際發生的轉換 浮點轉換

原則上，此運算可以表示去量化 (從 量化張量到一般張量，量化 (從一般張量轉換 張量和量化張量) 和重新量化 (在量化之間轉換 張量)，但目前我們有專門的作業 - 第一個用途為 uniform_dequantize，第一個用途則為 uniform_quantize 第二種用途日後，這兩項作業可能會合併 放入 convert (#1576)。

輸入

輸出

限制

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

範例

// %operand: [-1, 0, 1]
%result = "stablehlo.convert"(%operand) : (tensor<3xi64>) -> tensor<3xcomplex<f64>>
// %result: [(-1.0, 0.0), (0.0, 0.0), (1.0, 0.0)]

更多範例

卷積

語義學

計算 lhs 窗口與 rhs 分塊之間並產出的點積 result。下圖顯示如何計算 result 中的元素 lhs 和 rhs 使用具體範例。

更正式，請考慮以下有關 lhs 的輸入內容的重新解析方式 為了能夠表示 lhs 的視窗：

• lhs_window_dimensions = lhs_shape(dim(lhs, input_batch_dimension), dim(rhs, kernel_spatial_dimensions), dim(lhs, input_feature_dimension))。
• lhs_window_strides = lhs_shape(1, window_strides, 1)。
• lhs_padding = lhs_shape([0, 0], padding, [0, 0])。
• lhs_base_dilations = lhs_shape(1, lhs_dilation, 1)。
• lhs_window_dilations = lhs_shape(1, rhs_dilation, 1)。

這項修正會使用下列輔助函式：

• lhs_shape(n, hw, c) = permute([n] + hw + [c], [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension])。
• result_shape(n1, hw, c1) = permute([n1] + hw + [c1], [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension])。
• permute([j0, j1, ..., jR-1], permutation) = [i0, i1, ..., iR-1]，其中 j[d] = i[permutation[d]]。

如果 feature_group_count = 1 和 batch_group_count = 1，則針對所有 index_space(dim(result, output_spatial_dimensions...))的output_spatial_index， result[result_shape(:, output_spatial_index, :)] = dot_product，其中：

• padding_value = constant(0, element_type(lhs))。
• padded_lhs = pad(lhs, padding_value, lhs_padding[:, 0], lhs_padding[:, 1], lhs_base_dilations - 1)。
• lhs_window_start = lhs_shape(0, output_spatial_index, 0) * lhs_window_strides。
• lhs_window = slice(padded_lhs, lhs_window_start, lhs_window_start + lhs_window_dimensions, lhs_window_dilations)。
• reversed_lhs_window = reverse(lhs_window, [input_spatial_dimensions[dim] for dim in range(size(window_reversal)) if window_reversal[dim] = true])。 這項功能似乎未使用，因此我們計劃日後移除 (#1181)。
• dot_product = dot_general(reversed_lhs_window, rhs, lhs_batching_dimensions=[], lhs_contracting_dimensions=input_spatial_dimensions + [input_feature_dimension], rhs_batching_dimensions=[], rhs_contracting_dimensions=kernel_spatial_dimensions + [kernel_input_feature_dimension])。

如果為 feature_group_count > 1：

• lhses = split(lhs, feature_group_count, input_feature_dimension)。
• rhses = split(rhs, feature_group_count, kernel_output_feature_dimension)。
• results... = convolution(lhses..., rhses..., ..., feature_group_count=1, ...)。
• result = concatenate(results, output_feature_dimension)。

如果為 batch_group_count > 1：

• lhses = split(lhs, batch_group_count, input_batch_dimension)。
• rhses = split(rhs, batch_group_count, kernel_output_feature_dimension)。
• results... = convolution(lhses..., rhses..., ..., batch_group_count=1, ...)。
• result = concatenate(results, output_feature_dimension)。

如果是量化類型，請執行 dequantize_op_quantize( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs, type(result))。

如為混合量化類型，請執行 hybrid_dequantize_then_op( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs)。

輸入

輸出

限制

• (C1) N = rank(lhs) = rank(rhs)。
• (C2) size(window_strides) = N - 2。
• (C3) 0 < window_strides。
• (C4) shape(padding) = [N - 2, 2]。
• (C5) size(lhs_dilation) = N - 2。
• (C6) 0 < lhs_dilation。
• (C7) size(rhs_dilation) = N - 2。
• (C8) 0 < rhs_dilation。
• (C9) size(window_reversal) = N - 2。
• (C10) dim(lhs, input_batch_dimension) % batch_group_count = 0。
• (C11) dim(lhs, input_feature_dimension) % feature_group_count = 0。
• (C12) size(input_spatial_dimensions) = N - 2。
• (C13) 指定 input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]：
  • is_unique(input_dimensions)。
  • 0 <= input_dimensions < N。
• (C14) dim(rhs, kernel_input_feature_dimension) = dim(lhs, input_feature_dimension) / feature_group_count。
• (C15) dim(rhs, kernel_output_feature_dimension) % batch_group_count = 0。
• (C16) dim(rhs, kernel_output_feature_dimension) % feature_group_count = 0。
• (C17) size(kernel_spatial_dimensions) = N - 2。
• (C18) 指定 kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]：
  • is_unique(kernel_dimensions)。
  • 0 <= kernel_dimensions < N。
• (C19) size(output_spatial_dimensions) = N - 2。
• (C20) 指定 output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension] 時：
  • is_unique(output_dimensions)。
  • 0 <= output_dimensions < N。
• (C21) 0 < feature_group_count。
• (C22) 0 < batch_group_count。
• (C23) feature_group_count = 1 or batch_group_count = 1。
• (C24) size(precision_config) = 2。
• (C25) dim(result, result_dim) 定義為：
  • 如果 result_dim = output_batch_dimension，則為 dim(lhs, input_batch_dimension) / batch_group_count。
  • 如果 result_dim = output_feature_dimension，則為 dim(rhs, kernel_output_feature_dimension)。
  • 其他結果則是 num_windows，其中：
  • output_spatial_dimensions[spatial_dim] = result_dim。
  • lhs_dim = input_spatial_dimensions[spatial_dim]。
  • rhs_dim = kernel_spatial_dimensions[spatial_dim]。
  • dilated_input_shape[lhs_dim] = dim(lhs, lhs_dim) = 0 ? 0 : (dim(lhs, lhs_dim) - 1) * lhs_dilation[spatial_dim] + 1。
  • padded_input_shape[lhs_dim] = padding[spatial_dim, 0] + dilated_input_shape[lhs_dim] + padding[spatial_dim, 1]。
  • dilated_window_shape[lhs_dim] = dim(rhs, rhs_dim) = 0 ? 0 : (dim(rhs, rhs_dim) - 1) * rhs_dilation[spatial_dim] + 1。
  • is_empty_window[lhs_dim] = padded_input_shape[lhs_dim] = 0 || dilated_window_shape[lhs_dim] > padded_input_shape[lhs_dim]。
  • num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1。
• (C26) rank(result) = N。
• 如果作業使用非量化張量：
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result)。
• 如果作業使用量化張量：
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)。
  • (C29) 如果 is_per_axis_quantized(rhs)， 然後quantization_dimension(rhs) = kernel_output_feature_dimension。
  • (C30) 如果值為 is_per_axis_quantized(result)，則 quantization_dimension(result) = output_feature_dimension。
  • 如果為 is_quantized(lhs)：
  • (C31) storage_type(lhs) = storage_type(rhs)。
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)。
  • (C33) 如果值為 is_per_tensor_quantized(rhs)，則 is_per_tensor_quantized(result)。
  • 如果為 !is_quantized(lhs)：
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result)。

範例

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

更多範例

餘弦

語義學

對 operand 張執行元素的餘弦運算，並產生 result 張量。根據元素類型執行以下操作：

• 浮點值：IEEE-754 的 cos。
• 適用於複數：複數。
• 量化類型：dequantize_op_quantize(cosine, operand, type(result))。

輸入

輸出

限制

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

範例

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

更多範例

count_leading_zeros

語義學

根據元素計算 operand 中的開頭零位元數 並產生 result 張量

輸入

輸出

限制

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

範例

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

更多範例

custom_call

語義學

封裝實作定義的作業 call_target_name，其中包含 inputs 和 called_computations 並產生 results。has_side_effect, backend_config 和 api_version 可能會用於提供額外福利 實作定義的中繼資料

目前，這項作業包含一組非常組織的 中繼資料，反映其對應作業在 XLA 編譯器我們預計在日後整合這項中繼資料 (#741)。

輸入

輸出

範例

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

除號

語義學

執行被除數 lhs 和除數 rhs 張元素的元素除數，以及 會產生 result 張量。根據元素類型執行以下操作：

• 整數：產生代數商與任意 並捨棄小數部分
• 浮點值：IEEE-754 的 division。
• 適用於複數：複雜的除法。
• 量化類型：
  • dequantize_op_quantize(divide, lhs, rhs, type(result))。

輸入

輸出

限制

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

範例

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

更多範例

dot_general

語義學

在 lhs 和 rhs 切片之間計算內點積， result 張量。

更正式的 result[result_index] = dot_product，其中：

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

如果是量化類型，請執行 dequantize_op_quantize( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs, type(result))。

如為混合量化類型，請執行 hybrid_dequantize_then_op( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs)。

precision_config 會控制以下項目的速度和準確率的取捨 計算加速器後端運算作業可以是下列任一值 ( 我們無法明確指出這些列舉值的語意 #755)：

• DEFAULT：計算速度最快，但預估結果的準確度最低 。
• HIGH：計算速度較慢，但接近 。
• HIGHEST：計算速度最低，但最準確的近似值 。

DotAlgorithm 定義用於實作的演算法主要屬性 也會定義精確度。如果演算法屬性 欄位，則 precision_config 必須是 DEFAULT。DotAlgorithms 沒有預設值，因為預設的參數為實作 因此，所有點號演算法欄位可以設為 None，以指定 空白點演算法，因此會改用 precision_config 值。

DotAlgorithm 欄位包括：

• lhs_precision_type 和 rhs_precision_type，LHS 和 運算的 RHS 會四捨五入為。精確度類型與 儲存空間類型和輸出內容
• accumulation_type 用於累計的精確度。
• 「lhs_component_count」、「rhs_component_count」和「num_primitive_operations」 我們會將 LHS 和/或 RHS 分解為演算法 同時執行多個「原始」在這些 Pod 中 值 - 通常為了模擬較高的精確度 (例如 運用 bfloat16 人工智慧資料類型來進行高精確度運算： bf16_6x tf32_3x 等)。對於沒有分解的演算法，這些值 應設為 1。
• allow_imprecise_accumulation，用於指定是否以較低精確度進行累計 允許執行某些步驟 (例如：CUBLASLT_MATMUL_DESC_FAST_ACCUM)。

DotAlgorithm 屬性範例：

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


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


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

系統支援的組合完全取決於導入方式。於 一般而言，無法保證每種演算法都能支援 StableHLO 消費者的加速器類型。如果指定的演算法 錯誤，而不是復原為 。StableHLO 驗證會提供最適當的驗證 避免任何硬體並不支援的演算法。

查看「xla_data.proto > Algorithm」 部分支援的演算法值支援單 2483 需要製定計畫 是以後端為基礎的支援演算法集中查閱文件。

輸入

輸出

限制

• (C1) size(lhs_batching_dimensions) = size(rhs_batching_dimensions)。
• (C2) size(lhs_contracting_dimensions) = size(rhs_contracting_dimensions)。
• (C3) is_unique(lhs_batching_dimensions + lhs_contracting_dimensions)。
• (C4) is_unique(rhs_batching_dimensions + rhs_contracting_dimensions)。
• (C5) 0 <= lhs_batching_dimensions < rank(lhs)。
• (C6) 0 <= lhs_contracting_dimensions < rank(lhs)。
• (C7) 0 <= rhs_batching_dimensions < rank(rhs)。
• (C8) 0 <= rhs_contracting_dimensions < rank(rhs)。
• (C9) dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...)。
• (C10) dim(lhs, lhs_contracting_dimensions...) = dim(rhs, rhs_contracting_dimensions...)。
• (C11) size(precision_config) = 2。
• (C12) shape(result) = dim(lhs, lhs_batching_dimensions) + dim(lhs, lhs_result_dimensions) + dim(rhs, rhs_result_dimensions)。
• 如果作業使用非量化張量：
  • (C13) element_type(lhs) = element_type(rhs)。
• 如果作業使用量化張量：
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)。
  • (C15) zero_points(rhs) = 0。
  • (C16) 如果值為 is_per_axis_quantized(rhs)，則 quantization_dimension(rhs)不在rhs_contracting_dimensions中。
  • 如果為 is_quantized(lhs)：
  • (C17) storage_type(lhs) = storage_type(rhs)。
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)。
  • (C19) 如果值為 is_per_tensor_quantized(rhs)，則 is_per_tensor_quantized(result)。
  • 如果為 !is_quantized(lhs)：
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result)。
• 如果為 !is_empty_algorithm(lhs_precision_type, rhs_precision_type, accumulation_type, lhs_component_count, rhs_component_count, num_primitive_operations allow_imprecise_accumulation)：
  • (C21) precision_config... = DEFAULT。
  • (C22) 0 < lhs_component_count。
  • (C23) 0 < rhs_component_count。
  • (C24) 0 < num_primitive_operations。

範例

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

更多範例

dynamic_broadcast_in_dim

語義學

這項作業的功能與 broadcast_in_dim 運算，但結果形狀是透過 output_dimensions 動態指定。

這項作業也接受選用屬性 known_expanding_dimensions、known_non_expanding_dimensions 表達有關維度展開行為的靜態知識。 如果沒有指定，系統會假設所有維度都可能會擴大。

輸入

輸出

限制

• (C1) element_type(result) 的提供者：
  • 如果 !is_per_axis_quantized(operand)，則為 element_type(operand)。
  • element_type(operand)，但quantization_dimension(operand)， scales(operand)，zero_points(operand)可能與 quantization_dimension(result)、scales(result) 和 zero_points(result) 否則會怎麼樣
• (C2) size(broadcast_dimensions) = rank(operand)。
• (C3) 0 <= broadcast_dimensions < rank(result)。
• (C4) is_unique(broadcast_dimensions)。
• (C5) 對於 axes(operand) 中的所有 d：
  • dim(operand, d) = 1或
  • dim(operand, d) = dim(result, broadcast_dimensions[d])。
• (C6) 如果 is_per_axis_quantized(result)：
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)]。
  • 如果值為 dim(operand, quantization_dimension(operand)) = 1，則 scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result)))。
• (C7) size(output_dimensions) = rank(result)。
• (C8) is_unique(known_expanding_dimensions + known_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 動態指定。

輸入

輸出

限制

• (C1) N = rank(lhs) = rank(rhs)。
• (C2) size(window_strides) = N - 2。
• (C3) 0 < window_strides。
• (C4) shape(padding) = [N - 2, 2]。
• (C5) size(lhs_dilation) = N - 2。
• (C6) 0 < lhs_dilation。
• (C7) size(rhs_dilation) = N - 2。
• (C8) 0 < rhs_dilation。
• (C9) size(window_reversal) = N - 2。
• (C10) dim(lhs, input_batch_dimension) % batch_group_count = 0。
• (C11) dim(lhs, input_feature_dimension) % feature_group_count = 0。
• (C12) size(input_spatial_dimensions) = N - 2。
• (C13) 指定 input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]：
  • is_unique(input_dimensions)。
  • 0 <= input_dimensions < N。
• (C14) dim(rhs, kernel_input_feature_dimension) = dim(lhs, input_feature_dimension) / feature_group_count。
• (C15) dim(rhs, kernel_output_feature_dimension) % batch_group_count = 0。
• (C16) dim(rhs, kernel_output_feature_dimension) % feature_group_count = 0。
• (C17) size(kernel_spatial_dimensions) = N - 2。
• (C18) 指定 kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]：
  • is_unique(kernel_dimensions)。
  • 0 <= kernel_dimensions < N。
• (C19) size(output_spatial_dimensions) = N - 2。
• (C20) 指定 output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension] 時：
  • is_unique(output_dimensions)。
  • 0 <= output_dimensions < N。
• (C21) 0 < feature_group_count。
• (C22) 0 < batch_group_count。
• (C23) feature_group_count = 1 or batch_group_count = 1。
• (C24) size(precision_config) = 2。
• (C25) dim(result, result_dim) 定義為：
  • 如果 result_dim = output_batch_dimension，則為 dim(lhs, input_batch_dimension) / batch_group_count。
  • 如果 result_dim = output_feature_dimension，則為 dim(rhs, kernel_output_feature_dimension)。
  • 其他結果則是 num_windows，其中：
  • output_spatial_dimensions[spatial_dim] = result_dim。
  • lhs_dim = input_spatial_dimensions[spatial_dim]。
  • rhs_dim = kernel_spatial_dimensions[spatial_dim]。
  • dilated_input_shape[lhs_dim] = dim(lhs, lhs_dim) = 0 ? 0 : (dim(lhs, lhs_dim) - 1) * lhs_dilation[spatial_dim] + 1。
  • padded_input_shape[lhs_dim] = padding[spatial_dim, 0] + dilated_input_shape[lhs_dim] + padding[spatial_dim, 1]。
  • dilated_window_shape[lhs_dim] = dim(rhs, rhs_dim) = 0 ? 0 : (dim(rhs, rhs_dim) - 1) * rhs_dilation[spatial_dim] + 1。
  • is_empty_window[lhs_dim] = padded_input_shape[lhs_dim] = 0 || dilated_window_shape[lhs_dim] > padded_input_shape[lhs_dim]。
  • num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1。
• (C26) rank(result) = N。
• 如果作業使用非量化張量：
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result)。
• 如果作業使用量化張量：
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)。
  • (C29) 如果 is_per_axis_quantized(rhs)， 然後quantization_dimension(rhs) = kernel_output_feature_dimension。
  • (C30) 如果值為 is_per_axis_quantized(result)，則 quantization_dimension(result) = output_feature_dimension。
  • 如果為 is_quantized(lhs)：
  • (C31) storage_type(lhs) = storage_type(rhs)。
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)。
  • (C33) 如果值為 is_per_tensor_quantized(rhs)，則 is_per_tensor_quantized(result)。
  • 如果為 !is_quantized(lhs)：
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result)。

範例

// %lhs: [[
//        [[1], [2], [5], [6]],
//        [[3], [4], [7], [8]],
//        [[10], [11], [14], [15]],
//        [[12], [13], [16], [17]]
//      ]]
//
// %rhs: [
//         [[[1]], [[1]], [[1]]],
//         [[[1]], [[1]], [[1]]],
//         [[[1]], [[1]], [[1]]]
//        ]
// %padding: [[1, 1],
//            [1, 1]]
%result = "stablehlo.dynamic_conv"(%lhs, %rhs, %padding) {
  window_strides = array<i64: 4, 4>,
  lhs_dilation = array<i64: 2, 2>,
  rhs_dilation = array<i64: 1, 1>,
  window_reversal = array<i1: false, false>,
  dimension_numbers = #stablehlo.conv<raw
    input_batch_dimension = 0,
    input_feature_dimension = 3,
    input_spatial_dimensions = [0, 1],
    kernel_input_feature_dimension = 2,
    kernel_output_feature_dimension = 3,
    kernel_spatial_dimensions = [0, 1],
    output_batch_dimension = 0,
    output_feature_dimension = 3,
    output_spatial_dimensions = [1, 2]
  >,
  feature_group_count = 1 : i64,
  batch_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>, tensor<2x2xi64>) -> tensor<1x2x2x1xi64>
// %result: [[
//            [[1], [5]],
//            [[10], [14]]
//          ]]