Spesifikasi StableHLO

StabilHLO adalah operasi yang ditetapkan untuk operasi tingkat tinggi (HLO) di mesin machine learning (ML). StabilHLO berfungsi sebagai lapisan portabilitas antara berbagai Framework ML dan compiler ML: framework ML yang menghasilkan program StableHLO kompatibel dengan compiler ML yang menggunakan program StableHLO.

Tujuan kami adalah menyederhanakan dan mempercepat pengembangan ML dengan membuat lebih banyak interoperabilitas antara berbagai framework ML (seperti TensorFlow, JAX, PyTorch) dan compiler ML (seperti XLA dan IREE). Untuk mencapai tujuan tersebut, dokumen menyediakan spesifikasi untuk bahasa pemrograman StableHLO.

Spesifikasi ini memuat tiga bagian utama. Pertama, Bagian Programs menjelaskan struktur program StabilHLO yang terdiri dari fungsi StableHLO yang terdiri dari operasi StableHLO. Dalam struktur tersebut, bagian Ops menentukan semantik operasi individual. Bagian Execution menyediakan semantik untuk semua operasi ini yang dieksekusi bersama dalam sebuah program. Terakhir, Bagian Notasi membahas notasi yang digunakan dalam spesifikasi pendukung.

Untuk melihat spesifikasi dari rilis StableHLO sebelumnya, buka repo di rilis yang diberi tag minat. Misalnya, StableHLO v0.19.0 Spec. Untuk melihat perubahan yang terjadi di setiap bump versi minor StableHLO, lihat log versi di VhloDialect.td.

Program

Program ::= {Func}

Program SttableHLO terdiri dari sejumlah arbitrer fungsi StableHLO. Berikut adalah contoh program dengan fungsi @main yang memiliki 3 input (%image, %weights, dan %bias) dan 1 output. Isi fungsi memiliki 6 operasi.

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

Fungsi

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

Fungsi SttableHLO (yang juga disebut fungsi bernama) memiliki ID, input/{i>output<i} dan {i>body<i}. Di masa mendatang, kami berencana memperkenalkan metadata tambahan untuk fungsi guna mencapai kompatibilitas yang lebih baik dengan HLO (#425, #626, #740, #744).

Pengenal

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

ID SttableHLO mirip dengan ID dalam banyak pemrograman bahasa, dengan dua kekhasan: 1) semua pengenal memiliki {i>sigil<i} yang membedakan berbagai jenis pengenal, 2) pengenal nilai dapat lengkap bersifat numerik untuk menyederhanakan pembuatan program StableHLO.

Jenis

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

Jenis SttableHLO dikategorikan ke dalam jenis nilai (yang juga disebut jenis kelas satu) yang mewakili nilai StabilHLO dan jenis bukan nilai yang menggambarkan elemen program lainnya. Jenis StableHLO mirip dengan jenis pada banyak bahasa pemrograman, dengan keunikan utamanya adalah yang bersifat khusus domain yang memberikan hasil yang tidak biasa (mis. jenis skalar bukan jenis nilai).

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

Jenis tensor mewakili tensor, yaitu array multidimensi. Aset tersebut memiliki bentuk dan jenis elemen, dengan sebuah bentuk merepresentasikan bilangan non-negatif atau ukuran dimensi yang tidak diketahui dalam urutan menaik dimensi (yang juga disebut sumbu) yang diberi nomor dari 0 hingga R-1. Tujuan jumlah dimensi R disebut peringkat. Misalnya, tensor<2x3xf32> adalah jenis tensor dengan bentuk 2x3 dan jenis elemen f32. Model ini memiliki dua dimensi (atau, dengan kata lain, dua sumbu) - dimensi ke-0 dan dimensi ke-1 - yang ukurannya adalah 2 dan 3. Peringkatnya 2.

Bentuk bisa tidak diketahui sebagian atau seluruhnya (dinamis), mis. tensor<?x2xf64> tidak diketahui sebagian dan tensor<?x?xf64> sama sekali tidak diketahui. Dinamis ukuran dimensi ditampilkan menggunakan ?. Bentuk tidak boleh diberi peringkat.

Di masa mendatang, kami berencana mempelajari jenis tensor yang diperluas di luar ukuran dimensi dan jenis elemen, misalnya, untuk menyertakan tata letak (#629) dan ketersebaran (#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

Jenis elemen terkuantisasi mewakili nilai bilangan bulat dari jenis penyimpanan di rentang dari storage_min hingga storage_max (inklusif) yang sesuai dengan nilai floating point dari jenis yang dinyatakan. Untuk nilai bilangan bulat i yang diberikan, nilai floating point f yang sesuai dapat dikomputasi sebagai f = (i - zero_point) * scale, yang mana scale dan zero_point dipanggil parameter kuantisasi. storage_min dan storage_max bersifat opsional dalam tata bahasa, tetapi memiliki nilai default min_value(storage_type) dan max_value(storage_type). Jenis elemen terkuantisasi memiliki batasan berikut:

• (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) Jika is_empty(quantization_dimension), maka size(scales) = 1.
• (C11) 0 <= quantization_dimension.

Saat ini, QuantizationScale adalah konstanta floating point, tetapi ada kepentingan yang kuat dalam skala berbasis bilangan bulat, diwakili dengan pengganda dan {i>shift<i}. Kami berencana untuk menjelajahi fitur ini dalam waktu dekat (#1404).

Ada diskusi berkelanjutan tentang semantik QuantizationZeroPoint, termasuk jenisnya, nilai-nilainya, dan apakah mungkin hanya ada satu atau dan mungkin beberapa titik nol dalam tipe tensor terkuantisasi. Berdasarkan hasil diskusi ini, spesifikasi seputar titik nol dapat berubah pada masa mendatang (#1405).

Diskusi berkelanjutan lainnya melibatkan semantik QuantizationStorageMin dan QuantizationStorageMax untuk menentukan apakah ada batasan dikenakan pada nilai-nilai ini dan pada nilai-nilai tensor terkuantisasi (#1406).

Terakhir, kita berencana untuk merepresentasikan skala yang tidak diketahui dan nol poin, sama seperti perencanaan kita dalam mengeksplorasi cara untuk merepresentasikan ukuran dimensi (#1407).

Jenis tensor terkuantisasi mewakili tensor dengan elemen terkuantisasi. Ini tensor sama persis dengan tensor reguler, hanya saja memiliki jenis elemen terkuantisasi, bukan jenis elemen reguler.

Pada tensor terkuantisasi, kuantisasi bisa berupa per-tensor, artinya satu scale dan zero_point untuk seluruh tensor atau dapat berupa per-axis, artinya, memiliki beberapa scales dan zero_points, satu pasang per irisan dimensi tertentu quantization_dimension. Secara lebih formal, dalam tensor t dengan kuantisasi per sumbu, terdapat dim(t, quantization_dimension) irisan dari quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :], dll. Semua elemen dalam slice i menggunakan scales[i] dan zero_points[i] sebagai parameter kuantisasinya. Jenis tensor terkuantisasi memiliki batasan:

• Untuk kuantisasi per-tensor:
  • Tidak ada batasan tambahan.
• Untuk kuantisasi per sumbu:
  • (C12) quantization_dimension < rank(self).
  • (C13) dim(self, quantization_dimension) = size(scales).

TokenType ::= 'token'

Jenis token mewakili token, yaitu nilai buram yang dihasilkan dan dipakai oleh beberapa operasi. Token digunakan untuk menerapkan urutan eksekusi pada operasi seperti yang dijelaskan di bagian Eksekusi.

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

Jenis tuple mewakili tuple, yaitu daftar heterogen. Tuple adalah warisan yang hanya ada untuk kompatibilitas dengan HLO. Di HLO, tuple adalah digunakan untuk merepresentasikan input dan {i>output <i}variadic. Dalam StableHLO, input variadic dan didukung secara native, dan satu-satunya penggunaan tuple di StableHLO adalah untuk secara komprehensif merepresentasikan ABI HLO, mis. T, tuple<T>, dan tuple<tuple<T>> mungkin secara signifikan berbeda bergantung pada terlepas dari implementasi layanan. Di masa mendatang, kami berencana melakukan perubahan pada HLO ABI sehingga kami dapat menghapus jenis tuple dari 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'

Jenis elemen mewakili elemen jenis tensor. Tidak seperti di banyak pemrograman bahasa, jenis ini bukan kelas satu di StableHLO. Hal ini berarti bahwa Program StabilHLO tidak dapat langsung merepresentasikan nilai jenis ini (akibatnya, merepresentasikan nilai skalar jenis T dengan tensor 0 dimensi nilai dari jenis tensor<T>).

• Jenis boolean mewakili nilai boolean true dan false.
• Jenis bilangan bulat dapat ditandatangani (si) atau tidak ditandatangani (ui) dan memiliki salah satu lebar bit yang didukung (2, 4, 8, 16, 32, atau 64). Jenis siN bertanda tangan mewakili nilai bilangan bulat dari -2^(N-1) hingga 2^(N-1)-1 jenis uiN inklusif, dan tidak ditandatangani mewakili nilai bilangan bulat dari 0 hingga 2^N-1 inklusif.
• Jenis floating point dapat berupa salah satu dari hal berikut:
  • Jenis f8E4M3FN dan f8E5M2 yang sesuai dengan masing-masing Encoding E4M3 dan E5M2 untuk format FP8 yang dijelaskan di Format FP8 untuk Deep Learning.
  • Jenis f8E4M3FNUZ dan f8E5M2FNUZ yang sesuai dengan E4M3 dan E5M2 pengkodean format FP8 yang dijelaskan dalam Format Numerik 8-bit untuk Jaringan Neural Dalam.
  • Jenis f8E4M3B11FNUZ yang sesuai dengan encoding E4M3 untuk format FP8 dijelaskan dalam Pelatihan dan Inferensi Hybrid 8-bit Floating Point (HFP8) untuk Jaringan Neural Dalam.
  • Jenis bf16 yang sesuai dengan format bfloat16 yang dijelaskan di BFloat16: Rahasia performa tinggi pada Cloud TPU.
  • Jenis f16, f32, dan f64 yang sesuai masing-masing binary16 ("presisi setengah"), binary32 ("presisi tunggal") dan Format binary64 ("presisi ganda") yang dijelaskan di standar IEEE 754.
  • Jenis tf32 sesuai dengan format TensorFloat32 dan memiliki dukungan terbatas di StableHLO.
• Jenis kompleks merepresentasikan nilai kompleks yang memiliki bagian nyata dan bagian imajiner dari jenis elemen yang sama. Kompleks yang didukung adalah complex<f32> (kedua bagian dari jenis f32) dan complex<f64> (kedua bagian adalah jenis f64).

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

Jenis fungsi mewakili fungsi bernama dan anonim. Mereka memiliki jenis input (daftar jenis di sebelah kiri ->) dan jenis output (daftar jenis di sisi kanan ->). Dalam banyak pemrograman bahasa, jenis fungsi adalah class satu, tetapi tidak di StableHLO.

StringType ::= 'string'

Jenis string mewakili urutan byte. Tidak seperti di banyak pemrograman bahasa, jenis string bukan kelas satu di StableHLO dan hanya digunakan untuk menentukan {i>metadata<i} statis untuk elemen program.

Operasi

Operasi SttableHLO (yang juga disebut ops) mewakili kumpulan tertutup operasi tingkat tinggi dalam model machine learning. Seperti yang dibahas di atas, Sintaksis StableHLO sangat terinspirasi oleh MLIR, yang belum tentu merupakan ergonomis, tetapi mungkin lebih cocok untuk tujuan StableHLO sehingga meningkatkan interoperabilitas antara framework ML dan compiler ML.

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

Operasi SttableHLO (yang juga disebut ops) memiliki nama, input/{i>output<i} dan tanda tangan. Nama terdiri dari awalan stablehlo. dan mnemonic yang secara unik mengidentifikasi salah satu operasi yang didukung. Lihat di bawah untuk daftar lengkap dari semua operasi yang didukung.

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

Operasi menggunakan input dan menghasilkan output. Input dikategorikan ke dalam nilai input (dihitung selama eksekusi), fungsi input (disediakan secara statis, karena dalam StableHLO bukan merupakan nilai kelas satu) dan atribut input (juga disediakan secara statis). Jenis input dan output dikonsumsi dan diproduksi oleh operasi bergantung pada mnemoniknya. Misalnya, add op menggunakan 2 nilai input dan menghasilkan 1 nilai output. Sebagai perbandingan, Operasi select_and_scatter menggunakan 3 nilai input, 2 fungsi input, dan 3 atribut input.

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

Fungsi input (yang juga disebut fungsi anonim) sangat mirip dengan fungsi bernama kecuali bahwa: 1) mereka tidak memiliki ID (oleh karena itu nama "anonim"), 2) mereka tidak mendeklarasikan jenis output (jenis output disimpulkan dari operasi return dalam fungsi).

Sintaks untuk fungsi input menyertakan bagian yang saat ini tidak digunakan (lihat Unused di atas) yang tersedia untuk kompatibilitas dengan MLIR. Di MLIR, ada konsep yang lebih umum tentang "kawasan" yang dapat memiliki beberapa "blok" operasi yang terhubung bersama melalui operasi lompatan. Blok ini memiliki ID yang sesuai ke produksi Unused, agar dapat dibedakan satu sama lain. StabilHLO tidak memiliki operasi lompatan, jadi bagian yang sesuai dari sintaks MLIR adalah tidak digunakan (tetapi masih ada).

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

Atribut input memiliki nama dan nilai yang merupakan salah satu atribut yang didukung konstanta. Mereka adalah cara utama untuk menentukan metadata statis untuk yang kurang penting. Misalnya, operasi concatenate menggunakan atribut dimension untuk menentukan dimensi yang menggabungkan nilai inputnya. Demikian pula, Operasi slice menggunakan beberapa atribut seperti start_indices dan limit_indices untuk menentukan batas yang digunakan untuk memotong nilai input.

Saat ini, program StableHLO di luar organisasi terkadang berisi atribut yang tidak dijelaskan dalam dokumen ini. Di masa mendatang, kami berencana menyerap atribut ini ke dalam opset StableHLO atau melarangnya yang muncul di program StableHLO. Sementara itu, berikut daftar atribut:

• layout (#629).
• mhlo.frontend_attributes (#628).
• mhlo.sharding (#619).
• output_operand_aliases (#740).
• Metadata lokasi (#594).

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

Tanda tangan op terdiri dari jenis semua nilai input (daftar jenis di di sisi kiri ->) dan jenis semua nilai output (daftar jenis di sisi kanan ->). Sebenarnya, tipe input tipe output hampir selalu redundan (karena untuk sebagian besar operasi StableHLO, jenis output dapat disimpulkan dari input). Meskipun demikian, op signature sengaja dijadikan bagian dari sintaksis StableHLO untuk kompatibilitas dengan MLIR.

Di bawah ini adalah contoh operasi yang mnemoniknya adalah select_and_scatter. Menggunakan 3 nilai input (%operand, %source, dan %init_value), 2 fungsi input dan 3 atribut input (window_dimensions, window_strides, dan padding). Perhatikan bagaimana tanda tangan operasi hanya menyertakan jenis nilai inputnya (tetapi bukan jenis fungsi dan atribut input yang disediakan secara inline).

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

Konstanta

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

Konstanta SttableHLO memiliki literal dan jenis yang bersama-sama mewakili nilai StableHLO. Umumnya, jenis ini adalah bagian dari sintaks konstan, kecuali jika tidak ambigu (misalnya, konstanta boolean yang jelas memiliki jenis i1, sedangkan konstanta bilangan bulat dapat memiliki beberapa kemungkinan jenis).

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

Konstanta Boolean mewakili nilai boolean true dan false. Boolean konstanta memiliki jenis i1.

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

Konstanta bilangan bulat mewakili nilai bilangan bulat melalui string yang menggunakan bilangan desimal atau notasi heksadesimal. Basis lainnya, misalnya biner atau oktal, tidak didukung. Konstanta bilangan bulat memiliki batasan berikut:

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

Konstanta floating point merepresentasikan nilai floating point melalui string yang menggunakan desimal atau notasi ilmiah. Selain itu, notasi heksadesimal dapat digunakan untuk menentukan bit yang mendasarinya secara langsung dalam format floating point jenis yang sesuai. Konstanta floating point memiliki batasan berikut:

• (C1) Jika notasi non-heksadesimal digunakan, is_wellformed(float_literal, float_type).
• (C2) Jika notasi heksadesimal digunakan, size(hexadecimal_digits) = num_bits(float_type) / 4.

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

Konstanta kompleks merepresentasikan nilai kompleks menggunakan daftar bagian nyata (ditempatkan lebih dahulu) dan bagian imajiner (ditempatkan kedua). Misalnya, (1.0, 0.0) : complex<f32> mewakili 1.0 + 0.0i, dan (0.0, 1.0) : complex<f32> mewakili 0.0 + 1.0i. Urutan dari yang kemudian disimpan di memori adalah yang ditentukan implementasinya. Konstanta kompleks memiliki batasan berikut:

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

Konstanta Tensor mewakili nilai tensor menggunakan daftar bertingkat yang ditentukan melalui Notasi NumPy. Misalnya, dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> mewakili nilai tensor dengan pemetaan berikut dari indeks ke elemen: {0, 0} => 1, {0, 1} => 2, {0, 2} => 3, {1, 0} => 4, {1, 1} => 5, {1, 2} => 6. Urutan elemen-elemen ini kemudian disimpan dalam memori adalah ditentukan oleh implementasi. Konstanta tensor memiliki batasan berikut:

• (C1) has_syntax(tensor_literal, element_type(tensor_type)), dengan kondisi:
  • 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)), dengan kondisi:
  • 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:]).
  • jika tidak, false.

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

Konstanta tensor terkuantisasi mewakili nilai tensor terkuantisasi menggunakan notasi sebagai konstanta tensor, dengan elemen yang ditetapkan sebagai konstanta jenis penyimpanan. Konstanta tensor terkuantisasi memiliki batasan berikut:

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

Literal string terdiri dari byte yang ditentukan menggunakan karakter ASCII dan urutan escape. Model ini bersifat agnostik encoding, jadi interpretasinya byte ditentukan oleh implementasi. Literal string memiliki jenis string.

Operasi

abs

Semantik

Melakukan operasi abs berbasis element-wise pada tensor operand dan menghasilkan result tensor. Bergantung pada jenis elemen, melakukan hal berikut:

• Untuk bilangan bulat bertanda tangan: modulus bilangan bulat.
• Untuk float: abs dari IEEE-754.
• Untuk bilangan kompleks: modulus kompleks.
• Untuk jenis terkuantisasi: dequantize_op_quantize(abs, operand, type(result)).

Input

Output

Batasan

• (C1) shape(result) = shape(operand).
• (C2) baseline_element_type(result) didefinisikan sebagai:
  • complex_element_type(element_type(operand)) jika is_complex(operand).
  • baseline_element_type(operand) sebaliknya.

Contoh

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

Contoh Lainnya

tambahkan

Semantik

Melakukan penambahan berdasarkan elemen pada dua tensor lhs dan rhs serta menghasilkan Tensor result. Bergantung pada jenis elemen, melakukan hal berikut:

• Untuk boolean: OR logis.
• Untuk bilangan bulat: penambahan bilangan bulat.
• Untuk float: addition dari IEEE-754.
• Untuk bilangan kompleks: penambahan kompleks.
• Untuk jenis terkuantisasi: dequantize_op_quantize(add, lhs, rhs, type(result)).

Input

Output

Batasan

• Jika operasi tersebut menggunakan tensor non-kuantisasi:
  • (C1) type(lhs) = type(rhs) = type(result).
• Jika operasi tersebut menggunakan tensor terkuantisasi:
  • (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) Jika is_per_axis_quantized(lhs), maka quantization_dimension(lhs) = quantization_dimension(result).
  • (C7) Jika is_per_axis_quantized(rhs), maka quantization_dimension(rhs) = quantization_dimension(result).

Contoh

// %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]]

Contoh Lainnya

after_all

Semantik

Memastikan operasi yang menghasilkan inputs dijalankan sebelum operasi yang bergantung pada result. Eksekusi operasi ini tidak melakukan apa pun, fungsi tersebut hanya ada untuk membuat dependensi data dari result hingga inputs.

Input

Output

Contoh

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

Contoh Lainnya

all_gather

Semantik

Dalam setiap grup proses di grid proses StableHLO, merangkai nilai-nilai tensor operands dari setiap proses di sepanjang all_gather_dim dan menghasilkan Tensor results.

Operasi tersebut membagi grid proses StableHLO menjadi process_groups yang didefinisikan sebagai berikut:

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

Setelah itu, dalam setiap process_group:

• operands...@receiver = [operand@sender for sender in process_group] untuk semua receiver dalam process_group.
• results...@process = concatenate(operands...@process, all_gather_dim) untuk semua process dalam process_group.

Input

Output

Batasan

• (C1) 0 <= all_gather_dim < rank(operands...).
• (C2) is_unique(replica_groups).
• (C3) size(replica_groups) didefinisikan sebagai:
  • num_replicas jika cross_replica digunakan.
  • num_replicas jika cross_replica_and_partition digunakan.
  • num_processes jika flattened_ids digunakan.
• (C4) 0 <= replica_groups < size(replica_groups).
• (C5) Jika use_global_device_ids = true, maka channel_id > 0.
• (C6) type(results...) = type(operands...) kecuali:
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1).

Contoh

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

Contoh Lainnya

all_reduce

Semantik

Dalam setiap grup proses pada grid proses StableHLO, menerapkan pengurangan fungsi computation dengan nilai tensor operands dari setiap proses dan menghasilkan tensor results.

Operasi tersebut membagi grid proses StableHLO menjadi process_groups yang didefinisikan sebagai berikut:

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

Setelah itu, dalam setiap process_group:

• results...@process[result_index] = exec(schedule) untuk beberapa hierarki biner schedule dengan:
  • exec(node) = computation(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value.
• schedule adalah hierarki biner yang ditentukan penerapan yang sesuai urutan traversal adalah to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0])).

Input

Output

Batasan

• (C1) is_unique(replica_groups).
• (C2) size(replica_groups) didefinisikan sebagai:
  • num_replicas jika cross_replica digunakan.
  • num_replicas jika cross_replica_and_partition digunakan.
  • num_processes jika flattened_ids digunakan.
• (C3) 0 <= replica_groups < size(replica_groups).
• (C4) Jika use_global_device_ids = true, maka channel_id > 0.
• (C5) computation memiliki jenis (tensor<E>, tensor<E>) -> (tensor<E>), dengan is_promotable(element_type(operand), E).
• (C6) shape(results...) = shape(operands...).
• (C7) element_type(results...) = E.

Contoh

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

Contoh Lainnya

all_to_all

Semantik

Dalam setiap grup proses dalam {i>process grid<i} StableHLO, membagi nilai tensor operands di sepanjang split_dimension menjadi beberapa bagian, sehingga akan menyebarkan bagian di antara proses, menyambungkan bagian yang tersebar di concat_dimension dan menghasilkan tensor results. Operasi tersebut membagi grid proses StableHLO menjadi process_groups yang didefinisikan sebagai berikut:

• cross_replica(replica_groups) jika channel_id <= 0.
• cross_partition(replica_groups) jika channel_id > 0.

Setelah itu, dalam setiap process_group:

• split_parts...@sender = split(operands...@sender, split_count, split_dimension) untuk semua sender di process_group.
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group] di mana receiver_index = process_group.index(receiver).
• results...@process = concatenate(scattered_parts...@process, concat_dimension).

Input

Output

Batasan

• (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) ditentukan sebagai:
  • num_replicas jika cross_replica digunakan.
  • num_partitions jika cross_partition digunakan.
• (C7) 0 <= replica_groups < size(replica_groups).
• (C8) dim(replica_groups, 1) = split_count.
• (C9) type(results...) = type(operands...) kecuali, jika split_dimension != concat_dimension:
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count.
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count.

Contoh

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

Contoh Lainnya

dan

Semantik

Melakukan AND berdasarkan elemen dari dua tensor lhs dan rhs serta menghasilkan result tensor. Bergantung pada jenis elemen, melakukan hal berikut:

• Untuk boolean: logika AND.
• Untuk bilangan bulat: bitwise AND.

Input

Output

Batasan

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

Contoh

// %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]]

Contoh Lainnya

atan2

Semantik

Menjalankan operasi atan2 berdasarkan elemen pada tensor lhs dan rhs serta menghasilkan Tensor result. Bergantung pada jenis elemen, melakukan hal berikut:

• Untuk float: atan2 dari IEEE-754.
• Untuk bilangan kompleks: atan2 kompleks.
• Untuk jenis terkuantisasi: dequantize_op_quantize(atan2, lhs, rhs, type(result)).

Input

Output

Batasan

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

Contoh

// %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]

Contoh Lainnya

batch_norm_grad

Semantik

Menghitung gradien beberapa input propagasi mundur batch_norm_training dari grad_output, dan menghasilkan grad_operand, grad_scale, dan grad_offset tensor. Secara lebih formal, operasi ini bisa dinyatakan sebagai dekomposisi menjadi operasi StabilHLO yang ada menggunakan sintaksis Python sebagai berikut:

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

Untuk jenis terkuantisasi, menjalankan 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)).

Input

Output

Batasan

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, mean, variance, grad_output, grad_operand, grad_scale dan grad_offset memiliki baseline_element_type yang sama.
• (C3) operand, grad_output, dan grad_operand memiliki bentuk yang sama.
• (C4) scale, mean, variance, grad_scale, dan grad_offset memiliki bentuk yang sama.
• (C5) size(scale) = dim(operand, feature_index).

Contoh

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

Semantik

Menormalkan tensor operand di semua dimensi kecuali untuk feature_index dan menghasilkan tensor result. Secara lebih formal, operasi dapat dinyatakan sebagai dekomposisi untuk operasi StableHLO yang ada menggunakan sintaks Python sebagai berikut:

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)

Untuk jenis terkuantisasi, menjalankan 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)).

Input

Output

Batasan

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, mean, variance, dan result memiliki baseline_element_type yang sama.
• (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).

Contoh

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

Semantik

Menghitung rata-rata dan varians di semua dimensi kecuali untuk feature_index dan menormalisasi tensor operand yang menghasilkan output, batch_mean dan tensor batch_var. Secara lebih formal, operasi ini dapat dinyatakan sebagai dekomposisi ke operasi StableHLO yang ada menggunakan sintaksis Python sebagai berikut ini:

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

Untuk jenis terkuantisasi, menjalankan 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)).

Input

Output

Batasan

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, batch_mean, batch_var, dan output memiliki baseline_element_type yang sama.
• (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).

Contoh

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

Semantik

Melakukan operasi bitcast pada tensor operand dan menghasilkan tensor result dengan bit dari seluruh tensor operand diinterpretasikan ulang menggunakan dari tensor result.

Secara lebih formal, mengingat E = element_type(operand), E' = element_type(result), dan R = rank(operand):

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

bits menampilkan representasi dalam memori dari nilai tertentu, beserta perilakunya adalah implementasi-ditentukan karena representasi yang tepat dari tensor yang ditentukan implementasinya, dan representasi yang tepat dari tipe elemen adalah serta definisi implementasi.

Input

Output

Batasan

• (C1) Dengan mempertimbangkan E = is_quantized(operand) ? storage_type(operand) : element_type(operand), E' = is_quantized(result) ? storage_type(result) : element_type(result), dan R = rank(operand):
  • Jika num_bits(E') = num_bits(E), shape(result) = shape(operand).
  • Jika num_bits(E') < num_bits(E):
  • rank(result) = R + 1.
  • dim(result, i) = dim(operand, i) untuk semua 0 <= i < R.
  • dim(result, R) * num_bits(E') = num_bits(E).
  • Jika num_bits(E') > num_bits(E):
  • rank(result) = R - 1.
  • dim(result, i) = dim(operand, i) untuk semua 0 <= i < R.
  • dim(operand, R - 1) * num_bits(E) = num_bits(E').
• (C2) Jika is_complex(operand) or is_complex(result), maka is_complex(operand) and is_complex(result).

Contoh

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

Contoh Lainnya

broadcast_in_dim

Semantik

Memperluas dimensi dan/atau peringkat tensor input dengan menduplikasi data pada tensor operand dan menghasilkan tensor result. Secara lebih formal, result[result_index] = operand[operand_index] untuk semua d di axes(operand):

• operand_index[d] = 0 jika dim(operand, d) = 1.
• operand_index[d] = result_index[broadcast_dimensions[d]] sebaliknya.

Input

Output

Batasan

• (C1) element_type(result) diberikan oleh:
  • element_type(operand), jika !is_per_axis_quantized(operand).
  • element_type(operand) kecuali quantization_dimension(operand), scales(operand), dan zero_points(operand) mungkin berbeda dengan quantization_dimension(result), scales(result), dan zero_points(result) resp., jika tidak.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Untuk semua d di axes(operand):
  • dim(operand, d) = 1 atau
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Jika is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Jika dim(operand, quantization_dimension(operand)) = 1, maka scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))).

Contoh

// %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]
//            ]
//          ]

Contoh Lainnya

casing

Semantik

Menghasilkan output dari mengeksekusi tepat satu fungsi dari branches bergantung pada nilai index. Secara lebih formal, result = selected_branch() dalam hal ini:

• selected_branch = branches[index] jika 0 <= index < size(branches).
• selected_branch = branches[-1] sebaliknya.

Input

Output

Batasan

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

Contoh

// %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]

Contoh Lainnya

Cbrt

Semantik

Melakukan operasi root kubik berdasarkan elemen pada tensor operand dan menghasilkan Tensor result. Bergantung pada jenis elemen, melakukan hal berikut:

• Untuk float: rootn(x, 3) dari IEEE-754.
• Untuk bilangan kompleks: akar kubik kompleks.
• Untuk jenis terkuantisasi: dequantize_op_quantize(cbrt, operand, type(result))

Input

Output

Batasan

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

Contoh

// %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]

Contoh Lainnya

ceil

Semantik

Melakukan ceil berbasis elemen dari tensor operand dan menghasilkan tensor result. Menerapkan operasi roundToIntegralTowardPositive dari IEEE-754 spesifikasi pendukung. Untuk jenis terkuantisasi, menjalankan dequantize_op_quantize(ceil, operand, type(result)).

Input

Output

Batasan

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

Contoh

// %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]

Contoh Lainnya

Cholesky

Semantik

Menghitung dekomposisi Cholesky dari sekumpulan matriks.

Secara lebih formal, untuk semua i di index_space(result), result[i0, ..., iR-3, :, :] adalah dekomposisi Cholesky dari a[i0, ..., iR-3, :, :], dalam bentuk segitiga rendah (jika lower adalah true) atau matriks segitiga atas (jika lower adalah false). Nilai {i>output<i} dalam segitiga yang berlawanan, yaitu segitiga atas yang tegas atau segitiga bawah yang ketat, yang juga ditentukan oleh implementasi.

Jika ada i yang matriks inputnya bukan positif-definit Hermitian matriks, maka perilakunya tidak terdefinisi.

Untuk jenis terkuantisasi, menjalankan dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)).

Input

Output

Batasan

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

Contoh

// %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]
//          ]

klem

Semantik

Menjepit setiap elemen tensor operand di antara nilai minimum dan maksimum dan menghasilkan tensor result. Secara lebih formal, result[result_index] = minimum(maximum(operand[result_index], min_element), max_element), dengan min_element = rank(min) = 0 ? min[] : min[result_index], max_element = rank(max) = 0 ? max[] : max[result_index]. Untuk jenis terkuantisasi, menjalankan dequantize_op_quantize(clamp, min, operand, max, type(result)).

Memaksakan urutan pada bilangan kompleks melibatkan semantik yang mengejutkan, jadi di masa mendatang kami berencana menghapus dukungan untuk bilangan kompleks untuk operasi ini (#560).

Input

Output

Batasan

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

Contoh

// %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]

Contoh Lainnya

collective_broadcast

Semantik

Dalam setiap grup proses di grid proses StableHLO, kirim nilai proses Tensor operand dari proses sumber ke proses target dan menghasilkan Tensor result.

Operasi tersebut membagi grid proses StableHLO menjadi process_groups yang didefinisikan sebagai berikut:

• cross_replica(replica_groups) jika channel_id <= 0.
• cross_partition(replica_groups) jika channel_id > 0.

Setelah itu, result@process diberikan oleh:

• operand@process_groups[i, 0] jika ada i sehingga prosesnya menjadi di process_groups[i].
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) sebaliknya.

Input

Output

Batasan

• (C1) is_unique(replica_groups).
• (C2) 0 <= replica_groups < N dengan N ditentukan sebagai:
  • num_replicas jika cross_replica digunakan.
  • num_partitions jika cross_partition digunakan.
• (C3) type(result) = type(operand).

Contoh

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

Semantik

Dalam setiap grup proses pada grid proses StableHLO, mengirimkan nilai Tensor operand dari proses sumber ke proses target dan menghasilkan Tensor result.

Operasi tersebut membagi grid proses StableHLO menjadi process_groups yang didefinisikan sebagai berikut:

• cross_replica(source_target_pairs) jika channel_id <= 0.
• cross_partition(source_target_pairs) jika channel_id > 0.

Setelah itu, result@process diberikan oleh:

• operand@process_groups[i, 0], jika ada i sehingga process_groups[i, 1] = process.
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) sebaliknya.

Input

Output

Batasan

• (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, dengan N ditentukan sebagai:
  • num_replicas jika cross_replica digunakan.
  • num_partitions jika cross_partition digunakan.
• (C5) type(result) = type(operand).

Contoh

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

Contoh Lainnya

bandingkan

Semantik

Melakukan perbandingan tensor lhs dan rhs berdasarkan elemen sesuai dengan comparison_direction dan compare_type, serta menghasilkan tensor result.

Nilai comparison_direction dan compare_type memiliki hal berikut semantik:

Untuk jenis elemen boolean dan bilangan bulat:

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

Untuk jenis elemen floating point dengan compare_type = FLOAT, op akan mengimplementasikan operasi IEEE-754 berikut:

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

Untuk jenis elemen floating point dengan compare_type = TOTALORDER, op menggunakan kombinasi operasi totalOrder dan compareQuietEqual dari IEEE-754.

Untuk jenis elemen kompleks, perbandingan leksikografis pasangan (real, imag) dijalankan menggunakan comparison_direction dan compare_type yang disediakan. Memaksakan urutan pada bilangan kompleks melibatkan semantik yang mengejutkan, jadi di masa mendatang kami berencana menghapus dukungan untuk bilangan kompleks jika comparison_direction adalah GE, GT, LE, atau LT (#560).

Untuk jenis terkuantisasi. menjalankan dequantize_compare(lhs, rhs, comparison_direction).

Input

Output

Batasan

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs).
• (C2) shape(lhs) = shape(rhs) = shape(result).
• (C3) compare_type didefinisikan sebagai:
  • SIGNED jika is_signed_integer(element_type(lhs)).
  • UNSIGNED jika is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)).
  • FLOAT atau TOTALORDER jika is_float(element_type(lhs)).
  • FLOAT jika is_complex(element_type(lhs)).

Contoh

// %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]

Contoh Lainnya

kompleks

Semantik

Melakukan konversi {i>element<i}-{i>wise<i} ke nilai kompleks dari pasangan real dan nilai imajiner, lhs dan rhs, serta menghasilkan tensor result.

Input

Output

Batasan

• (C1) type(lhs) = type(rhs).
• (C2) shape(result) = shape(lhs).
• (C3) element_type(result) memiliki jenis complex<E>, dengan E = element_type(lhs).

Contoh

// %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)]

Contoh Lainnya

gabungan

Semantik

Mengenkapsulasi operasi yang terdiri (tersusun) dari operasi StableHLO lainnya, mengambil inputs dan composite_attributes, serta memproduksi results. Tujuan semantik op diimplementasikan oleh atribut decomposition. Tujuan op composite dapat diganti dengan dekomposisinya tanpa mengubah program semantik. Dalam kasus saat penyisipan dekomposisi tidak memberikan hasil semantik operasi, sebaiknya gunakan custom_call.

Kolom version (default-nya adalah 0) digunakan untuk menunjukkan kapan kolom semantik berubah.

Input

Output

Batasan

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

Contoh

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

Contoh Lainnya

concatenate

Semantik

Menggabungkan inputs di sepanjang dimensi dimension dalam urutan yang sama seperti yang ditentukan argumen dan menghasilkan tensor result. Secara lebih formal, result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1], dengan:

1. id = d0 + ... + dk-1 + kd.
2. d sama dengan dimension, dan d0, ... adalah ukuran dimensi ke-d dari inputs.

Input

Output

Batasan

• (C1) same(element_type(inputs...)).
• (C2) same(shape(inputs...)) kecuali untuk 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]) kecuali untuk:
  • dim(result, dimension) = dim(inputs[0], dimension) + ....

Contoh

// %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]]

Contoh Lainnya

konstanta

Semantik

Menghasilkan tensor output dari value yang konstan.

Input

Output

Batasan

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

Contoh

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

Contoh Lainnya

melakukan konversi

Semantik

Melakukan konversi berbasis elemen dari satu jenis elemen ke jenis elemen lainnya pada Tensor operand dan menghasilkan tensor result.

Untuk konversi boolean-to-any-supported-type, nilai false adalah dikonversi menjadi nol, dan nilai true dikonversi menjadi satu. Sebagai any-supported-type-to-boolean, nilai nol dikonversi menjadi false, dan nilai selain nol dikonversi menjadi true. Lihat di bawah untuk mengetahui cara bekerja untuk tipe-tipe yang kompleks.

Untuk konversi yang melibatkan integer-ke-bilangan bulat, bilangan bulat-ke-floating-point atau floating-point-to-floating-point, jika nilai sumber bisa tepat direpresentasikan dalam jenis tujuan, nilai hasil adalah nilai persis merepresentasinya. Jika tidak, perilaku akan ditentukan nanti (#180).

Untuk konversi yang melibatkan floating-point-to-integer, bagian pecahannya adalah terpotong. Jika nilai yang terpotong tidak dapat ditampilkan dalam jenis tujuan, perilakunya akan ditentukan nanti (#180).

Konversi yang melibatkan kompleks-ke-kompleks mengikuti perilaku yang sama dari konversi floating-point-to-floating-point untuk mengonversi real time dan bagian imajiner.

Untuk konversi complex-to-any-other-type dan complex-to-any-other-type, nilai imajiner sumber diabaikan atau nilai imajiner tujuannya adalah masing-masing menjadi nol. Konversi dari bagian riil mengikuti konversi floating point.

Pada prinsipnya, operasi ini dapat mengekspresikan dekuantisasi (konversi dari tensor terkuantisasi ke tensor reguler), kuantisasi (konversi dari tensor ke tensor terkuantisasi) dan rekuantisasi (konversi antara tensor), tetapi saat ini kami memiliki operasi khusus untuk itu - uniform_dequantize untuk kasus penggunaan pertama dan uniform_quantize untuk kasus penggunaan pertama kedua dan ketiga. Di masa mendatang, kedua operasi ini dapat digabungkan ke dalam convert (#1576).

Input

Output

Batasan

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

Contoh

// %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)]

Contoh Lainnya

konvolusi

Semantik

Menghitung produk titik antara jendela lhs dan irisan rhs serta menghasilkan result. Diagram berikut menunjukkan cara penghitungan elemen pada result lhs dan rhs menggunakan contoh konkret.

Secara lebih formal, pertimbangkan framing ulang input berikut dalam hal lhs agar dapat mengekspresikan jendela 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).

Penyesuaian frame ini menggunakan fungsi bantuan berikut:

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

Jika feature_group_count = 1 dan batch_group_count = 1, maka untuk semua output_spatial_index di index_space(dim(result, output_spatial_dimensions...)), result[result_shape(:, output_spatial_index, :)] = dot_product dengan:

• 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]). Fitur ini tampaknya tidak digunakan, jadi di masa mendatang kami berencana untuk menghapusnya (#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]).

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

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

Untuk jenis terkuantisasi, menjalankan 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)).

Untuk jenis terkuantisasi hybrid, menjalankan 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).

Input

Output

Batasan

• (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) Dengan mempertimbangkan 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) Dengan mempertimbangkan 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) Dengan mempertimbangkan 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) didefinisikan sebagai:
  • dim(lhs, input_batch_dimension) / batch_group_count jika result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) jika result_dim = output_feature_dimension.
  • num_windows jika tidak, jika:
  • 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.
• Jika operasi tersebut menggunakan tensor non-kuantisasi:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Jika operasi tersebut menggunakan tensor terkuantisasi:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Jika is_per_axis_quantized(rhs), lalu quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Jika is_per_axis_quantized(result), maka quantization_dimension(result) = output_feature_dimension.
  • Jika is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Jika is_per_tensor_quantized(rhs), maka is_per_tensor_quantized(result).
  • Jika !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Contoh

// %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]]
//          ]]

Contoh Lainnya

kosinus

Semantik

Melakukan operasi kosinus element-wise pada tensor operand dan menghasilkan Tensor result. Bergantung pada jenis elemen, melakukan hal berikut:

• Untuk float: cos dari IEEE-754.
• Untuk bilangan kompleks: kosinus kompleks.
• Untuk jenis terkuantisasi: dequantize_op_quantize(cosine, operand, type(result)).

Input

Output

Batasan

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

Contoh

// %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]]

Contoh Lainnya

count_leading_zeros

Semantik

Melakukan penghitungan berbasis elemen dari jumlah bit nol di awal dalam operand dan menghasilkan tensor result.

Input

Output

Batasan

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

Contoh

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

Contoh Lainnya

custom_call

Semantik

Mengenkapsulasi call_target_name operasi yang ditentukan implementasi yang menggunakan inputs dan called_computations serta menghasilkan results. has_side_effect, backend_config dan api_version dapat digunakan untuk memberikan metadata yang ditentukan implementasi.

Saat ini, operasi ini berisi kumpulan data yang tidak terorganisir {i>metadata <i}yang mencerminkan evolusi organik dari operasi kompilator XLA. Di masa mendatang, kami berencana untuk menyatukan metadata ini (#741).

Input

Output

Contoh

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

bagi

Semantik

Melakukan pembagian menurut elemen dari tensor lhs dan pembagi rhs serta menghasilkan tensor result. Bergantung pada jenis elemen, melakukan hal berikut:

• Untuk bilangan bulat: pembagian bilangan bulat yang menghasilkan hasil bagi aljabar dengan bagian pecahan dibuang.
• Untuk float: division dari IEEE-754.
• Untuk bilangan kompleks: pembagian kompleks.
• Untuk jenis terkuantisasi:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)).

Input

Output

Batasan

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

Contoh

// %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]

Contoh Lainnya

dot_general

Semantik

Menghitung produk titik di antara irisan lhs dan irisan rhs serta menghasilkan Tensor result.

Secara lebih formal, result[result_index] = dot_product, dengan:

• 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 dengan size(result_batching_index) = size(lhs_batching_dimensions), size(result_lhs_index) = size(lhs_result_dimensions) dan 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)).

Untuk jenis terkuantisasi, menjalankan 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)).

Untuk jenis terkuantisasi hybrid, menjalankan 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 mengontrol keseimbangan antara kecepatan dan akurasi untuk komputasi pada backend akselerator. Kolom ini dapat berupa salah satu dari hal berikut (di semantik, semantik dari nilai enum ini kurang ditentukan, namun kita berencana untuk mengatasinya dalam #755):

• DEFAULT: Penghitungan tercepat, tetapi perkiraan yang paling tidak akurat ke nomor aslinya.
• HIGH: Penghitungan lebih lambat, tetapi perkiraan yang lebih akurat terhadap nomor aslinya.
• HIGHEST: Penghitungan paling lambat, tetapi perkiraan paling akurat ke nomor aslinya.

DotAlgorithm menentukan properti utama algoritma yang digunakan untuk menerapkan operasi titik, yang juga menentukan presisi. Jika atribut algoritma kolom ditetapkan, maka precision_config harus DEFAULT. DotAlgorithms tidak memiliki nilai default, karena parameter default adalah implementasi didefinisikan. Dengan demikian, semua kolom algoritme titik dapat ditetapkan ke None untuk menentukan algoritma titik kosong, yang akan menggunakan nilai precision_config.

Kolom DotAlgorithm mencakup:

• lhs_precision_type dan rhs_precision_type, presisi yang ditentukan oleh LHS dan RHS operasi dibulatkan ke. Jenis presisi tidak bergantung pada jenis penyimpanan input dan output-nya.
• accumulation_type presisi yang digunakan untuk akumulasi.
• lhs_component_count, rhs_component_count, dan num_primitive_operations terapkan saat kita melakukan algoritma yang menguraikan LHS dan/atau RHS menjadi banyak komponen dan melakukan beberapa “primitif” operasi dot pada nilai - biasanya untuk mengemulasi presisi yang lebih tinggi (mis. Memanfaatkan Jenis Data Kecerdasan Buatan bfloat16 untuk Komputasi Presisi Lebih Tinggi: bf16_6x tf32_3x, dll.). Untuk algoritma tanpa dekomposisi, nilai ini harus ditetapkan ke 1.
• allow_imprecise_accumulation untuk menentukan apakah akumulasi dalam presisi lebih rendah diizinkan untuk beberapa langkah (mis. CUBLASLT_MATMUL_DESC_FAST_ACCUM).

Contoh atribut 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}

Implementasi bebas untuk memutuskan kombinasi mana yang didukung. Di beberapa umum, tidak ada jaminan bahwa setiap algoritma didukung pada setiap jenis akselerator oleh konsumen StableHLO. Jika algoritma tertentu tidak didukung, maka kesalahan akan dimunculkan, bukan untuk kembali ke alternatif. Verifikasi StabilHLO akan memberikan upaya verifikasi terbaik, mencegah algoritme yang tidak diketahui didukung pada perangkat keras apa pun.

Lihat xla_data.proto > Algorithm untuk beberapa nilai algoritma yang didukung. Tiket #2483 berisi rencana untuk membuat dokumen terpusat pada algoritma yang didukung oleh backend.

Input

Output

Batasan

• (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).
• Jika operasi tersebut menggunakan tensor non-kuantisasi:
  • (C13) element_type(lhs) = element_type(rhs).
• Jika operasi tersebut menggunakan tensor terkuantisasi:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C15) zero_points(rhs) = 0.
  • (C16) Jika is_per_axis_quantized(rhs), maka quantization_dimension(rhs) tidak ada di rhs_contracting_dimensions.
  • Jika is_quantized(lhs):
  • (C17) storage_type(lhs) = storage_type(rhs).
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C19) Jika is_per_tensor_quantized(rhs), maka is_per_tensor_quantized(result).
  • Jika !is_quantized(lhs):
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result).
• Jika !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.

Contoh

// %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]]
//          ]

Contoh Lainnya

dynamic_broadcast_in_dim

Semantik

Operasi ini secara fungsional identik dengan broadcast_in_dim berjalan, tetapi bentuk hasilnya ditentukan secara dinamis melalui output_dimensions.

Operasi tersebut juga menerima atribut opsional known_expanding_dimensions, known_non_expanding_dimensions untuk mengekspresikan pengetahuan statis tentang perilaku dimensi yang meluas. Jika tidak ditentukan, semua dimensi diasumsikan mungkin dapat diperluas.

Input

Output

Batasan

• (C1) element_type(result) diberikan oleh:
  • element_type(operand), jika !is_per_axis_quantized(operand).
  • element_type(operand) kecuali quantization_dimension(operand), scales(operand), dan zero_points(operand) mungkin berbeda dengan quantization_dimension(result), scales(result), dan zero_points(result) resp., jika tidak.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Untuk semua d di axes(operand):
  • dim(operand, d) = 1 atau
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Jika is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Jika dim(operand, quantization_dimension(operand)) = 1, maka 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).

Contoh

// %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]
//            ]
//          ]

Contoh Lainnya

dynamic_conv

Semantik

Operasi ini secara fungsional identik dengan konvolusi op, tetapi padding ditentukan secara dinamis melalui padding.

Input

Output

Batasan

• (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) Dengan mempertimbangkan 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) Dengan mempertimbangkan 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) Dengan mempertimbangkan 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) didefinisikan sebagai:
  • dim(lhs, input_batch_dimension) / batch_group_count jika result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) jika result_dim = output_feature_dimension.
  • num_windows jika tidak, jika:
  • 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.
• Jika operasi tersebut menggunakan tensor non-kuantisasi:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Jika operasi tersebut menggunakan tensor terkuantisasi:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Jika is_per_axis_quantized(rhs), lalu quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Jika is_per_axis_quantized(result), maka quantization_dimension(result) = output_feature_dimension.
  • Jika is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Jika is_per_tensor_quantized(rhs), maka is_per_tensor_quantized(result).
  • Jika !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Contoh

// %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]]
//          ]]