Spesifikasi StableHLO

StableHLO adalah kumpulan operasi untuk operasi tingkat tinggi (HLO) dalam model machine learning (ML). StableHLO 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.

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

Spesifikasi ini memuat tiga bagian utama. Pertama, bagian Program menjelaskan struktur program StableHLO yang terdiri dari fungsi StableHLO yang terdiri dari operasi StableHLO. Dalam struktur tersebut, bagian Ops menentukan semantik setiap operasi. Bagian Execution menyediakan semantik untuk semua operasi ini yang dijalankan bersama dalam program. Terakhir, bagian Notasi membahas notasi yang digunakan di seluruh spesifikasi.

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

Program

Program ::= {Func}

Program StableHLO terdiri dari sejumlah fungsi StableHLO yang tidak ditentukan. 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 StableHLO (yang juga disebut fungsi bernama) memiliki ID, input/output, dan isi. Di masa mendatang, kami berencana untuk 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 StableHLO mirip dengan ID dalam banyak bahasa pemrograman, dengan dua keunikan: 1) semua ID memiliki sigil yang membedakan berbagai jenis ID, 2) ID nilai dapat sepenuhnya numerik untuk menyederhanakan pembuatan program StableHLO.

Jenis

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

Jenis StableHLO dikategorikan ke dalam jenis nilai (yang juga disebut jenis kelas satu) yang mewakili nilai StableHLO dan jenis non-nilai yang menjelaskan elemen program lainnya. Jenis StableHLO mirip dengan jenis dalam banyak bahasa pemrograman, dengan keunikan utamanya adalah sifat khusus domain StableHLO yang menghasilkan beberapa hasil yang tidak biasa (misalnya, jenis skalar bukan merupakan jenis nilai).

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

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

Bentuk dapat sebagian atau sepenuhnya tidak diketahui (dinamis), misalnya tensor<?x2xf64> sebagian tidak diketahui dan tensor<?x?xf64> sepenuhnya tidak diketahui. Ukuran dimensi dinamis direpresentasikan menggunakan ?. Bentuk tidak boleh diberi peringkat.

Di masa mendatang, kami berencana untuk mempelajari perluasan jenis tensor di luar ukuran dimensi dan jenis elemen, misalnya, untuk menyertakan tata letak (#629) dan sparsitas (#1078).

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

Jenis elemen kuantisasi mewakili nilai bilangan bulat dari jenis penyimpanan dalam rentang dari storage_min hingga storage_max (inklusif) yang sesuai dengan nilai floating point dari jenis yang dinyatakan. Untuk nilai bilangan bulat i tertentu, nilai floating point yang sesuai dengan f dapat dikomputasi sebagai f = (i - zero_point) * scale, dengan scale dan zero_point disebut 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 minat yang kuat pada skala berbasis bilangan bulat, yang direpresentasikan dengan pengali dan pergeseran. Kami berencana untuk mempelajarinya dalam waktu dekat (#1404).

Ada diskusi yang sedang berlangsung tentang semantik QuantizationZeroPoint, termasuk jenis, nilai, dan apakah hanya ada satu atau berpotensi beberapa titik nol dalam jenis tensor kuantisasi. Berdasarkan hasil diskusi ini, spesifikasi seputar titik nol dapat berubah di masa mendatang (#1405).

Diskusi lain yang sedang berlangsung melibatkan semantik QuantizationStorageMin dan QuantizationStorageMax untuk menentukan apakah batasan apa pun harus diterapkan pada nilai ini dan pada nilai tensor kuantisasi (#1406).

Terakhir, kami berencana untuk mengeksplorasi representasi skala dan titik nol yang tidak diketahui, mirip dengan cara kami berencana untuk mengeksplorasi representasi ukuran dimensi yang tidak diketahui (#1407).

Jenis tensor terkuantisasi merepresentasikan tensor dengan elemen terkuantisasi. Tensor ini sama persis dengan tensor reguler, kecuali elemennya memiliki jenis elemen yang dikuantisasi, bukan jenis elemen reguler.

Dalam tensor kuantisasi, kuantisasi dapat berupa per-tensor, yang berarti memiliki satu scale dan zero_point untuk seluruh tensor atau dapat berupa per-sumbu, yang berarti memiliki beberapa scales dan zero_points, satu pasangan per slice dimensi tertentu quantization_dimension. Secara lebih formal, dalam tensor t dengan kuantisasi per sumbu, ada dim(t, quantization_dimension) slice dari quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :], dll. Semua elemen dalam slice ke-i menggunakan scales[i] dan zero_points[i] sebagai parameter kuantisasi. Jenis tensor terkuantisasi memiliki batasan berikut:

• 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 digunakan 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 fitur lama yang hanya ada untuk kompatibilitas dengan HLO. Di HLO, tuple digunakan untuk merepresentasikan input dan output variadik. Di StableHLO, input dan output variadik didukung secara native, dan satu-satunya penggunaan tuple di StableHLO adalah untuk mewakili HLO ABI secara komprehensif, misalnya T, tuple<T>, dan tuple<tuple<T>> mungkin secara material berbeda bergantung pada implementasi tertentu. Di masa mendatang, kami berencana untuk membuat perubahan pada HLO ABI yang dapat memungkinkan kita 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 ::= 'f4E2M1FN' | 'f6E2M3FN' | 'f6E3M2FN' | 'f8E3M4' | 'f8E4M3'
            | 'f8E4M3FN' | 'f8E4M3FNUZ' | 'f8E4M3B11FNUZ' | 'f8E5M2'
            | 'f8E5M2FNUZ' | 'f8E8M0FNU' | 'bf16' | 'f16' | 'f32' | 'f64'
TensorFloat32 ::= 'tf32'
ComplexType ::= 'complex' '<' ComplexElementType '>'
ComplexElementType ::= 'f32' | 'f64'

Jenis elemen mewakili elemen jenis tensor. Tidak seperti di banyak bahasa pemrograman, jenis ini bukan kelas satu di StableHLO. Artinya, program StableHLO tidak dapat langsung merepresentasikan nilai dari jenis ini (akibatnya, idiomatik untuk merepresentasikan nilai skalar dari jenis T dengan nilai tensor dimensi 0 dari jenis tensor<T>).

• Jenis boolean mewakili nilai boolean true dan false.
• Jenis bilangan bulat dapat bertanda (si) atau tidak bertanda (ui) dan memiliki salah satu lebar bit yang didukung (2, 4, 8, 16, 32, atau 64). Jenis siN bertanda mewakili nilai bilangan bulat dari -2^(N-1) hingga 2^(N-1)-1 inklusif, dan jenis uiN tanpa tanda mewakili nilai bilangan bulat dari 0 hingga 2^N-1 inklusif.
• Jenis floating point dapat berupa salah satu dari berikut:
  • Bilangan floating point 8-bit f8E3M4, f8E4M3, dan f8E5M2 mengikuti konvensi IEEE-754.
  • Jenis f8E4M3FN dan f8E5M2 masing-masing sesuai dengan encoding E4M3 dan E5M2 dari format FP8 yang dijelaskan dalam Format FP8 untuk Deep Learning.
  • Jenis f8E4M3FNUZ dan f8E5M2FNUZ yang sesuai dengan encoding E4M3 dan E5M2 dari format FP8 yang dijelaskan dalam Format Numerik 8-bit untuk Jaringan Neural Dalam.
  • Jenis f8E4M3B11FNUZ yang sesuai dengan encoding E4M3 dari format FP8 yang dijelaskan dalam Pelatihan dan Inferensi Floating Point 8-bit Hybrid (HFP8) untuk Deep Neural Networks.
  • Jenis bf16 yang sesuai dengan format bfloat16 yang dijelaskan dalam BFloat16: Rahasia performa tinggi di Cloud TPU.
  • Jenis f16, f32, dan f64 masing-masing sesuai dengan format binary16 ("presisi setengah"), binary32 ("presisi tunggal"), dan binary64 ("presisi ganda") yang dijelaskan dalam standar IEEE 754.
  • Jenis tf32 sesuai dengan format TensorFloat32 dan memiliki dukungan terbatas di StableHLO.
  • Jenis MX (penskalaan mikro) f4E2M1FN, f6E2M3FN, f6E3M2FN, dan f8E8M0FNU dijelaskan dalam Spesifikasi Format Penskalaan Mikro OCP.
• Jenis kompleks mewakili nilai kompleks yang memiliki bagian riil dan bagian imajiner dari jenis elemen yang sama. Jenis kompleks yang didukung adalah complex<f32> (kedua bagiannya berjenis f32) dan complex<f64> (kedua bagiannya berjenis f64).

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

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

StringType ::= 'string'

Jenis string mewakili urutan byte. Tidak seperti dalam banyak bahasa pemrograman, jenis string bukan kelas pertama di StableHLO dan hanya digunakan untuk menentukan metadata statis untuk elemen program.

Operasi

Operasi StableHLO (yang juga disebut ops) mewakili kumpulan tertutup operasi tingkat tinggi dalam model machine learning. Seperti yang telah dibahas di atas, sintaksis StableHLO sangat terinspirasi oleh MLIR, yang belum tentu merupakan alternatif paling ergonomis, tetapi dapat dibilang paling sesuai dengan sasaran StableHLO untuk membuat lebih banyak interoperabilitas antara framework ML dan compiler ML.

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

Operasi StableHLO (yang juga disebut ops) memiliki nama, input/output, dan tanda tangan. Nama terdiri dari awalan stablehlo. dan mnemonik yang secara unik mengidentifikasi salah satu operasi yang didukung. Lihat di bawah untuk mengetahui daftar lengkap semua operasi yang didukung.

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

Ops 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 yang digunakan dan dihasilkan oleh op bergantung pada mnemoninya. Misalnya, operasi add menggunakan 2 nilai input dan menghasilkan 1 nilai output. Sebagai perbandingan, op 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: 1) tidak memiliki ID (sehingga namanya "anonim"), 2) tidak mendeklarasikan jenis output (jenis output disimpulkan dari operasi return dalam fungsi).

Sintaksis untuk fungsi input mencakup bagian yang saat ini tidak digunakan (lihat produksi Unused di atas) yang ada untuk kompatibilitas dengan MLIR. Di MLIR, ada konsep "region" yang lebih umum yang dapat memiliki beberapa "blok" operasi yang terhubung bersama melalui operasi lompat. Blok ini memiliki ID yang sesuai dengan produksi Unused, sehingga dapat dibedakan satu sama lain. StableHLO tidak memiliki operasi lompat, sehingga bagian yang sesuai dari sintaksis MLIR tidak digunakan (tetapi masih ada).

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

Atribut input memiliki nama dan nilai yang merupakan salah satu konstanta yang didukung. Class ini adalah cara utama dalam menentukan metadata statis untuk elemen program. Misalnya, operasi concatenate menggunakan atribut dimension untuk menentukan dimensi tempat nilai inputnya digabungkan. Demikian pula, op slice menggunakan beberapa atribut seperti start_indices dan limit_indices untuk menentukan batas yang digunakan untuk memotong nilai input.

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

• 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 sisi kiri ->) dan jenis semua nilai output (daftar jenis di sisi kanan ->). Sederhananya, jenis input bersifat redundan, dan jenis output hampir selalu redundan (karena untuk sebagian besar operasi StableHLO, jenis output dapat disimpulkan dari input). Meskipun demikian, tanda tangan op sengaja menjadi bagian dari sintaksis StableHLO untuk kompatibilitas dengan MLIR.

Berikut adalah contoh operasi yang mnemoninya adalah select_and_scatter. Fungsi ini 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 StableHLO memiliki literal dan jenis yang bersama-sama mewakili nilai StableHLO. Umumnya, jenis ini adalah bagian dari sintaksis konstanta, kecuali jika tidak ambigu (misalnya, konstanta boolean memiliki jenis i1 secara tidak ambigu, sedangkan konstanta bilangan bulat dapat memiliki beberapa kemungkinan jenis).

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

Konstanta Boolean mewakili nilai boolean true dan false. Konstanta Boolean 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 merepresentasikan nilai bilangan bulat melalui string yang menggunakan notasi desimal atau heksadesimal. Basis lain, 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 mewakili nilai floating point melalui string yang menggunakan notasi desimal atau notasi ilmiah. Selain itu, notasi heksadesimal dapat digunakan untuk menentukan secara langsung bit yang mendasarinya dalam format floating point dari 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 mewakili nilai kompleks menggunakan daftar bagian real (didahulukan) dan bagian imajiner (didahulukan). 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 penyimpanan bagian-bagian ini dalam memori ditentukan oleh implementasi. 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 merepresentasikan nilai tensor menggunakan daftar bertingkat yang ditentukan melalui notasi NumPy. Misalnya, dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> merepresentasikan 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 penyimpanan elemen ini dalam memori ditentukan oleh implementasi. Konstanta tensor memiliki batasan berikut:

• (C1) has_syntax(tensor_literal, element_type(tensor_type)), dengan:
  • 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:
  • 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 kuantisasi mewakili nilai tensor kuantisasi menggunakan notasi yang sama dengan konstanta tensor, dengan elemen yang ditentukan sebagai konstanta dari jenis penyimpanannya. 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. Byte ini tidak bergantung pada encoding, sehingga interpretasi byte ini ditentukan oleh implementasi. String literal memiliki jenis string.

Operasi

abs

Semantik

Melakukan operasi abs element-wise pada tensor operand dan menghasilkan tensor result. Bergantung pada jenis elemen, lakukan tindakan 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) ditentukan sebagai:
  • complex_element_type(element_type(operand)) jika is_complex(operand).
  • baseline_element_type(operand) jika tidak.

Contoh

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

 Contoh Lainnya

tambahkan

Semantik

Melakukan penambahan element-wise dari dua tensor lhs dan rhs dan menghasilkan tensor result. Bergantung pada jenis elemen, lakukan tindakan 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 menggunakan tensor yang tidak dikuantisasi:
  • (C1) type(lhs) = type(rhs) = type(result).
• Jika operasi menggunakan tensor kuantisasi:
  • (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 bahwa operasi yang menghasilkan inputs dieksekusi sebelum operasi apa pun yang bergantung pada result. Eksekusi operasi ini tidak melakukan apa pun, hanya ada untuk menetapkan dependensi data dari result ke 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 petak proses StableHLO, gabungkan nilai tensor operands dari setiap proses di sepanjang all_gather_dim dan hasilkan tensor results.

Operasi ini membagi petak proses StableHLO menjadi process_groups yang ditentukan sebagai berikut:

• cross_replica(replica_groups) if 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) if 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 di process_group.
• results...@process = concatenate(operands...@process, all_gather_dim) untuk semua process di process_group.

Input

Output

Batasan

• (C1) 0 <= all_gather_dim < rank(operands...).
• (C2) is_unique(replica_groups).
• (C3) size(replica_groups) ditentukan 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 di petak proses StableHLO, menerapkan fungsi pengurangan computation ke nilai tensor operands dari setiap proses dan menghasilkan tensor results.

Operasi ini membagi petak proses StableHLO menjadi process_groups yang ditentukan sebagai berikut:

• cross_replica(replica_groups) if 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) if 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 implementasi yang traversal urutannya 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) ditentukan 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 di petak proses StableHLO, bagi nilai tensor operands di sepanjang split_dimension menjadi beberapa bagian, sebar bagian yang dibagi di antara proses, gabungkan bagian yang tersebar di sepanjang concat_dimension, dan hasilkan tensor results. Operasi ini membagi petak proses StableHLO menjadi process_groups yang ditentukan sebagai berikut:

• cross_replica(replica_groups) if channel_id <= 0.
• cross_partition(replica_groups) if 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] dengan 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 element-wise dari dua tensor lhs dan rhs serta menghasilkan tensor result. Bergantung pada jenis elemen, lakukan tindakan berikut:

• Untuk boolean: AND logika.
• 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

Melakukan operasi atan2 element-wise pada tensor lhs dan rhs serta menghasilkan tensor result. Bergantung pada jenis elemen, lakukan tindakan berikut:

• Untuk float: atan2 dari IEEE-754.
• Untuk bilangan kompleks: atan2 kompleks.
• Untuk jenis kuantisasi: 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 batch_norm_training yang melakukan backpropagation dari grad_output, dan menghasilkan tensor grad_operand, grad_scale, dan grad_offset. Secara lebih formal, operasi ini dapat dinyatakan sebagai dekomposisi ke operasi StableHLO 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 kuantisasi, lakukan 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

Menormalisasi tensor operand di semua dimensi kecuali dimensi feature_index dan menghasilkan tensor result. Secara lebih formal, operasi ini dapat dinyatakan sebagai dekomposisi ke operasi StableHLO yang ada menggunakan sintaksis 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 kuantisasi, lakukan 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 varian di semua dimensi kecuali untuk dimensi feature_index dan menormalisasi tensor operand yang menghasilkan tensor output, batch_mean, dan batch_var. Secara lebih formal, operasi ini dapat dinyatakan sebagai dekomposisi ke operasi StableHLO yang ada menggunakan sintaksis Python sebagai berikut:

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 kuantisasi, lakukan 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 ditafsirkan ulang menggunakan jenis tensor result.

Secara lebih formal, dengan 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, dan perilakunya ditentukan oleh implementasi karena representasi persis tensor ditentukan oleh implementasi, dan representasi persis jenis elemen juga ditentukan oleh implementasi.

Input

Output

Batasan

• (C1) Dengan 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 dalam tensor operand dan menghasilkan tensor result. Secara lebih formal, result[result_index] = operand[operand_index] dengan untuk semua d di axes(operand):

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

Input

Output

Batasan

• (C1) element_type(result) diberikan oleh:
  • element_type(operand), jika !is_per_axis_quantized(operand).
  • element_type(operand) kecuali bahwa quantization_dimension(operand), scales(operand), dan zero_points(operand) mungkin berbeda dari quantization_dimension(result), scales(result), dan zero_points(result) masing-masing.
• (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() dengan:

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

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 akar kubik element-wise pada tensor operand dan menghasilkan tensor result. Bergantung pada jenis elemen, lakukan tindakan berikut:

• Untuk float: rootn(x, 3) dari IEEE-754.
• Untuk bilangan kompleks: akar pangkat tiga 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 element-wise dari tensor operand dan menghasilkan tensor result. Mengimplementasikan operasi roundToIntegralTowardPositive dari spesifikasi IEEE-754. Untuk jenis kuantisasi, lakukan 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 batch 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 matriks segitiga bawah (jika lower adalah true) atau matriks segitiga atas (jika lower adalah false). Nilai output di segitiga yang berlawanan, yaitu segitiga atas yang ketat atau segitiga bawah yang ketat, ditentukan oleh implementasi.

Jika ada i di mana matriks input bukan matriks positif-definit Hermitian, perilakunya tidak ditentukan.

Untuk jenis kuantisasi, lakukan 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

Membatasi setiap elemen tensor operand 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 pengurutan pada angka 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 pada petak proses StableHLO, kirim nilai tensor operand dari proses sumber ke proses target dan hasilkan tensor result.

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

• cross_replica(replica_groups) if 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 berada di process_groups[i].
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) jika tidak.

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 petak proses StableHLO, mengirimkan nilai tensor operand dari proses sumber ke proses target dan menghasilkan tensor result.

Operasi ini membagi petak proses StableHLO menjadi process_groups yang ditentukan sebagai berikut:

• cross_replica(source_target_pairs) jika channel_id <= 0.
• cross_partition(source_target_pairs) if 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)) jika tidak.

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 element-wise dari tensor lhs dan rhs sesuai dengan comparison_direction dan compare_type, serta menghasilkan tensor result.

Nilai comparison_direction dan compare_type memiliki semantik berikut:

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 menerapkan 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 yang kompleks, perbandingan leksikografis pasangan (real, imag) dilakukan menggunakan comparison_direction dan compare_type yang disediakan. Menerapkan pengurutan pada bilangan kompleks melibatkan semantik yang mengejutkan, jadi pada masa mendatang, kami berencana untuk menghapus dukungan untuk bilangan kompleks saat comparison_direction adalah GE, GT, LE, atau LT (#560).

Untuk jenis kuantisasi. melakukan 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 ditentukan sebagai:
  • SIGNED if 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 elemen ke nilai kompleks dari sepasang nilai riil dan imaginer, 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, menggunakan inputs dan composite_attributes, serta menghasilkan results. Semantik op diterapkan oleh atribut decomposition. Operasi composite dapat diganti dengan dekomposisinya tanpa mengubah semantik program. Jika menyisipkan dekomposisi tidak memberikan semantik op yang sama, sebaiknya gunakan custom_call.

Kolom version (default-nya adalah 0) digunakan untuk menunjukkan kapan semantik komposit 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 dengan argumen yang diberikan 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 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 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 dikonversi menjadi nol, dan nilai true dikonversi menjadi satu. Untuk konversi any-supported-type-to-boolean, nilai nol dikonversi menjadi false, dan nilai non-nol dikonversi menjadi true. Lihat di bawah untuk mengetahui cara kerja ini untuk jenis yang kompleks.

Untuk konversi yang melibatkan bilangan bulat ke bilangan bulat, bilangan bulat ke floating point, atau floating point ke floating point, jika nilai sumber dapat direpresentasikan secara tepat dalam jenis tujuan, nilai hasilnya adalah representasi yang tepat tersebut. Jika tidak, perilakunya adalah TBD (#180).

Untuk konversi yang melibatkan floating-point-to-integer, bagian pecahan akan terpotong. Jika nilai yang dipangkas tidak dapat direpresentasikan dalam jenis tujuan, perilakunya adalah TBD (#180).

Konversi yang melibatkan kompleks ke kompleks mengikuti perilaku yang sama dengan konversi floating point ke floating point untuk mengonversi bagian riil dan imajiner.

Untuk konversi complex-to-any-other-type dan any-other-type-to-complex, nilai imajiner sumber diabaikan atau nilai imajiner tujuan disetel ke nol. Konversi bagian riil mengikuti konversi floating point.

Pada prinsipnya, operasi ini dapat mengekspresikan dekuantisasi (konversi dari tensor kuantisasi menjadi tensor reguler), kuantisasi (konversi dari tensor reguler menjadi tensor kuantisasi), dan rekuantisasi (konversi antara tensor kuantisasi), tetapi saat ini kami memiliki operasi khusus untuk itu - uniform_dequantize untuk kasus penggunaan pertama dan uniform_quantize untuk kasus penggunaan kedua dan ketiga. Nantinya, 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 perkalian titik antara jendela lhs dan slice rhs serta menghasilkan result. Diagram berikut menunjukkan cara elemen di result dihitung dari lhs dan rhs menggunakan contoh konkret.

Secara lebih formal, pertimbangkan penyusunan 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).

Pembingkaian ulang 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 pada 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 kuantisasi campuran, lakukan 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 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 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) Diberikan 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) ditentukan sebagai:
  • dim(lhs, input_batch_dimension) / batch_group_count if result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) jika result_dim = output_feature_dimension.
  • num_windows jika tidak, dengan:
  • 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 menggunakan tensor yang tidak dikuantisasi:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Jika operasi menggunakan tensor kuantisasi:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Jika is_per_axis_quantized(rhs), maka 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 berbasis elemen pada tensor operand dan menghasilkan tensor result. Bergantung pada jenis elemen, lakukan tindakan berikut:

• Untuk float: cos dari IEEE-754.
• Untuk bilangan kompleks: kosinus kompleks.
• Untuk jenis kuantisasi: 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 element-wise dari jumlah bit nol di awal dalam tensor 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 operasi call_target_name 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 tambahan yang ditentukan implementasi.

Saat ini, operasi ini berisi kumpulan metadata yang cukup tidak teratur yang mencerminkan evolusi organik dari operasi yang setara di compiler XLA. Pada 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 element-wise dari tensor dividen lhs dan pembagi rhs dan menghasilkan tensor result. Bergantung pada jenis elemen, lakukan tindakan berikut:

• Untuk bilangan bulat: pembagian bilangan bulat yang menghasilkan hasil bagi aljabar dengan bagian pecahan yang dihapus.
• Untuk float: division dari IEEE-754.
• Untuk bilangan kompleks: pembagian kompleks.
• Untuk jenis kuantisasi:
  • 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 perkalian titik antara slice lhs dan slice 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 kuantisasi, lakukan 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 kuantisasi campuran, lakukan 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 kompromi antara kecepatan dan akurasi untuk komputasi di backend akselerator. Ini dapat berupa salah satu dari berikut ini (saat ini, semantik nilai enum ini tidak ditentukan, tetapi kami berencana untuk mengatasinya di #755):

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

DotAlgorithm menentukan properti utama algoritma yang digunakan untuk menerapkan operasi titik, yang juga menentukan presisi. Jika kolom atribut algoritma ditetapkan, precision_config harus berupa DEFAULT. DotAlgorithms tidak memiliki nilai default, karena parameter default ditentukan oleh implementasi. Dengan demikian, semua kolom algoritma 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 digunakan untuk membulatkan LHS dan RHS operasi. Jenis presisi tidak bergantung pada jenis penyimpanan input dan output.
• accumulation_type presisi yang digunakan untuk akumulasi.
• lhs_component_count, rhs_component_count, dan num_primitive_operations berlaku saat kita melakukan algoritma yang menguraikan LHS dan/atau RHS menjadi beberapa komponen dan melakukan beberapa operasi titik "primitif" pada nilai tersebut, biasanya untuk mengemulasi presisi yang lebih tinggi (misalnya Memanfaatkan Jenis Data Kecerdasan Buatan bfloat16 Untuk Komputasi Presisi Tinggi: bf16_6x, tf32). Untuk algoritma tanpa dekomposisi, nilai ini harus ditetapkan ke 1.
• allow_imprecise_accumulation untuk menentukan apakah akumulasi dalam presisi yang lebih rendah diizinkan untuk beberapa langkah (misalnya, 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. Secara umum, tidak dijamin bahwa setiap algoritma didukung di setiap jenis akselerator oleh konsumen StableHLO. Jika algoritma tertentu tidak didukung, error harus ditampilkan, bukan kembali ke alternatif. Verifikasi StableHLO akan memberikan verifikasi upaya terbaik, yang mencegah algoritma yang tidak diketahui didukung di hardware mana pun.

Lihat xla_data.proto > Algorithm untuk mengetahui beberapa nilai algoritma yang didukung. Tiket #2483 berisi rencana untuk membuat dokumen terpusat tentang 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 menggunakan tensor yang tidak dikuantisasi:
  • (C13) element_type(lhs) = element_type(rhs).
• Jika operasi menggunakan tensor kuantisasi:
  • (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 dalam 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 operasi broadcast_in_dim, tetapi bentuk hasilnya ditetapkan secara dinamis melalui output_dimensions.

Operasi tersebut juga menerima atribut opsional known_expanding_dimensions, known_nonexpanding_dimensions untuk menyatakan pengetahuan statis tentang perilaku dimensi yang diperluas. Jika tidak ditentukan, semua dimensi diasumsikan dapat diperluas.

Input

Output

Batasan

• (C1) element_type(result) diberikan oleh:
  • element_type(operand), jika !is_per_axis_quantized(operand).
  • element_type(operand) kecuali bahwa quantization_dimension(operand), scales(operand), dan zero_points(operand) mungkin berbeda dari quantization_dimension(result), scales(result), dan zero_points(result) masing-masing.
• (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_nonexpanding_dimensions).
• (C9) 0 <= known_expanding_dimensions < rank(operand).
• (C10) 0 <= known_nonexpanding_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_nonexpanding_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 op konvolusi, 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 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 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) Diberikan 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) ditentukan sebagai:
  • dim(lhs, input_batch_dimension) / batch_group_count if result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) jika result_dim = output_feature_dimension.
  • num_windows jika tidak, dengan:
  • 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 menggunakan tensor yang tidak dikuantisasi:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Jika operasi menggunakan tensor kuantisasi:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Jika is_per_axis_quantized(rhs), maka 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]]
//          ]]