مشخصات StableHLO

StableHLO مجموعه ای از عملیات برای عملیات سطح بالا (HLO) در مدل های یادگیری ماشین (ML) است. StableHLO به عنوان یک لایه قابل حمل بین چارچوب های مختلف ML و کامپایلرهای ML کار می کند: چارچوب های ML که برنامه های StableHLO را تولید می کنند با کامپایلرهای ML که برنامه های StableHLO را مصرف می کنند سازگار هستند.

هدف ما ساده سازی و تسریع توسعه ML با ایجاد قابلیت همکاری بیشتر بین چارچوب های مختلف ML (مانند TensorFlow، JAX و PyTorch) و کامپایلرهای ML (مانند XLA و IREE) است. برای این منظور، این سند مشخصات زبان برنامه نویسی StableHLO را ارائه می دهد.

این مشخصات شامل سه بخش عمده است. ابتدا، بخش برنامه‌ها ساختار برنامه‌های StableHLO را توصیف می‌کند که از توابع StableHLO تشکیل شده‌اند که خود از عملیات StableHLO تشکیل شده‌اند. در آن ساختار، بخش Ops معنای عملیات های فردی را مشخص می کند. بخش Execution معنایی را برای همه این عملیات ها که با هم در یک برنامه اجرا می شوند ارائه می دهد. در نهایت، بخش Notation در مورد نماد استفاده شده در تمام مشخصات بحث می کند.

برای مشاهده مشخصات نسخه قبلی StableHLO، مخزن را در نسخه مورد علاقه با برچسب باز کنید. به عنوان مثال، StableHLO v0.19.0 Spec . برای مشاهده تغییراتی که در هر برآمدگی جزئی نسخه StableHLO رخ داده است، به گزارش نسخه در VhloDialect.td مراجعه کنید.

برنامه ها

Program ::= {Func}

برنامه های StableHLO از تعداد دلخواه توابع StableHLO تشکیل شده اند. در زیر یک برنامه نمونه با تابع @main وجود دارد که دارای 3 ورودی ( %image ، %weights و %bias ) و 1 خروجی است. بدنه تابع دارای 6 عملیات است.

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

کارکرد

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

توابع StableHLO (که توابع نامگذاری شده نیز نامیده می شوند) دارای شناسه، ورودی/خروجی و بدنه هستند. در آینده، ما در حال برنامه ریزی برای معرفی ابرداده های اضافی برای توابع برای دستیابی به سازگاری بهتر با HLO هستیم ( #425 ، #626 ، #740 ، #744 ).

شناسه ها

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

شناسه های StableHLO مشابه شناسه ها در بسیاری از زبان های برنامه نویسی هستند، با دو ویژگی: 1) همه شناسه ها دارای علامت هایی هستند که انواع مختلف شناسه ها را متمایز می کنند، 2) شناسه های مقدار می توانند کاملاً عددی باشند تا تولید برنامه های StableHLO را ساده کنند.

انواع

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

انواع StableHLO به انواع ارزش (که انواع درجه یک نیز نامیده می شوند) دسته بندی می شوند که مقادیر StableHLO را نشان می دهند و انواع غیر ارزشی که عناصر دیگر برنامه را توصیف می کنند. انواع StableHLO مشابه انواع در بسیاری از زبان های برنامه نویسی است، با ویژگی اصلی آن طبیعت خاص دامنه StableHLO است که منجر به برخی از نتایج غیرعادی می شود (مثلاً انواع اسکالر نوع ارزش نیستند).

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

انواع تانسور نشان دهنده تانسورها، یعنی آرایه های چند بعدی هستند. آنها دارای یک شکل و یک نوع عنصر هستند، که در آن یک شکل اندازه ابعاد غیر منفی یا مجهول را به ترتیب صعودی ابعاد مربوطه (که محورها نیز نامیده می شوند) را نشان می دهد که از 0 تا R-1 شماره گذاری شده اند. به تعداد ابعاد R رتبه می گویند. به عنوان مثال، tensor<2x3xf32> یک نوع تانسور با شکل 2x3 و نوع عنصر f32 است. دارای دو بعد (یا به عبارتی دو محور) - بعد 0 و بعد 1 - که اندازه آنها 2 و 3 است. رتبه آن 2 است.

اشکال می توانند تا حدی یا کاملاً ناشناخته باشند (دینامیک)، به عنوان مثال tensor<?x2xf64> تا حدی ناشناخته است و tensor<?x?xf64> کاملاً ناشناخته است. اندازه های بعد پویا با استفاده از ? . شکل ها را نمی توان بدون رتبه بندی کرد.

در آینده، ما در حال برنامه ریزی برای کاوش در گسترش انواع تانسور فراتر از اندازه ابعاد و انواع عناصر هستیم، به عنوان مثال، شامل طرح بندی ( #629 ) و پراکندگی ( #1078 ).

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

انواع عناصر کوانتیزه مقادیر صحیح یک نوع ذخیره‌سازی را در محدوده از storage_min تا storage_max (شامل) نشان می‌دهند که با مقادیر ممیز شناور یک نوع بیان شده مطابقت دارد. برای یک مقدار صحیح معین i ، مقدار ممیز شناور مربوطه f می توان به صورت f = (i - zero_point) * scale محاسبه کرد، که در آن scale و zero_point پارامترهای کوانتیزه شدن نامیده می شوند. storage_min و storage_max در دستور زبان اختیاری هستند، اما مقادیر پیش‌فرض به ترتیب min_value(storage_type) و max_value(storage_type) هستند. انواع عناصر کوانتیزه دارای محدودیت های زیر هستند:

• (C1) type(storage_min) = storage_type .
• (C2) type(storage_max) = storage_type .
• (C3) min_value(storage_type) <= storage_min < storage_max <= max_value(storage_type) .
• (C4) type(scales...) = expressed_type .
• (C5) 0 < scales .
• (C6) is_finite(scales...) .
• (C7) storage_min <= zero_points <= storage_max .
• (C8) type(zero_points...) = storage_type .
• (C9) size(scales) = size(zero_points) .
• (C10) اگر is_empty(quantization_dimension) است، آنگاه size(scales) = 1 .
• (C11) 0 <= quantization_dimension .

در حال حاضر، QuantizationScale یک ثابت ممیز شناور است، اما علاقه زیادی به مقیاس‌های مبتنی بر اعداد صحیح وجود دارد که با ضرب‌کننده‌ها و شیفت‌ها نشان داده می‌شوند. ما در حال برنامه ریزی برای کشف این موضوع در آینده نزدیک هستیم ( #1404 ).

یک بحث مداوم در مورد معنایی QuantizationZeroPoint ، از جمله نوع، مقادیر و اینکه آیا می‌تواند فقط یک یا چند نقطه صفر در یک نوع تانسور کوانتیزه وجود داشته باشد، وجود دارد. بر اساس نتایج این بحث، مشخصات حول نقاط صفر ممکن است در آینده تغییر کند ( #1405 ).

بحث در حال انجام دیگر شامل معنایی QuantizationStorageMin و QuantizationStorageMax برای تعیین اینکه آیا محدودیتی باید بر روی این مقادیر و مقادیر تانسورهای کوانتیزه اعمال شود ( #1406 ) است.

در نهایت، ما در حال برنامه‌ریزی برای کشف مقیاس‌های ناشناخته و نقاط صفر هستیم، مشابه نحوه برنامه‌ریزی ما برای کاوش نمایش اندازه‌های ابعاد ناشناخته ( #1407 ).

انواع تانسور کوانتیزه نشان دهنده تانسورهایی با عناصر کوانتیزه است. این تانسورها دقیقاً مشابه تانسورهای معمولی هستند، با این تفاوت که عناصر آنها به جای انواع المان معمولی، دارای انواع عناصر کوانتیزه شده هستند.

در تانسورهای کوانتیزه، کوانتیزاسیون می‌تواند به ازای هر تانسور ، یعنی داشتن یک scale و zero_point برای کل تانسور یا می‌تواند در هر محور باشد، به این معنی که دارای scales متعدد و zero_points ، یک جفت در هر برش از یک بعد quantization_dimension خاص است. به طور رسمی تر، در یک تانسور t با کوانتیزاسیون هر محور، برش های dim(t, quantization_dimension) از quantization_dimension وجود دارد: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :] و غیره. همه عناصر در برش i scales[i] و zero_points[i] به عنوان پارامترهای کوانتیزاسیون خود استفاده می کنند. انواع تانسورهای کوانتیزه دارای محدودیت های زیر هستند:

• برای کوانتیزاسیون بر تانسور:
  • بدون محدودیت اضافی
• برای کوانتیزاسیون هر محور:
  • (C13) quantization_dimension < rank(self) .
  • (C14) dim(self, quantization_dimension) = size(scales) .

TokenType ::= 'token'

انواع توکن نشان دهنده توکن ها هستند، یعنی مقادیر مات تولید و مصرف شده توسط برخی عملیات. توکن ها برای اعمال دستور اجرا بر روی عملیات همانطور که در بخش اجرا توضیح داده شده استفاده می شود.

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

انواع تاپی نشان دهنده تاپل ها، یعنی لیست های ناهمگن هستند. Tuples یک ویژگی قدیمی است که فقط برای سازگاری با HLO وجود دارد. در HLO از تاپل ها برای نمایش ورودی ها و خروجی های متغیر استفاده می شود. در StableHLO، ورودی‌ها و خروجی‌های متغیر به صورت بومی پشتیبانی می‌شوند، و تنها استفاده از تاپل‌ها در StableHLO نمایش جامع HLO ABI است که در آن T , tuple<T> و tuple<tuple<T>> ممکن است بسته به یک پیاده‌سازی خاص از نظر مادی متفاوت باشند. . در آینده، ما در حال برنامه ریزی برای ایجاد تغییراتی در HLO ABI هستیم که ممکن است به ما اجازه دهد تا انواع تاپل را از StableHLO ( #598 ) حذف کنیم.

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

انواع عنصر عناصری از انواع تانسور را نشان می دهد. برخلاف بسیاری از زبان های برنامه نویسی، این نوع ها در StableHLO درجه یک نیستند. این بدان معنی است که برنامه های StableHLO نمی توانند مستقیماً مقادیر این نوع را نشان دهند (در نتیجه، نمایش مقادیر اسکالر نوع T با مقادیر تانسور 0 بعدی از نوع tensor<T> اصطلاحی است).

• نوع Boolean مقادیر بولی true و false را نشان می دهد.
• انواع عدد صحیح می توانند علامت دار ( si ) یا بدون علامت ( ui ) باشند و دارای یکی از عرض بیت های پشتیبانی شده ( 4 ، 8 ، 16 ، 32 یا 64 ) باشند. انواع siN علامت دار مقادیر صحیح را از -2^(N-1) تا 2^(N-1)-1 شامل، و انواع unsigned uiN مقادیر صحیح از 0 تا 2^N-1 را نشان می دهند.
• انواع ممیز شناور می تواند یکی از موارد زیر باشد:
  • انواع f8E4M3FN و f8E5M2 مربوط به کدگذاری های E4M3 و E5M2 فرمت FP8 که در فرمت های FP8 برای یادگیری عمیق توضیح داده شده است.
  • انواع f8E4M3FNUZ و f8E5M2FNUZ مربوط به کدگذاری های E4M3 و E5M2 فرمت های FP8 که در قالب های عددی 8 بیتی برای شبکه های عصبی عمیق توضیح داده شده است.
  • نوع f8E4M3B11FNUZ مربوط به کدگذاری E4M3 فرمت‌های FP8 که در آموزش و استنتاج هیبریدی 8 بیتی شناور نقطه شناور (HFP8) برای شبکه‌های عصبی عمیق توضیح داده شده است.
  • نوع bf16 مربوط به فرمت bfloat16 شرح داده شده در BFloat16: راز عملکرد بالا در Cloud TPU .
  • انواع f16 ، f32 و f64 مربوط به فرمت‌های binary16 ("نیمه دقت")، binary32 ("دقت تک") و binary64 ("دقت دوگانه") که در استاندارد IEEE 754 توضیح داده شده‌اند.
• انواع مختلط مقادیر پیچیده ای را نشان می دهند که یک قسمت واقعی و یک قسمت خیالی از همان نوع عنصر دارند. انواع پیچیده پشتیبانی شده complex<f32> (هر دو قسمت از نوع f32 هستند) و complex<f64> (هر دو قسمت از نوع f64 هستند).

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

انواع توابع هر دو توابع نامدار و ناشناس را نشان می دهند. آنها دارای انواع ورودی (لیست انواع در سمت چپ -> ) و انواع خروجی (لیست انواع در سمت راست -> ) هستند. در بسیاری از زبان های برنامه نویسی، انواع تابع درجه یک هستند، اما در StableHLO نه.

StringType ::= 'string'

نوع رشته نشان دهنده دنباله ای از بایت ها است. برخلاف بسیاری از زبان‌های برنامه‌نویسی، نوع رشته در StableHLO درجه یک نیست و فقط برای تعیین ابرداده استاتیک برای عناصر برنامه استفاده می‌شود.

عملیات

عملیات StableHLO (که به آنها ops نیز می گویند) مجموعه بسته ای از عملیات سطح بالا را در مدل های یادگیری ماشین نشان می دهد. همانطور که در بالا مورد بحث قرار گرفت، نحو StableHLO به شدت از MLIR الهام گرفته شده است، که لزوما ارگونومیک ترین جایگزین نیست، اما مسلما بهترین مناسب برای هدف StableHLO برای ایجاد قابلیت همکاری بیشتر بین چارچوب های ML و کامپایلرهای ML است.

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

عملیات StableHLO (که به آنها ops نیز می گویند) دارای یک نام، ورودی/خروجی و یک امضا هستند. نام از stablehlo. پیشوند و یادگاری که به طور منحصر به فرد یکی از عملیات های پشتیبانی شده را شناسایی می کند. برای یک لیست جامع از همه عملیات های پشتیبانی شده به زیر مراجعه کنید.

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

عملیات ها ورودی ها را مصرف می کنند و خروجی ها را تولید می کنند. ورودی ها به مقادیر ورودی (محاسبه شده در طول اجرا)، توابع ورودی (به صورت ایستا ارائه می شوند، زیرا در StableHLO توابع مقادیر درجه یک نیستند) و ویژگی های ورودی (همچنین به صورت ایستا ارائه می شوند) دسته بندی می شوند. نوع ورودی و خروجی های مصرف شده و تولید شده توسط یک op به حافظه آن بستگی دارد. برای مثال، add op 2 مقدار ورودی مصرف می کند و 1 مقدار خروجی تولید می کند. در مقایسه، گزینه select_and_scatter 3 مقدار ورودی، 2 تابع ورودی و 3 ویژگی ورودی مصرف می کند.

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

توابع ورودی (که توابع ناشناس نیز نامیده می شوند) بسیار شبیه توابع نامگذاری شده هستند با این تفاوت که: 1) آنها شناسه ندارند (از این رو نام "ناشناس")، 2) انواع خروجی را اعلام نمی کنند (انواع خروجی عبارتند از استنتاج از عملیات return در تابع).

نحو برای توابع ورودی شامل یک بخش استفاده نشده در حال حاضر (به تولید Unused در بالا مراجعه کنید) که برای سازگاری با MLIR وجود دارد. در MLIR، یک مفهوم کلی تر از "منطقه ها" وجود دارد که می تواند چندین "بلوک" از عملیات را از طریق عملیات پرش به هم متصل کند. این بلوک ها دارای شناسه هایی هستند که با تولید Unused مطابقت دارد، به طوری که می توان آنها را از یکدیگر متمایز کرد. StableHLO عملیات پرش ندارد، بنابراین بخش مربوطه از نحو MLIR استفاده نشده است (اما هنوز وجود دارد).

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

ویژگی های ورودی یک نام و یک مقدار دارند که یکی از ثابت های پشتیبانی شده است. آنها راه اصلی برای تعیین ابرداده استاتیک برای عناصر برنامه هستند. به عنوان مثال، concatenate op dimension ویژگی برای مشخص کردن بعدی که مقادیر ورودی آن در امتداد آن الحاق می شود، استفاده می کند. به طور مشابه، slice op از چندین ویژگی مانند start_indices و limit_indices برای تعیین محدوده هایی که برای برش مقدار ورودی استفاده می شود، استفاده می کند.

در حال حاضر، برنامه های StableHLO در طبیعت گاهی دارای ویژگی هایی هستند که در این سند توضیح داده نشده اند. در آینده، ما در حال برنامه ریزی برای جذب این ویژگی ها در opset StableHLO یا منع ظاهر شدن آنها در برنامه های StableHLO هستیم. در همین حال، لیست این ویژگی ها به شرح زیر است:

• layout ( #629 ).
• mhlo.frontend_attributes ( #628 ).
• mhlo.sharding ( #619 ).
• output_operand_aliases ( #740 ).
• فراداده مکان ( #594 ).

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

امضای op شامل انواع همه مقادیر ورودی (فهرست انواع در سمت چپ -> ) و انواع همه مقادیر خروجی (لیست انواع در سمت راست -> ) است. به بیان دقیق، انواع ورودی اضافی هستند، و انواع خروجی نیز تقریباً همیشه اضافی هستند (زیرا برای اکثر عملیات StableHLO، انواع خروجی را می توان از ورودی ها استنباط کرد). با این وجود، امضای op عمداً بخشی از نحو StableHLO برای سازگاری با MLIR است.

در زیر یک عملیات مثالی وجود دارد که حافظه آن select_and_scatter است. 3 مقدار ورودی ( %operand ، %source و %init_value )، 2 تابع ورودی و 3 ویژگی ورودی ( window_dimensions ، window_strides و padding ) مصرف می‌کند. توجه داشته باشید که چگونه امضای op فقط شامل انواع مقادیر ورودی آن می شود (اما نه انواع توابع ورودی و ویژگی هایی که به صورت درون خطی ارائه می شوند).

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

ثابت ها

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

ثابت های StableHLO دارای یک نوع حرف و یک نوع هستند که با هم یک مقدار StableHLO را نشان می دهند. به طور کلی، نوع بخشی از نحو ثابت است، به جز زمانی که بدون ابهام باشد (مثلاً یک ثابت بولی به طور واضح دارای نوع i1 است، در حالی که یک ثابت عدد صحیح می تواند چندین نوع ممکن داشته باشد).

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

ثابت های بولی مقادیر بولی true و false را نشان می دهند. ثابت های بولی دارای نوع i1 هستند.

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

ثابت های صحیح مقادیر صحیح را از طریق رشته هایی نشان می دهند که از نماد اعشاری یا هگزادسیمال استفاده می کنند. پایه های دیگر، به عنوان مثال باینری یا اکتال، پشتیبانی نمی شوند. ثابت های عدد صحیح دارای محدودیت های زیر هستند:

• (C1) is_wellformed(integer_literal, integer_type) .

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

ثابت های ممیز شناور مقادیر ممیز شناور را از طریق رشته هایی نشان می دهند که از نماد اعشاری یا علمی استفاده می کنند. علاوه بر این، نماد هگزادسیمال را می توان برای تعیین مستقیم بیت های زیرین در قالب ممیز شناور از نوع مربوطه استفاده کرد. ثابت های ممیز شناور دارای محدودیت های زیر هستند:

• (C1) اگر از نماد غیرهگزادسیمال استفاده شود، is_wellformed(float_literal, float_type) .
• (C2) اگر از نماد هگزادسیمال استفاده می شود، size(hexadecimal_digits) = num_bits(float_type) / 4 .

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

ثابت های مختلط مقادیر مختلط را با استفاده از لیست های یک قسمت واقعی (اول می آید) و یک قسمت خیالی (در درجه دوم) نشان می دهند. به عنوان مثال، (1.0, 0.0) : complex<f32> نشان دهنده 1.0 + 0.0i و (0.0, 1.0) : complex<f32> نشان دهنده 0.0 + 1.0i است. ترتیبی که در آن این قطعات سپس در حافظه ذخیره می شوند، پیاده سازی تعریف شده است. ثابت های مختلط دارای محدودیت های زیر هستند:

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

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

ثابت های تانسور مقادیر تانسور را با استفاده از لیست های تودرتو مشخص شده از طریق نماد NumPy نشان می دهند. برای مثال، dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> یک مقدار تانسور را با نگاشت زیر از شاخص ها به عناصر نشان می دهد: {0, 0} => 1 , {0, 1} => 2 ، {0, 2} => 3 ، {1, 0} => 4 ، {1, 1} => 5 ، {1, 2} => 6 . ترتیبی که این عناصر سپس در حافظه ذخیره می شوند به صورت پیاده سازی تعریف شده است. ثابت های تانسور دارای محدودیت های زیر هستند:

• (C1) has_syntax(tensor_literal, element_type(tensor_type)) ، که در آن:
  • has_syntax(element_literal: Syntax, element_type: Type) = is_wellformed(element_literal, type) .
  • has_syntax(tensor_literal: List, element_type: Type) = has_syntax(tensor_literal..., element_type) .
• (C2) has_shape(tensor_literal, shape(tensor_type)) ، که در آن:
  • has_shape(element_literal: Syntax, []) = true .
  • has_shape(tensor_literal: List, shape: List) = size(tensor_literal) = shape[0] and has_shape(tensor_literal..., shape[1:]) .
  • در غیر این صورت، false

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

ثابت‌های تانسور کوانتیزه‌شده مقادیر تانسور کوانتیزه‌شده را با استفاده از نمادهای ثابت تانسور نشان می‌دهند، با عناصری که به‌عنوان ثابت‌های نوع ذخیره‌شان مشخص می‌شوند. ثابت های تانسور کوانتیزه دارای محدودیت های زیر هستند:

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

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

لفظ رشته ای از بایت هایی تشکیل شده است که با استفاده از کاراکترهای ASCII و دنباله های فرار مشخص شده اند. آنها کدگذاری-آگنوستیک هستند، بنابراین تفسیر این بایت ها به صورت پیاده سازی تعریف شده است. حرفهای رشته ای دارای نوع string هستند.

عملیات

عضلات شکم

مفاهیم

عملیات abs از نظر عنصر را روی تانسور operand انجام می دهد و یک تانسور result تولید می کند. بسته به نوع عنصر، موارد زیر را انجام می دهد:

• برای اعداد صحیح امضا شده: مدول عدد صحیح.
• برای شناورها: abs از IEEE-754.
• برای اعداد مختلط: مدول مختلط.
• برای انواع کوانتیزه شده: dequantize_op_quantize(abs, operand, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

• (C1) shape(result) = shape(operand) .
• (C2) baseline_element_type(result) به صورت زیر تعریف می شود:
  • complex_element_type(element_type(operand)) if is_complex(operand) .
  • baseline_element_type(operand) در غیر این صورت.

مثال ها

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

نمونه های بیشتر

اضافه کردن

مفاهیم

جمع عنصری دو تانسور lhs و rhs را انجام می دهد و یک تانسور result تولید می کند. بسته به نوع عنصر، موارد زیر را انجام می دهد:

• برای Booleans: منطقی OR.
• برای اعداد صحیح: جمع اعداد صحیح.
• برای شناورها: addition از IEEE-754.
• برای اعداد مختلط: جمع مختلط.
• برای انواع کوانتیزه شده: dequantize_op_quantize(add, lhs, rhs, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

• اگر عملیات از تانسورهای غیر کوانتیزه استفاده می کند:
  • (C1) type(lhs) = type(rhs) = type(result) .
• اگر عملیات از تانسورهای کوانتیزه استفاده می کند:
  • (C2) is_quantized(lhs) and is_quantized(rhs) and is_quantized(result) .
  • (C3) storage_type(lhs) = storage_type(rhs) = storage_type(result) .
  • (C4) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C5) (is_per_axis_quantized(lhs) or is_per_axis_quantized(rhs)) = is_per_axis_quantized(result) .
  • (C6) اگر is_per_axis_quantized(lhs) ، سپس quantization_dimension(lhs) = quantization_dimension(result) .
  • (C7) اگر is_per_axis_quantized(rhs) ، سپس quantization_dimension(rhs) = quantization_dimension(result) .

مثال ها

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

نمونه های بیشتر

گذشته از همه اینها

مفاهیم

اطمینان حاصل می کند که عملیات تولید inputs قبل از هر عملیاتی که به result بستگی دارد اجرا می شود. اجرای این عملیات کاری انجام نمی دهد، فقط برای ایجاد وابستگی داده ها از result به inputs وجود دارد.

ورودی ها

خروجی ها

مثال ها

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

نمونه های بیشتر

همه_جمع آوری

مفاهیم

در هر گروه فرآیندی در شبکه فرآیند StableHLO، مقادیر تانسور operand از هر فرآیند در امتداد all_gather_dim به هم پیوسته و یک تانسور result تولید می‌کند.

این عملیات، شبکه فرآیند StableHLO را به process_groups تقسیم می‌کند که به صورت زیر تعریف می‌شوند:

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

پس از آن، در هر process_group :

• operands@receiver = [operand@sender for sender in process_group] برای همه receiver در process_group .
• result@process = concatenate(operands@process, all_gather_dim) برای همه process در process_group .

ورودی ها

خروجی ها

محدودیت ها

• (C1) 0 <= all_gather_dim < rank(operand) .
• (C2) is_unique(replica_groups) .
• size(replica_groups) به صورت زیر تعریف می شود:
  • num_replicas اگر cross_replica استفاده شود.
  • num_replicas اگر cross_replica_and_partition استفاده شود.
  • اگر flattened_ids استفاده شود، num_processes .
• (C4) 0 <= replica_groups < size(replica_groups) .
• (C5) اگر use_global_device_ids = true ، channel_id > 0 .
• (C6) type(result) = type(operand) به جز:
  • dim(result, all_gather_dim) = dim(operand, all_gather_dim) * dim(process_groups, 1) .

مثال ها

// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [[1, 2], [3, 4]]
// %operand@(1, 0): [[5, 6], [7, 8]]
%result = "stablehlo.all_gather"(%operand) {
  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<2x4xi64>
// %result@(0, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result@(1, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]

نمونه های بیشتر

all_reduce

مفاهیم

در هر گروه فرآیند در شبکه فرآیند StableHLO، یک computation تابع کاهشی را برای مقادیر تانسور operand از هر فرآیند اعمال می‌کند و یک تانسور result تولید می‌کند.

این عملیات، شبکه فرآیند StableHLO را به process_groups تقسیم می‌کند که به صورت زیر تعریف می‌شوند:

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

پس از آن، در هر process_group :

• result@process[result_index] = exec(schedule) برای برخی از schedule درخت دودویی که در آن:
  • exec(node) = computation(exec(node.left), exec(node.right)) .
  • exec(leaf) = leaf.value .
• schedule یک درخت باینری تعریف شده توسط پیاده سازی است که پیمایش ترتیبی آن to_destination_type(operands@process_group...[result_index], type(func_inputs(computation)[0])) است.

ورودی ها

خروجی ها

محدودیت ها

• (C1) is_unique(replica_groups) .
• size(replica_groups) به صورت زیر تعریف می شود:
  • num_replicas اگر cross_replica استفاده شود.
  • num_replicas اگر cross_replica_and_partition استفاده شود.
  • اگر flattened_ids استفاده شود، num_processes .
• (C3) 0 <= replica_groups < size(replica_groups) .
• (C4) اگر use_global_device_ids = true ، channel_id > 0 .
• (C5) computation دارای نوع (tensor<E>, tensor<E>) -> (tensor<E>) که در آن is_promotable(element_type(operand), E) .
• (C6) shape(result) = shape(operand) .
• (C7) element_type(result) = E

مثال ها

// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [1, 2, 3, 4]
// %operand@(1, 0): [5, 6, 7, 8]
%result = "stablehlo.all_reduce"(%operand) ({
  ^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_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor<4xi64>) -> tensor<4xi64>
// %result@(0, 0): [6, 8, 10, 12]
// %result@(1, 0): [6, 8, 10, 12]

نمونه های بیشتر

همه_به_همه

مفاهیم

در هر گروه فرآیندی در شبکه فرآیند StableHLO، مقادیر تانسور operand در امتداد split_dimension را به قطعات تقسیم می‌کند، قسمت‌های تقسیم شده را بین فرآیندها پراکنده می‌کند، قطعات پراکنده را در امتداد concat_dimension به هم متصل می‌کند و یک تانسور result تولید می‌کند.

این عملیات، شبکه فرآیند StableHLO را به process_groups تقسیم می‌کند که به صورت زیر تعریف می‌شوند:

• cross_replica(replica_groups) if channel_id <= 0 .
• cross_partition(replica_groups) اگر channel_id > 0 .

پس از آن، در هر process_group :

• split_parts@sender = split(operand@sender, split_count, split_dimension) برای همه sender در process_group .
• scattered_parts@receiver = [split_parts@sender[receiver_index] for sender in process_group] که در آن receiver_index = process_group.index(receiver) .
• result@process = concatenate(scattered_parts@process, concat_dimension) .

ورودی ها

خروجی ها

محدودیت ها

• (C1) 0 <= split_dimension < rank(operand) .
• (C2) dim(operand, split_dimension) % split_count = 0 .
• (C3) 0 <= concat_dimension < rank(operand) .
• (C4) 0 < split_count .
• (C5) is_unique(replica_groups) .
• size(replica_groups) به صورت زیر تعریف می شود:
  • num_replicas اگر cross_replica استفاده شود.
  • num_partitions اگر cross_partition استفاده شود.
• (C7) 0 <= replica_groups < size(replica_groups) .
• (C8) dim(replica_groups, 1) = split_count .
• (C9) type(result) = type(operand) به جز:
  • dim(result, split_dimension) = dim(operand, split_dimension) / split_count .
  • dim(result, concat_dimension) = dim(operand, concat_dimension) * split_count .

مثال ها

// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [[1, 2, 3, 4],
//                   [5, 6, 7, 8]]
// %operand@(1, 0): [[9, 10, 11, 12],
//                   [13, 14, 15, 16]]
%result = "stablehlo.all_to_all"(%operand) {
  split_dimension = 1 : i64,
  concat_dimension = 0 : i64,
  split_count = 2 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
} : (tensor<2x4xi64>) -> tensor<4x2xi64>
// %result@(0, 0): [[1, 2],
//                  [5, 6],
//                  [9, 10],
//                  [13, 14]]
// %result@(1, 0): [[3, 4],
//                  [7, 8],
//                  [11, 12],
//                  [15, 16]]

نمونه های بیشتر

و

مفاهیم

از نظر عنصر AND دو تانسور lhs و rhs را اجرا می کند و یک تانسور result تولید می کند. بسته به نوع عنصر، موارد زیر را انجام می دهد:

• برای Booleans: منطقی و.
• برای اعداد صحیح: بیتی AND.

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

آتان 2

مفاهیم

عملیات atan2 را از نظر عنصر روی تانسور lhs و rhs انجام می دهد و یک تانسور result تولید می کند. بسته به نوع عنصر، موارد زیر را انجام می دهد:

• برای شناورها: atan2 از IEEE-754.
• برای اعداد مختلط: atan2 مختلط.
• برای انواع کوانتیزه شده: dequantize_op_quantize(atan2, lhs, rhs, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

batch_norm_grad

مفاهیم

گرادیان های چندین ورودی batch_norm_training را که از grad_output پس انتشار می یابند محاسبه می کند و تانسورهای grad_operand ، grad_scale و grad_offset تولید می کند. به طور رسمی تر، این عملیات را می توان به صورت تجزیه به عملیات StableHLO موجود با استفاده از نحو Python به صورت زیر بیان کرد:

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

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

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

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

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

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

  return grad_operand, grad_scale, grad_offset

برای انواع کوانتیزه، dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, mean, variance, grad_output: batch_norm_grad(operand, scale, mean, variance, grad_output, epsilon, feature_index), operand, scale, mean, variance, grad_output, type(grad_operand), type(grad_scale), type(feature_index)) .

ورودی ها

خروجی ها

محدودیت ها

• (C1) 0 <= feature_index < rank(operand) .
• (C2) operand ، scale ، mean ، variance ، grad_output ، grad_operand ، grad_scale و grad_offset دارای همان baseline_element_type هستند.
• (C3) operand ، grad_output و grad_operand شکل یکسانی دارند.
• (C4) scale ، mean ، variance ، grad_scale و grad_offset شکل یکسانی دارند.
• (C5) size(scale) = dim(operand, feature_index) .

مثال ها

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

دسته_هنجار_استنتاج

مفاهیم

تانسور operand در تمام ابعاد به جز بعد feature_index عادی می کند و یک تانسور result تولید می کند. به طور رسمی تر، این عملیات را می توان به صورت تجزیه به عملیات StableHLO موجود با استفاده از نحو Python به صورت زیر بیان کرد:

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

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

برای انواع کوانتیزه dequantize_op_quantize(lambda operand, scale, offset, mean, variance: batch_norm_inference(operand, scale, offset, mean, variance, epsilon, feature_index), operand, scale, offset, mean, variance, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

• (C1) 0 <= feature_index < rank(operand) .
• (C2) operand ، scale ، offset ، mean ، variance و result یکسان baseline_element_type دارند.
• (C3) size(scale) = dim(operand, feature_index) .
• (C4) size(offset) = dim(operand, feature_index) .
• (C5) size(mean) = dim(operand, feature_index) .
• (C6) size(variance) = dim(operand, feature_index) .
• (C7) baseline_type(operand) = baseline_type(result) .

مثال ها

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

آموزش هنجارهای دسته ای

مفاهیم

میانگین و واریانس را در تمام ابعاد به جز بعد feature_index محاسبه می کند و تانسور operand تولید کننده output ، batch_mean و batch_var را عادی می کند. به طور رسمی تر، این عملیات را می توان به صورت تجزیه به عملیات StableHLO موجود با استفاده از نحو Python به صورت زیر بیان کرد:

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

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

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

برای انواع کوانتیزه، dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, offset: batch_norm_training(operand, scale, offset, epsilon, feature_index), operand, scale, offset, type(output), type(batch_mean), type(batch_var)) انجام می دهد.

ورودی ها

خروجی ها

محدودیت ها

• (C1) 0 <= feature_index < rank(operand) .
• (C2) operand ، scale ، offset ، batch_mean ، batch_var و output baseline_element_type یکسان دارند.
• (C3) size(scale) = dim(operand, feature_index) .
• (C4) size(offset) = dim(operand, feature_index) .
• (C5) size(batch_mean) = dim(operand, feature_index) .
• (C6) size(batch_var) = dim(operand, feature_index) .
• (C7) baseline_type(output) = baseline_type(operand) .

مثال ها

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

bitcast_convert

مفاهیم

یک عملیات بیت‌کست را روی تانسور operand انجام می‌دهد و یک تانسور result تولید می‌کند که در آن بیت‌های کل تانسور operand با استفاده از نوع تانسور result تفسیر مجدد می‌شوند.

به طور رسمی تر، با توجه به E = element_type(operand) ، E' = element_type(result) ، و R = rank(operand) :

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

bits نمایش درون حافظه یک مقدار داده شده را برمی گرداند، و رفتار آن به صورت پیاده سازی تعریف شده است، زیرا نمایش دقیق تانسورها توسط پیاده سازی تعریف شده است، و نمایش دقیق انواع عناصر نیز به صورت پیاده سازی تعریف شده است.

ورودی ها

خروجی ها

محدودیت ها

• (C1) با توجه به E = is_quantized(operand) ? storage_type(operand) : element_type(operand) , E' = is_quantized(result) ? storage_type(result) : element_type(result) ، و R = rank(operand) :
  • اگر num_bits(E') = num_bits(E) ، shape(result) = shape(operand) .
  • اگر num_bits(E') < num_bits(E) :
  • rank(result) = R + 1 .
  • dim(result, i) = dim(operand, i) برای همه 0 <= i < R .
  • dim(result, R) * num_bits(E') = num_bits(E) .
  • اگر num_bits(E') > num_bits(E) :
  • rank(result) = R - 1 .
  • dim(result, i) = dim(operand, i) برای همه 0 <= i < R .
  • dim(operand, R - 1) * num_bits(E) = num_bits(E') .
• (C2) اگر is_complex(operand) or is_complex(result) ، پس is_complex(operand) and is_complex(result) .

مثال ها

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

نمونه های بیشتر

broadcast_in_dim

مفاهیم

ابعاد و/یا رتبه یک تانسور ورودی را با کپی کردن داده ها در تانسور operand گسترش می دهد و یک تانسور result تولید می کند. به طور رسمی تر، result[result_index] = operand[operand_index] که در آن برای همه d در axes(operand) :

• operand_index[d] = 0 اگر dim(operand, d) = 1 .
• operand_index[d] = result_index[broadcast_dimensions[d]] در غیر این صورت.

ورودی ها

خروجی ها

محدودیت ها

• (C1) element_type(result) توسط:
  • element_type(operand) ، اگر !is_per_axis_quantized(operand) .
  • element_type(operand) به استثنای آن quantization_dimension(operand) ، scales(operand) و zero_points(operand) ممکن است با quantization_dimension(result) ، scales(result) و zero_points(result) احترام باشد ، در غیر این صورت متفاوت است.
• (C2) size(broadcast_dimensions) = rank(operand) .
• (C3) 0 <= broadcast_dimensions < rank(result) .
• (C4) is_unique(broadcast_dimensions) .
• (C5) برای همه d در axes(operand) :
  • dim(operand, d) = 1 یا
  • dim(operand, d) = dim(result, broadcast_dimensions[d]) .
• (c6) اگر is_per_axis_quantized(result) :
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)] .
  • اگر dim(operand, quantization_dimension(operand)) = 1 ، سپس scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))) .

مثال ها

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

نمونه های بیشتر

مورد

مفاهیم

بسته به مقدار index ، خروجی را از اجرای دقیقاً یک عملکرد از branches تولید می کند. به طور رسمی ، result = selected_branch() که در آن:

• selected_branch = branches[index] اگر 0 <= index < size(branches) .
• selected_branch = branches[-1] در غیر این صورت.

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

cbrt

مفاهیم

عملکرد ریشه مکعب عناصر را بر روی تانسور operand انجام می دهد و تانسور result تولید می کند. بسته به نوع عنصر ، موارد زیر را انجام می دهد:

• برای شناورها: rootn(x, 3) از IEEE-754.
• برای اعداد پیچیده: ریشه مکعب پیچیده.
• برای انواع کمیت: dequantize_op_quantize(cbrt, operand, type(result))

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

سقف

مفاهیم

مجاری عناصر عناصر تانسور operand را انجام می دهد و یک تانسور result تولید می کند. عملکرد roundToIntegralTowardPositive را از مشخصات IEEE-754 پیاده سازی می کند. برای انواع کم ، dequantize_op_quantize(ceil, operand, type(result)) را انجام می دهد.

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

بی پروا

مفاهیم

تجزیه چولسکی یک دسته از ماتریس ها را محاسبه می کند.

به طور رسمی تر ، برای همه i در index_space(result) ، result[i0, ..., iR-3, :, :] یک تجزیه چولسکی از a[i0, ..., iR-3, :, :] ] به شکل هر یک از ماتریس های مثلثی پایین تر (اگر lower true ) یا مثلثی فوقانی (اگر lower false باشد). مقادیر خروجی در مثلث مخالف ، یعنی مثلث فوقانی بالا یا مثلث پایین سخت به طور متناوب ، اجرا شده است.

اگر وجود داشته باشد i در جایی که ماتریس ورودی یک ماتریس دینهای مثبت هرمیتی نیست ، رفتار تعریف نشده است.

برای انواع کمیت ، dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)) را انجام می دهد.

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

گیره

مفاهیم

هر عنصر از تانسور operand را بین حداقل و حداکثر مقدار گیره می کند و تانسور result تولید می کند. به طور رسمی تر ، result[result_index] = minimum(maximum(operand[result_index], min_element), max_element) ، جایی که min_element = rank(min) = 0 ? min[] : min[result_index] ، max_element = rank(max) = 0 ? max[] : max[result_index] . برای انواع کم ، dequantize_op_quantize(clamp, min, operand, max, type(result)) را انجام می دهد.

تحمیل سفارش بر روی اعداد پیچیده شامل معناشناسی شگفت آور است ، بنابراین در آینده قصد داریم پشتیبانی از شماره های پیچیده را برای این عملیات حذف کنیم ( شماره 560 ).

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

مجموعه_

مفاهیم

در هر گروه فرآیند در شبکه فرآیند StableHlo ، مقدار تانسور operand را از فرآیند منبع به فرآیندهای هدف ارسال کرده و یک تانسور result تولید می کند.

این عملیات شبکه فرآیند StableHlo را به process_groups تقسیم می کند که به شرح زیر تعریف شده است:

• cross_replica(replica_groups) اگر channel_id <= 0 .
• cross_partition(replica_groups) اگر channel_id > 0 .

پس از آن ، result@process توسط:

• operand@process_groups[i, 0] اگر یک i وجود داشته باشد به گونه ای که روند در process_groups[i] .
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) در غیر این صورت.

ورودی ها

خروجی ها

محدودیت ها

• (C1) is_unique(replica_groups) .
• (c2) 0 <= replica_groups < N که در آن N به عنوان تعریف شده است:
  • num_replicas اگر از cross_replica استفاده می شود.
  • num_partitions اگر از cross_partition استفاده می شود.
• (C3) type(result) = type(operand) .

مثال ها

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

مجموعه_پرموت

مفاهیم

در هر گروه فرآیند در شبکه فرآیند StableHlo ، مقدار تانسور operand را از فرآیند منبع به فرآیند هدف می فرستد و تانسور result تولید می کند.

این عملیات شبکه فرآیند StableHlo را به process_groups تقسیم می کند که به شرح زیر تعریف شده است:

• cross_replica(source_target_pairs) اگر channel_id <= 0 .
• cross_partition(source_target_pairs) اگر channel_id > 0 .

پس از آن ، result@process توسط:

• operand@process_groups[i, 0] ، اگر یک i به گونه ای وجود داشته باشد که این process_groups[i, 1] = process .
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) در غیر این صورت.

ورودی ها

خروجی ها

محدودیت ها

• (C1) dim(source_target_pairs, 1) = 2 .
• (C2) is_unique(source_target_pairs[:, 0]) .
• (C3) is_unique(source_target_pairs[:, 1]) .
• (C4) 0 <= source_target_pairs < N ، جایی که N به صورت تعریف شده است:
  • num_replicas اگر از cross_replica استفاده می شود.
  • num_partitions اگر از cross_partition استفاده می شود.
• (c5) type(result) = type(operand) .

مثال ها

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

نمونه های بیشتر

مقایسه کنید

مفاهیم

مقایسه عناوین عناصر تنشور lhs و rhs را با توجه به comparison_direction و compare_type انجام می دهد و تانسور result تولید می کند.

مقادیر comparison_direction و compare_type دارای معناشناسی زیر است:

برای انواع عناصر بولی و عدد صحیح:

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

برای انواع عنصر نقطه شناور با compare_type = FLOAT ، OP عملیات IEEE-754 زیر را پیاده سازی می کند:

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

برای انواع عنصر نقطه شناور با compare_type = TOTALORDER ، OP از ترکیبی از totalOrder و عملیات compareQuietEqual از IEEE-754 استفاده می کند.

برای انواع عناصر پیچیده ، مقایسه واژگانی از جفت (real, imag) با استفاده از comparison_direction ارائه شده و compare_type انجام می شود. تحمیل سفارش بر روی اعداد پیچیده شامل معناشناسی غافلگیرکننده است ، بنابراین در آینده قصد داریم پشتیبانی از اعداد پیچیده را حذف کنیم وقتی comparison_direction GE ، GT ، LE یا LT ( #560 ) است.

برای انواع کمکی dequantize_compare(lhs, rhs, comparison_direction) .

ورودی ها

خروجی ها

محدودیت ها

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs) .
• (C2) shape(lhs) = shape(rhs) = shape(result) .
• (C3) compare_type به صورت تعریف شده است:
  • اگر is_signed_integer(element_type(lhs)) SIGNED .
  • UNSIGNED اگر is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)) .
  • FLOAT یا TOTALORDER اگر is_float(element_type(lhs)) .
  • FLOAT is_complex(element_type(lhs)) .

مثال ها

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

نمونه های بیشتر

مجتمع

مفاهیم

تبدیل عنصر عاقلانه به یک مقدار پیچیده از یک جفت از مقادیر واقعی و خیالی ، lhs و rhs را انجام می دهد و یک تانسور result تولید می کند.

ورودی ها

خروجی ها

محدودیت ها

• (C1) type(lhs) = type(rhs) .
• (C2) shape(result) = shape(lhs) .
• (C3) element_type(result) دارای نوع complex<E> است که در آن E = element_type(lhs) .

مثال ها

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

نمونه های بیشتر

کامپوزیت

مفاهیم

عملیات ساخته شده (تشکیل شده) از سایر عملیات StableHlo را محاصره می کند ، از inputs و composite_attributes استفاده می کند و results تولید می کند. معناشناسی OP توسط ویژگی decomposition اجرا می شود. composite OP را می توان با تجزیه آن بدون تغییر معانی برنامه جایگزین کرد. در مواردی که در حال تجزیه تجزیه ، همان معانی OP را ارائه نمی دهد ، استفاده از custom_call ترجیح می دهید.

قسمت version (پیش فرض 0 ) برای نشان دادن هنگام تغییر معانی کامپوزیت استفاده می شود.

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

به هم پیوستن

مفاهیم

inputs در dimension بعد به همان ترتیب آرگومان های داده شده جمع می کند و یک تانسور result تولید می کند. به طور رسمی ، result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1] ، کجا:

1. id = d0 + ... + dk-1 + kd .
2. d برابر با dimension است ، و d0 ، ... اندازه ابعاد inputs d .

ورودی ها

خروجی ها

محدودیت ها

• (C1) same(element_type(inputs...)) .
• (C2) same(shape(inputs...)) به جز dim(inputs..., dimension) .
• (C3) 0 < size(inputs) .
• (C4) 0 <= dimension < rank(inputs[0]) .
• (c5) element_type(result) = element_type(inputs[0]) .
• (c6) shape(result) = shape(inputs[0]) به جز:
  • dim(result, dimension) = dim(inputs[0], dimension) + ...

مثال ها

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

نمونه های بیشتر

ثابت

مفاهیم

یک تانسور output از یک value ثابت تولید می کند.

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

تبدیل

مفاهیم

تبدیل عناصر عاقلانه از یک عنصر به نوع دیگر در تانسور operand را انجام می دهد و یک تانسور result تولید می کند.

برای تبدیل های بولی به هر نوع پشتیبانی ، مقدار false به صفر تبدیل می شود و مقدار true به یک تبدیل می شود. برای تبدیل های هر نوع پشتیبانی از نوع به بولایی ، یک مقدار صفر به false تبدیل می شود و مقادیر غیر صفر به true تبدیل می شوند. در زیر نحوه عملکرد این کار برای انواع پیچیده را مشاهده کنید.

برای تبدیل های مربوط به عدد صحیح ، عدد صحیح به شناور یا نقطه شناور نقطه به شناور ، اگر مقدار منبع دقیقاً در نوع مقصد نشان داده شود ، مقدار نتیجه آن نمایش دقیق است. در غیر این صورت ، رفتار TBD است ( شماره 180 ).

برای تبدیل های مربوط به شناور-نقطه به دین ، ​​قسمت کسری کوتاه می شود. اگر مقدار کوتاه شده در نوع مقصد قابل نمایش نباشد ، رفتار TBD است ( #180 ).

تبدیل شامل پیچیده به پیچیده ، از همان رفتار تبدیل های نقطه شناور به شناور برای تبدیل قطعات واقعی و خیالی پیروی می کند.

برای تبدیل های پیچیده به هر نوع و هر نوع دیگری از نوع دیگر ، ارزش تخیل منبع به ترتیب نادیده گرفته می شود یا مقدار خیالی مقصد به ترتیب صفر می شود. تبدیل قسمت واقعی از تبدیل نقطه شناور پیروی می کند.

در اصل ، این عملیات می تواند بیان مجدد (تبدیل از تنسورهای کمیت به تانسرهای منظم) ، کمیت (تبدیل از تنسورهای منظم به تنسورهای کمیت) و نیاز (تبدیل بین تنسورهای کمیت) ، اما در حال حاضر ما عملیات اختصاصی برای آن داریم - یکنواخت برای آن - uniform_dequantize برای آن اولین مورد استفاده و uniform_quantize برای موارد استفاده دوم و سوم. در آینده ، این دو گزینه ممکن است به convert ( شماره 1576 ) ادغام شوند.

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

پیچیدگی

مفاهیم

محصولات DOT را بین ویندوز lhs و برش های rhs محاسبه می کند و result تولید می کند. نمودار زیر نشان می دهد که چگونه عناصر در result از lhs و rhs با استفاده از یک مثال بتونی محاسبه می شوند.

به طور رسمی تر ، تغییر زیر از ورودی ها را از نظر lhs در نظر بگیرید تا بتوانید ویندوز lhs بیان کنید:

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

این تغییر مکان از توابع یاور زیر استفاده می کند:

• lhs_shape(n, hw, c) = permute([n] + hw + [c], [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]) .
• result_shape(n1, hw, c1) = permute([n1] + hw + [c1], [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]) .
• permute([j0, j1, ..., jR-1], permutation) = [i0, i1, ..., iR-1] که در آن j[d] = i[permutation[d]] .

اگر feature_group_count = 1 و batch_group_count = 1 ، پس برای همه output_spatial_index در index_space(dim(result, output_spatial_dimensions...)) ، result[result_shape(:, output_spatial_index, :)] = dot_product در کجا:

• padding_value = constant(0, element_type(lhs)) .
• padded_lhs = pad(lhs, padding_value, lhs_padding[:, 0], lhs_padding[:, 1], lhs_base_dilations - 1) .
• lhs_window_start = lhs_shape(0, output_spatial_index, 0) * lhs_window_strides .
• lhs_window = slice(padded_lhs, lhs_window_start, lhs_window_start + lhs_window_dimensions, lhs_window_dilations) .
• reversed_lhs_window = reverse(lhs_window, [input_spatial_dimensions[dim] for dim in range(size(window_reversal)) if window_reversal[dim] = true]) . به نظر می رسد این ویژگی استفاده نشده است ، بنابراین در آینده قصد داریم آن را حذف کنیم ( شماره 1181 ).
• dot_product = dot_general(reversed_lhs_window, rhs, lhs_batching_dimensions=[], lhs_contracting_dimensions=input_spatial_dimensions + [input_feature_dimension], rhs_batching_dimensions=[], rhs_contracting_dimensions=kernel_spatial_dimensions + [kernel_input_feature_dimension]) .

اگر feature_group_count > 1 :

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

اگر batch_group_count > 1 :

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

برای انواع کمتری ، dequantize_op_quantize( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs, type(result)) .

For hybrid quantized types, performs hybrid_dequantize_then_op( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs) .

ورودی ها

خروجی ها

محدودیت ها

• (C1) N = rank(lhs) = rank(rhs) .
• (C2) size(window_strides) = N - 2 .
• (C3) 0 < window_strides .
• (C4) shape(padding) = [N - 2, 2] .
• (C5) size(lhs_dilation) = N - 2 .
• (C6) 0 < lhs_dilation .
• (C7) size(rhs_dilation) = N - 2 .
• (c8) 0 < rhs_dilation .
• (C9) size(window_reversal) = N - 2 .
• (C10) dim(lhs, input_batch_dimension) % batch_group_count = 0 .
• (C11) dim(lhs, input_feature_dimension) % feature_group_count = 0 .
• (C12) size(input_spatial_dimensions) = N - 2 .
• (C13) با توجه به input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension] :
  • is_unique(input_dimensions) .
  • 0 <= input_dimensions < N .
• dim(rhs, kernel_input_feature_dimension) = dim(lhs, input_feature_dimension) / feature_group_count
• (C15) dim(rhs, kernel_output_feature_dimension) % batch_group_count = 0 .
• (C16) dim(rhs, kernel_output_feature_dimension) % feature_group_count = 0 .
• (C17) size(kernel_spatial_dimensions) = N - 2 .
• (c18) با توجه به kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension] :
  • is_unique(kernel_dimensions) .
  • 0 <= kernel_dimensions < N .
• (C19) size(output_spatial_dimensions) = N - 2 .
• (c20) با توجه به output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension] :
  • is_unique(output_dimensions) .
  • 0 <= output_dimensions < N .
• (C21) 0 < feature_group_count .
• (C22) 0 < batch_group_count .
• (C23) feature_group_count = 1 or batch_group_count = 1 .
• (C24) size(precision_config) = 2 .
• (C25) dim(result, result_dim) به این صورت تعریف شده است:
  • dim(lhs, input_batch_dimension) / batch_group_count اگر result_dim = output_batch_dimension .
  • dim(rhs, kernel_output_feature_dimension) اگر result_dim = output_feature_dimension .
  • num_windows در غیر این صورت ، جایی که:
  • output_spatial_dimensions[spatial_dim] = result_dim .
  • lhs_dim = input_spatial_dimensions[spatial_dim] .
  • rhs_dim = kernel_spatial_dimensions[spatial_dim] .
  • dilated_input_shape[lhs_dim] = dim(lhs, lhs_dim) = 0 ? 0 : (dim(lhs, lhs_dim) - 1) * lhs_dilation[spatial_dim] + 1 .
  • padded_input_shape[lhs_dim] = padding[spatial_dim, 0] + dilated_input_shape[lhs_dim] + padding[spatial_dim, 1] .
  • dilated_window_shape[lhs_dim] = dim(rhs, rhs_dim) = 0 ? 0 : (dim(rhs, rhs_dim) - 1) * rhs_dilation[spatial_dim] + 1 .
  • is_empty_window[lhs_dim] = padded_input_shape[lhs_dim] = 0 || dilated_window_shape[lhs_dim] > padded_input_shape[lhs_dim] .
  • num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1 .
• (C26) rank(result) = N .
• اگر این عملیات از تانسور غیرقانونی استفاده می کند:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result) .
• در صورت استفاده از تانسور کمیته:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs) .
  • (C29) اگر is_per_axis_quantized(rhs) ، سپس quantization_dimension(rhs) = kernel_output_feature_dimension .
  • (C30) اگر is_per_axis_quantized(result) ، سپس quantization_dimension(result) = output_feature_dimension .
  • اگر is_quantized(lhs) :
  • (C31) storage_type(lhs) = storage_type(rhs) .
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C33) اگر is_per_tensor_quantized(rhs) ، سپس is_per_tensor_quantized(result) .
  • اگر !is_quantized(lhs) :
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result) .

مثال ها

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

نمونه های بیشتر

کسینوس

مفاهیم

عملکرد کنجین عناصر را بر روی تانسور operand انجام می دهد و تانسور result تولید می کند. بسته به نوع عنصر ، موارد زیر را انجام می دهد:

• برای شناورها: cos از IEEE-754.
• برای اعداد پیچیده: كسین پیچیده.
• برای انواع کمیت: dequantize_op_quantize(cosine, operand, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

count_leading_zeros

مفاهیم

تعداد عناصر عاقلانه از تعداد بیت های صفر پیشرو در تانسور operand را انجام می دهد و یک تانسور result تولید می کند.

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

CUSTION_CALL

مفاهیم

یک عملیات تعریف شده توسط اجرای call_target_name را که inputs را می گیرد و called_computations شود ، محاصره می کند و results می گیرد. has_side_effect ، backend_config و api_version ممکن است برای ارائه ابرداده تعریف شده اضافی استفاده شود.

در حال حاضر ، این عملیات شامل مجموعه ای نسبتاً بی نظم از ابرداده است که نشان دهنده تکامل ارگانیک عملکرد همتای خود در کامپایلر XLA است. در آینده ، ما قصد داریم این ابرداده را متحد کنیم ( شماره 741 ).

ورودی ها

خروجی ها

مثال ها

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

تقسیم کنید

مفاهیم

بخش عناصر عاقلانه از سود سهام lhs و تانسور Divisor rhs را انجام می دهد و تانسور result تولید می کند. بسته به نوع عنصر ، موارد زیر را انجام می دهد:

• برای عدد صحیح: تقسیم عدد صحیح که میزان جبری را با هر قسمت کسری دور انداخته می شود.
• برای شناورها: division از IEEE-754.
• برای اعداد پیچیده: تقسیم پیچیده.
• برای انواع کمیت:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

dot_general

مفاهیم

محصولات DOT را بین برش های lhs و برش های rhs محاسبه می کند و تانسور result تولید می کند.

به طور رسمی تر ، result[result_index] = dot_product ، جایی که:

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

برای انواع کمیت ، dequantize_op_quantize( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs, type(result)) .

برای انواع کمیت شده ترکیبی ، hybrid_dequantize_then_op( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs)

precision_config تجارت بین سرعت و دقت محاسبات در باکتری های شتاب دهنده را کنترل می کند. این می تواند یکی از موارد زیر باشد (در حال حاضر ، معناشناسی این مقادیر enum مشخص نشده است ، اما ما در حال برنامه ریزی برای پرداختن به این موضوع در شماره 755 هستیم):

• DEFAULT : سریعترین محاسبه ، اما حداقل تقریب دقیق به شماره اصلی.
• HIGH : محاسبه آهسته تر ، اما تقریب دقیق تر به شماره اصلی.
• HIGHEST : کمترین محاسبه ، اما دقیق ترین تقریب به تعداد اصلی.

ورودی ها

خروجی ها

محدودیت ها

• (C1) size(lhs_batching_dimensions) = size(rhs_batching_dimensions) .
• (C2) size(lhs_contracting_dimensions) = size(rhs_contracting_dimensions) .
• (C3) is_unique(lhs_batching_dimensions + lhs_contracting_dimensions) .
• (C4) is_unique(rhs_batching_dimensions + rhs_contracting_dimensions) .
• (c5) 0 <= lhs_batching_dimensions < rank(lhs) .
• (C6) 0 <= lhs_contracting_dimensions < rank(lhs) .
• (C7) 0 <= rhs_batching_dimensions < rank(rhs) .
• (c8) 0 <= rhs_contracting_dimensions < rank(rhs) .
• (C9) dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...) .
• (C10) dim(lhs, lhs_contracting_dimensions...) = dim(rhs, rhs_contracting_dimensions...) .
• (C11) size(precision_config) = 2 .
• (C12) shape(result) = dim(lhs, lhs_batching_dimensions) + dim(lhs, lhs_result_dimensions) + dim(rhs, rhs_result_dimensions) .
• اگر این عملیات از تانسور غیرقانونی استفاده می کند:
  • (C13) element_type(lhs) = element_type(rhs) .
• در صورت استفاده از تانسور کمیته:
  • (c14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs) .
  • (C15) zero_points(rhs) = 0 .
  • (C16) اگر is_per_axis_quantized(rhs) ، سپس quantization_dimension(rhs) در rhs_contracting_dimensions نیست.
  • اگر is_quantized(lhs) :
  • (C17) storage_type(lhs) = storage_type(rhs) .
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C19) اگر is_per_tensor_quantized(rhs) ، سپس is_per_tensor_quantized(result) .
  • اگر !is_quantized(lhs) :
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result) .

مثال ها

// %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>]
} : (tensor<2x2x2xi64>, tensor<2x2x2xi64>) -> tensor<2x2x2xi64>
// %result: [
//           [[1, 2],
//            [3, 4]],
//           [[5, 6],
//            [7, 8]]
//          ]

نمونه های بیشتر

dynamic_broadcast_in_dim

مفاهیم

این عمل از نظر عملکردی با Broadcast_in_dim OP یکسان است ، اما شکل نتیجه از طریق output_dimensions به صورت پویا مشخص می شود.

این عملیات همچنین ویژگی های اختیاری known_expanding_dimensions می پذیرد_کسپندینگ_دیمنس ، known_non_expanding_dimensions برای بیان دانش استاتیک در مورد رفتار در حال گسترش ابعاد. اگر مشخص نشده باشد ، فرض بر این است که همه ابعاد در حال گسترش هستند.

ورودی ها

خروجی ها

محدودیت ها

• (C1) element_type(result) توسط:
  • element_type(operand) ، اگر !is_per_axis_quantized(operand) .
  • element_type(operand) به استثنای این که quantization_dimension(operand) ، scales(operand) و zero_points(operand) ممکن است با quantization_dimension(result) ، scales(result) و zero_points(result) احترام باشد ، در غیر این صورت متفاوت است.
• (C2) size(broadcast_dimensions) = rank(operand) .
• (C3) 0 <= broadcast_dimensions < rank(result) .
• (C4) is_unique(broadcast_dimensions) .
• (C5) برای همه d در axes(operand) :
  • dim(operand, d) = 1 یا
  • dim(operand, d) = dim(result, broadcast_dimensions[d]) .
• (c6) اگر is_per_axis_quantized(result) :
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)] .
  • اگر dim(operand, quantization_dimension(operand)) = 1 ، سپس scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))) .
• (C7) size(output_dimensions) = rank(result) .
• is_unique(known_expanding_dimensions + known_non_expanding_dimensions)
• (c9) 0 <= known_expanding_dimensions < rank(operand) .
• (C10) 0 <= known_non_expanding_dimensions < rank(operand) .

مثال ها

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

نمونه های بیشتر

dynamic_conv

مفاهیم

این عمل از نظر عملکردی با Convolution OP یکسان است ، اما بالشتک به صورت پویا از طریق padding مشخص می شود.

ورودی ها

خروجی ها

محدودیت ها

• (C1) N = rank(lhs) = rank(rhs) .
• (C2) size(window_strides) = N - 2 .
• (C3) 0 < window_strides .
• (C4) shape(padding) = [N - 2, 2] .
• (C5) size(lhs_dilation) = N - 2 .
• (C6) 0 < lhs_dilation .
• (C7) size(rhs_dilation) = N - 2 .
• (c8) 0 < rhs_dilation .
• (C9) size(window_reversal) = N - 2 .
• (C10) dim(lhs, input_batch_dimension) % batch_group_count = 0 .
• (C11) dim(lhs, input_feature_dimension) % feature_group_count = 0 .
• (C12) size(input_spatial_dimensions) = N - 2 .
• (C13) با توجه به input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension] :
  • is_unique(input_dimensions) .
  • 0 <= input_dimensions < N .
• 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) Given 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) Given 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) is defined as:
  • dim(lhs, input_batch_dimension) / batch_group_count if result_dim = output_batch_dimension .
  • dim(rhs, kernel_output_feature_dimension) if result_dim = output_feature_dimension .
  • num_windows otherwise, where:
  • 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 .
• If the operation uses non-quantized tensors:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result) .
• If the operation uses quantized tensors:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs) .
  • (C29) If is_per_axis_quantized(rhs) , then quantization_dimension(rhs) = kernel_output_feature_dimension .
  • (C30) If is_per_axis_quantized(result) , then quantization_dimension(result) = output_feature_dimension .
  • If is_quantized(lhs) :
  • (C31) storage_type(lhs) = storage_type(rhs) .
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C33) If is_per_tensor_quantized(rhs) , then is_per_tensor_quantized(result) .
  • If !is_quantized(lhs) :
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result) .

مثال ها

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

نمونه های بیشتر

dynamic_gather

مفاهیم

This operation is functionally identical to gather op, with the slice_sizes specified dynamically as a value.

ورودی ها

خروجی ها

محدودیت ها

• (C1) rank(operand) = size(offset_dims) + size(collapsed_slice_dims) .
• (C2) 0 <= index_vector_dim <= rank(start_indices) .
• (C3) size(start_index_map) = index_vector_dim < rank(start_indices) ? dim(start_indices, index_vector_dim) : 1 .
• (C4) is_unique(offset_dims) and is_sorted(offset_dims) .
• (C5) 0 <= offset_dims < rank(result) .
• (C6) is_unique(collapsed_slice_dims) and is_sorted(collapsed_slice_dims) .
• (C7) 0 <= collapsed_slice_dims < rank(operand) .
• (C8) slice_sizes[collapsed_slice_dims...] <= 1 .
• (C9) is_unique(start_index_map) .
• (C10) 0 <= start_index_map < rank(operand) .
• (C11) size(slice_sizes) = rank(operand) .
• (C12) 0 <= slice_sizes <= shape(operand) .
• (C13) shape(result) = combine(batch_dim_sizes, offset_dim_sizes) where:
  • batch_dim_sizes = shape(start_indices) except that the dimension size of start_indices corresponding to index_vector_dim is not included.
  • offset_dim_sizes = shape(slice_sizes) except that the dimension sizes in slice_sizes corresponding to collapsed_slice_dims are not included.
  • combine puts batch_dim_sizes at axes corresponding to batch_dims and offset_dim_sizes at axes corresponding to offset_dims .
• (C14) element_type(operand) = element_type(result) .

مثال ها

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

نمونه های بیشتر

dynamic_iota

مفاهیم

This operation is functionally identical to iota op, but the result shape is specified dynamically via output_shape .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

dynamic_pad

مفاهیم

This operation is functionally identical to pad op, but with edge_padding_low , edge_padding_high , and interior_padding specified dynamically as values.

ورودی ها

خروجی ها

محدودیت ها

• (C1) element_type(operand) = element_type(padding_value) = element_type(result) .
• (C2) size(edge_padding_low) = size(edge_padding_high) = size(interior_padding) = rank(operand) .
• (C3) 0 <= interior_padding .
• (C4) shape(result) = shape(operand) + edge_padding_low + max(shape(operand) - 1, 0) * interior_padding + edge_padding_high .

مثال ها

// %operand: [
//            [1, 2, 3],
//            [4, 5, 6]
//           ]
// %padding_value: 0
// %edge_padding_low: [0, 1]
// %edge_padding_high: [2, 1]
// %interior_padding: [1, 2]
%result = "stablehlo.dynamic_pad"(%operand, %padding_value,
  %edge_padding_low, %edge_padding_high, %interior_padding
) : (tensor<2x3xi64>, tensor<i64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> tensor<5x9xi64>
// %result: [
//           [0, 1, 0, 0, 2, 0, 0, 3, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0],
//           [0, 4, 0, 0, 5, 0, 0, 6, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0]
//          ]

نمونه های بیشتر

dynamic_reshape

مفاهیم

This operation is functionally identical to reshape op, but the result shape is specified dynamically via output_shape .

ورودی ها

خروجی ها

محدودیت ها

• (C1) element_type(result) is given by:
  • element_type(operand) , if !is_per_axis_quantized(operand) .
  • element_type(operand) except that quantization_dimension(operand) and quantization_dimension(result) may differ, otherwise.
• (C2) size(operand) = size(result) .
• (C3) If is_per_axis_quantized(operand) :
  • reduce(dims(operand, [0, 1, ..., quantization_dimension(operand) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) = reduce(dims(result, [0, 1, ..., quantization_dimension(result) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) .
  • dim(operand, quantization_dimension(operand)) = dim(result, quantization_dimension(result)) .
  • reduce(dims(operand, [quantization_dimension(operand) + 1, ..., rank(operand) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) = reduce(dims(result, [quantization_dimension(result) + 1, ..., rank(result) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) .
• (C4) size(output_shape) = rank(result) .

مثال ها

// %operand: [[1, 2, 3], [4, 5, 6]]
// %output_shape: [3, 2]
%result = "stablehlo.dynamic_reshape"(%operand, %output_shape) : (tensor<2x3xi64>, tensor<2xi64>) -> tensor<3x2xi64>
// %result: [[1, 2], [3, 4], [5, 6]]

نمونه های بیشتر

dynamic_slice

مفاهیم

Extracts a slice from the operand using dynamically-computed starting indices and produces a result tensor. start_indices contain the starting indices of the slice for each dimension subject to potential adjustment, and slice_sizes contain the sizes of the slice for each dimension. More formally, result[result_index] = operand[operand_index] where:

• adjusted_start_indices = clamp(0, start_indices, shape(operand) - slice_sizes) .
• operand_index = adjusted_start_indices + result_index .

ورودی ها

خروجی ها

محدودیت ها

• (C1) element_type(operand) = element_type(result) .
• (C2) size(start_indices) = size(slice_sizes) = rank(operand) .
• (C3) same(type(start_indices...)) .
• (C4) 0 <= slice_sizes <= shape(operand) .
• (C5) shape(result) = slice_sizes .

مثال ها

// %operand: [
//            [0, 0, 1, 1],
//            [0, 0, 1, 1],
//            [0, 0, 0, 0],
//            [0, 0, 0, 0]
//           ]
// %start_indices0: -1
// %start_indices1: 3
%result = "stablehlo.dynamic_slice"(%operand, %start_indices0, %start_indices1) {
  slice_sizes = array<i64: 2, 2>
} : (tensor<4x4xi32>, tensor<i64>, tensor<i64>) -> tensor<2x2xi32>
// %result: [
//           [1, 1],
//           [1, 1]
//          ]

نمونه های بیشتر

dynamic_update_slice

مفاهیم

Produces a result tensor which is equal to the operand tensor except that the slice starting at start_indices is updated with the values in update . More formally, result[result_index] is defined as:

• update[update_index] if 0 <= update_index < shape(update) where:
  • adjusted_start_indices = clamp(0, start_indices, shape(operand) - shape(update)) .
  • update_index = result_index - adjusted_start_indices .
• operand[result_index] otherwise.

ورودی ها

خروجی ها

محدودیت ها

• (C1) type(operand) = type(result) .
• (C2) element_type(update) = element_type(operand) .
• (C3) rank(update) = rank(operand) .
• (C4) size(start_indices) = rank(operand) .
• (C5) same(type(start_indices...)) .
• (C6) 0 <= shape(update) <= shape(operand) .

مثال ها

// %operand: [
//            [1, 1, 0, 0],
//            [1, 1, 0, 0],
//            [1, 1, 1, 1],
//            [1, 1, 1, 1]
//           ]
// %update: [
//           [1, 1],
//           [1, 1]
//          ]
// %start_indices0: -1
// %start_indices1: 3
%result = "stablehlo.dynamic_update_slice"(%operand, %update, %start_indices0, %start_indices1)
  : (tensor<4x4xi32>, tensor<2x2xi32>, tensor<i64>, tensor<i64>) -> tensor<4x4xi32>
// %result: [
//           [1, 1, 1, 1],
//           [1, 1, 1, 1],
//           [1, 1, 1, 1],
//           [1, 1, 1, 1]
//          ]

نمونه های بیشتر

نمایی

مفاهیم

Performs element-wise exponential operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: exp from IEEE-754.
• For complex numbers: complex exponential.
• For quantized types: dequantize_op_quantize(exponential, operand, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

exponential_minus_one

مفاهیم

Performs element-wise exponential minus one operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: expm1 from IEEE-754.
• For complex numbers: complex exponential minus one.
• For quantized types: dequantize_op_quantize(exponential_minus_one, operand, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

// %operand: [0.0, 1.0]
%result = "stablehlo.exponential_minus_one"(%operand) : (tensor<2xf64>) -> tensor<2xf64>
// %result: [0.0, 1.71828187]

نمونه های بیشتر

fft

مفاهیم

Performs the forward and inverse Fourier transforms for real and complex inputs/outputs.

fft_type is one of the following:

• FFT : Forward complex-to-complex FFT.
• IFFT : Inverse complex-to-complex FFT.
• RFFT : Forward real-to-complex FFT.
• IRFFT : Inverse real-to-complex FFT (ie takes complex, returns real).

More formally, given the function fft which takes 1-dimensional tensors of complex types as input, produces 1-dimensional tensors of same types as output and computes the discrete Fourier transform:

For fft_type = FFT , result is defined as the final result of a series of L computations where L = size(fft_length) . For example, for L = 3 :

• result1[i0, ..., :] = fft(operand[i0, ..., :]) .
• result2[i0, ..., :, iR-1] = fft(result1[i0, ..., :, iR-1]) .
• result[i0, ..., :, iR-2, iR-1] = fft(result2[i0, ..., :, iR-2, iR-1]) .

Furthermore, given the function ifft which has the same type signature and computes the inverse of fft :

For fft_type = IFFT , result is defined as the inverse of the computations for fft_type = FFT . For example, for L = 3 :

• result1[i0, ..., :, iR-2, iR-1] = ifft(operand[i0, ..., :, iR-2, iR-1]) .
• result2[i0, ..., :, iR-1] = ifft(result1[i0, ..., :, iR-1]) .
• result[i0, ..., :] = ifft(result2[i0, ..., :]) .

Furthermore, given the function rfft which takes 1-dimensional tensors of floating-point types, produces 1-dimensional tensors of complex types of the same floating-point semantics and works as follows:

• rfft(real_operand) = truncated_result where
• complex_operand... = (real_operand..., 0.0) .
• complex_result = fft(complex_operand) .
• truncated_result = complex_result[:(rank(complex_result) / 2 + 1)] .

(When the discrete Fourier transform is computed for real operands, the first N/2 + 1 elements of the result unambiguously define the rest of the result, so the result of rfft is truncated to avoid computing redundant elements).

For fft_type = RFFT , result is defined as the final result of a series of L computations where L = size(fft_length) . For example, for L = 3 :

• result1[i0, ..., :] = rfft(operand[i0, ..., :]) .
• result2[i0, ..., :, iR-1] = fft(result1[i0, ..., :, iR-1]) .
• result[i0, ..., :, iR-2, iR-1] = fft(result2[i0, ..., :, iR-2, iR-1]) .

Finally, given the function irfft which has the same type signature and computes the inverse of rfft :

For fft_type = IRFFT , result is defined as the inverse of the computations for fft_type = RFFT . For example, for L = 3 :

• result1[i0, ..., :, iR-2, iR-1] = ifft(operand[i0, ..., :, iR-2, iR-1]) .
• result2[i0, ..., :, iR-1] = ifft(result1[i0, ..., :, iR-1]) .
• result[i0, ..., :] = irfft(result2[i0, ..., :]) .

ورودی ها

خروجی ها

محدودیت ها

• (C1) size(fft_length) <= rank(operand) .
• (C2) The relationship between operand and result element types varies:
  • If fft_type = FFT , element_type(operand) and element_type(result) have the same complex type.
  • If fft_type = IFFT , element_type(operand) and element_type(result) have the same complex type.
  • If fft_type = RFFT , element_type(operand) is a floating-point type and element_type(result) is a complex type of the same floating-point semantics.
  • If fft_type = IRFFT , element_type(operand) is a complex type and element_type(result) is a floating-point type of the same floating-point semantics.
• (C3) 1 <= size(fft_length) <= 3 .
• (C4) If among operand and result , there is a tensor real of a floating-point type, then shape(real)[-size(fft_length):] = fft_length .
• (C5) shape(result) = shape(operand) except for:
  • If fft_type = RFFT , dim(result, -1) = dim(operand, -1) = 0 ? 0 : dim(operand, -1) / 2 + 1 .
  • If fft_type = IRFFT , dim(operand, -1) = dim(result, -1) = 0 ? 0 : dim(result, -1) / 2 + 1 .

مثال ها

// %operand: [(1.0, 0.0), (0.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
%result = "stablehlo.fft"(%operand) {
  fft_type = #stablehlo<fft_type FFT>,
  fft_length = array<i64: 4>
} : (tensor<4xcomplex<f32>>) -> tensor<4xcomplex<f32>>
// %result: [(1.0, 0.0), (1.0, 0.0), (1.0, 0.0), (1.0, 0.0)]

کف

مفاهیم

Performs element-wise floor of operand tensor and produces a result tensor. Implements the roundToIntegralTowardNegative operation from the IEEE-754 specification. For quantized types, performs dequantize_op_quantize(floor, operand, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

جمع آوری

مفاهیم

Gathers slices from operand tensor from offsets specified in start_indices and produces a result tensor.

The following diagram shows how elements in result map on elements in operand using a concrete example. The diagram picks a few example result indices and explains in detail which operand indices they correspond to.

More formally, result[result_index] = operand[operand_index] where:

• batch_dims = [d for d in axes(result) and d not in offset_dims] .
• batch_index = result_index[batch_dims...] .
• start_index is defined as:
  • start_indices[bi0, ..., :, ..., biN] where bi are individual elements in batch_index and : is inserted at the index_vector_dim index, if index_vector_dim < rank(start_indices) .
  • [start_indices[batch_index]] otherwise.
• For d_operand in axes(operand) ,
  • full_start_index[d_operand] = clamp(start_index[d_start], 0, dim(operand, d_operand) - slice_sizes[d_operand]) if d_operand = start_index_map[d_start] .
  • full_start_index[d_operand] = 0 otherwise.
• For d_operand in axes(operand) ,
  • full_batching_index[d_operand] = batch_index[d_start - (d_start < index_vector_dim ? 0 : 1)] if d_operand = operand_batching_dims[i_batching] and d_start = start_indices_batching_dims[i_batching] .
  • full_batching_index[d_operand] = 0 otherwise.
• offset_index = result_index[offset_dims...] .
• full_offset_index = [oi0, ..., 0, ..., oiN] where oi are individual elements in offset_index , and 0 is inserted at indices from collapsed_slice_dims and operand_batching_dims .
• operand_index = full_start_index + full_batching_index + full_offset_index .

If indices_are_sorted is true then the implementation can assume that start_indices are sorted with respect to start_index_map , otherwise the behavior is undefined. More formally, for all i1 < i2 from indices(result) , full_start_index(i1) <= full_start_index(i2) .

ورودی ها

خروجی ها

محدودیت ها

• (C1) rank(operand) = size(offset_dims) + size(collapsed_slice_dims) + size(operand_batching_dims) .
• (C2) 0 <= index_vector_dim <= rank(start_indices) .
• (C3) size(start_index_map) = index_vector_dim < rank(start_indices) ? dim(start_indices, index_vector_dim) : 1 .
• (C4) is_unique(offset_dims) and is_sorted(offset_dims) .
• (C5) 0 <= offset_dims < rank(result) .
• (C6) is_unique(concatenate(collapsed_slice_dims, operand_batching_dims))
• (C7) is_sorted(collapsed_slice_dims) .
• (C8) 0 <= collapsed_slice_dims < rank(operand) .
• (C9) slice_sizes[collapsed_slice_dims...] <= 1 .
• (C10) is_sorted(operand_batching_dims) .
• (C11) 0 <= operand_batching_dims < rank(operand) .
• (C12) slice_sizes[operand_batching_dims...] <= 1 .
• (C13) is_unique(start_indices_batching_dims) .
• (C14) 0 <= start_indices_batching_dims < rank(start_indices) .
• (C15) index_vector_dim not in start_indices_batching_dims .
• (C16) size(operand_batching_dims) == size(start_indices_batching_dims) .
• (C17) dim(operand, operand_batching_dims...) = dim(start_indices, start_indices_batching_dims...) .
• (C18) is_unique(concatenate(start_index_map, operand_batching_dims)) .
• (C19) 0 <= start_index_map < rank(operand) .
• (C20) size(slice_sizes) = rank(operand) .
• (C21) 0 <= slice_sizes <= shape(operand) .
• (C22) shape(result) = combine(batch_dim_sizes, offset_dim_sizes) where:
  • batch_dim_sizes = shape(start_indices) except that the dimension size of start_indices corresponding to index_vector_dim is not included.
  • offset_dim_sizes = slice_sizes except that the dimension sizes in slice_sizes corresponding to collapsed_slice_dims and operand_batching_dims are not included.
  • combine puts batch_dim_sizes at axes corresponding to batch_dims and offset_dim_sizes at axes corresponding to offset_dims .
• (C23) element_type(operand) = element_type(result) .

مثال ها

// %operand: [
//            [
//             [[1, 2], [3, 4], [5, 6], [7, 8]],
//             [[9, 10],[11, 12], [13, 14], [15, 16]],
//             [[17, 18], [19, 20], [21, 22], [23, 24]]
//            ],
//            [
//             [[25, 26], [27, 28], [29, 30], [31, 32]],
//             [[33, 34], [35, 36], [37, 38], [39, 40]],
//             [[41, 42], [43, 44], [45, 46], [47, 48]]
//            ]
//           ]
// %start_indices: [
//                  [
//                   [[0, 0], [1, 0], [2, 1]],
//                   [[0, 1], [1, 1], [0, 9]]
//                  ],
//                  [
//                   [[0, 0], [2, 1], [2, 2]],
//                   [[1, 2], [0, 1], [1, 0]]
//                  ]
//                 ]
%result = "stablehlo.gather"(%operand, %start_indices) {
  dimension_numbers = #stablehlo.gather<
    offset_dims = [3, 4],
    collapsed_slice_dims = [1],
    operand_batching_dims = [0],
    start_indices_batching_dims = [1],
    start_index_map = [2, 1],
    index_vector_dim = 3>,
  slice_sizes = array<i64: 1, 1, 2, 2>,
  indices_are_sorted = false
} : (tensor<2x3x4x2xi32>, tensor<2x2x3x2xi64>) -> tensor<2x2x3x2x2xi32>
// %result: [
//           [
//            [
//             [[1, 2], [3, 4]],
//             [[3, 4], [5, 6]],
//             [[13, 14], [15, 16]]
//            ],
//            [
//             [[33, 34], [35, 36]],
//             [[35, 36], [37, 38]],
//             [[41, 42], [43, 44]]
//            ]
//           ],
//           [
//            [
//             [[1, 2], [3, 4]],
//             [[13, 14], [15, 16]],
//             [[21, 22], [23, 24]]
//            ],
//            [
//             [[43, 44], [45, 46]],
//             [[33, 34], [35, 36]],
//             [[27, 28], [29, 30]]
//            ]
//           ]
//          ]

نمونه های بیشتر

get_dimension_size

مفاهیم

Produces the size of the given dimension of the operand . More formally, result = dim(operand, dimension) . The Semantics concerns only with the shape component of the type. The element-type could be anything.

ورودی ها

خروجی ها

محدودیت ها

• (C1) 0 <= dimension < rank(operand) .

مثال ها

// %operand: [[1, 2, 3], [4, 5, 6]]
%result = "stablehlo.get_dimension_size"(%operand) {
  dimension = 1 : i64
} : (tensor<2x3xi64>) -> tensor<i32>
// %result: 3

نمونه های بیشتر

get_tuple_element

مفاهیم

Extracts element at index position of the operand tuple and produces a result . More formally, result = operand[index] .

ورودی ها

خروجی ها

محدودیت ها

• (C1) 0 <= index < size(operand) .
• (C2) type(result) = tuple_element_types(operand)[index] .

مثال ها

// %operand: ([1.0, 2.0], (3))
  index = 0 : i32
} : (tuple<tensor<2xf32>, tuple<tensor<i32>>>) -> tensor<2xf32>
// %result: [1.0, 2.0]

نمونه های بیشتر

اگر

مفاهیم

Produces the output from executing exactly one function from true_branch or false_branch depending on the value of pred . More formally, result = pred ? true_branch() : false_branch() .

ورودی ها

خروجی ها

محدودیت ها

• (C1) input_types(true_branch) = input_types(false_branch) = [] .
• (C2) output_types(true_branch) = output_types(false_branch) .
• (C3) type(results...) = output_types(true_branch) .

مثال ها

// %result_true_branch: 10
// %result_false_branch: 11
// %pred: true
%result = "stablehlo.if"(%pred) ({
  "stablehlo.return"(%result_true_branch) : (tensor<i32>) -> ()
}, {
  "stablehlo.return"(%result_false_branch) : (tensor<i32>) -> ()
}) : (tensor<i1>) -> tensor<i32>
// %result: 10

نمونه های بیشتر

تصویر

مفاهیم

Extracts the imaginary part, element-wise, from the operand and produces a result tensor. More formally, for each element x : imag(x) = is_complex(x) ? imaginary_part(x) : constant(0, element_type(result)) .

ورودی ها

خروجی ها

محدودیت ها

• (C1) shape(result) = shape(operand) .
• (C2) element_type(result) is defined as:
  • complex_element_type(element_type(operand)) if is_complex(operand) .
  • element_type(operand) otherwise.

مثال ها

// %operand: [(1.0, 2.0), (3.0, 4.0)]
%result = "stablehlo.imag"(%operand) : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
// %result: [2.0, 4.0]

نمونه های بیشتر

تغذیه

مفاهیم

Reads data from the infeed and produces results .

Semantics of infeed_config is implementation-defined.

results consist of payload values which come first and a token which comes last. In the future, we are planning to split the payload and the token into two separate outputs to improve clarity ( #670 ).

ورودی ها

خروجی ها

محدودیت ها

• (C1) 0 < size(results) .
• (C2) is_empty(result[:-1]) or is_tensor(type(results[:-1])) .
• (C3) is_token(type(results[-1])) .

مثال ها

// %token: !stablehlo.token
// infeed_queue[0]: [[1, 2], [3, 4]]
// infeed_queue[1]: [[5, 6], [7, 8]]
%results0:2 = "stablehlo.infeed"(%token) {
  infeed_config = ""
} : (!stablehlo.token) -> (tensor<2x2xi64>, !stablehlo.token)
// results0#0: [[1, 2], [3, 4]]
%results1:2 = "stablehlo.infeed"(%token) {
  infeed_config = ""
} : (!stablehlo.token) -> (tensor<2x2xi64>, !stablehlo.token)
// results1#0: [[5, 6], [7, 8]]

نمونه های بیشتر

یوتا

مفاهیم

Fills an output tensor with values in increasing order starting from zero along the iota_dimension dimension. به طور رسمی تر،

output[output_index] = constant(is_quantized(output) ? quantize(output_index[iota_dimension], element_type(output)) : output_index[iota_dimension], element_type(output)) .

ورودی ها

خروجی ها

محدودیت ها

• (C1) 0 <= iota_dimension < rank(output) .

مثال ها

%output = "stablehlo.iota"() {
  iota_dimension = 0 : i64
} : () -> tensor<4x5xi32>
// %output: [
//           [0, 0, 0, 0, 0],
//           [1, 1, 1, 1, 1],
//           [2, 2, 2, 2, 2],
//           [3, 3, 3, 3, 3]
//          ]

%output = "stablehlo.iota"() {
  iota_dimension = 1 : i64
} : () -> tensor<4x5xi32>
// %output: [
//           [0, 1, 2, 3, 4],
//           [0, 1, 2, 3, 4],
//           [0, 1, 2, 3, 4],
//           [0, 1, 2, 3, 4]
//          ]

نمونه های بیشتر

is_finite

مفاهیم

Performs element-wise check whether the value in x is finite (ie is neither +Inf, -Inf, nor NaN) and produces a y tensor. Implements the isFinite operation from the IEEE-754 specification. For quantized types, the result is always true .

ورودی ها

خروجی ها

محدودیت ها

• (C1) shape(x) = shape(y) .

مثال ها

// Logical values: -Inf, +Inf, NaN, ...
// %x: [0xFFF0000000000000, 0x7FF0000000000000, 0x7FF8000000000000, -10.0, -0.0, 0.0, 10.0]
%y = "stablehlo.is_finite"(%x) : (tensor<7xf64) -> tensor<7xi1>
// %y: [false, false, false, true, true, true, true]

نمونه های بیشتر

ورود به سیستم

مفاهیم

Performs element-wise logarithm operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: log from IEEE-754.
• For complex numbers: complex logarithm.
• For quantized types: dequantize_op_quantize(log, operand, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

// %operand: [[1.0, 2.0], [3.0, 4.0]]
%result = "stablehlo.log"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
// %result: [[0.0, 0.69314718055994529], [1.0986122886681098, 1.3862943611198906]]

نمونه های بیشتر

log_plus_one

مفاهیم

Performs element-wise logarithm plus one operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: logp1 from IEEE-754.
• For complex numbers: complex logarithm plus one.
• For quantized types: dequantize_op_quantize(log_plus_one, operand, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

// %operand: [0.0, -0.999, 7.0, 6.38905621, 15.0]
%result = "stablehlo.log_plus_one"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// %result: [0.0, -6.90776825, 2.07944155, 2.0, 2.77258873]

نمونه های بیشتر

لجستیکی

مفاهیم

Performs element-wise logistic operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: division(1, addition(1, exp(-x))) from IEEE-754.
• For complex numbers: complex logistic.
• For quantized types: dequantize_op_quantize(logistic, operand, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

// %operand: [[0.0, 1.0], [2.0, 3.0]]
%result = "stablehlo.logistic"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
// %result: [[0.5, 0.73105858], [0.88079708, 0.95257413]]

نمونه های بیشتر

نقشه

مفاهیم

Applies a map function computation to inputs along the dimensions and produces a result tensor.

More formally, result[result_index] = computation(inputs...[result_index]) .

ورودی ها

خروجی ها

محدودیت ها

• (C1) shape(inputs...) = shape(result) .
• (C2) 0 < size(inputs) = N .
• (C3) dimensions = range(rank(inputs[0])) .
• (C4) computation has type (tensor<E0>, ..., tensor<EN-1>) -> tensor<E'> where Ei = element_type(inputs[i]) and E' = element_type(result) .

مثال ها

// %input0: [[0, 1], [2, 3]]
// %input1: [[4, 5], [6, 7]]
%result = "stablehlo.map"(%input0, %input1) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = stablehlo.multiply %arg0, %arg1 : tensor<i64>
    stablehlo.return %0 : tensor<i64>
}) {
  dimensions = array<i64: 0, 1>
} : (tensor<2x2xi64>, tensor<2x2xi64>) -> tensor<2x2xi64>
// %result: [[0, 5], [12, 21]]

نمونه های بیشتر

بیشترین

مفاهیم

Performs element-wise max operation on tensors lhs and rhs and produces a result tensor. Depending on the element type, does the following:

• For booleans: logical OR.
• For integers: integer maximum.
• For floats: maximum from IEEE-754.
• For complex numbers: lexicographic maximum for the (real, imaginary) pair. Imposing an ordering on complex numbers involves surprising semantics, so in the future we are planning to remove support for complex numbers for this operation ( #560 ).
• For quantized types:
  • dequantize_op_quantize(maximum, lhs, rhs, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

کمترین

مفاهیم

Performs element-wise min operation on tensors lhs and rhs and produces a result tensor. Depending on the element type, does the following:

• For booleans: logical AND.
• For integers: integer minimum.
• For floats: minimum from IEEE-754.
• For complex numbers: lexicographic minimum for the (real, imaginary) pair. Imposing an ordering on complex numbers involves surprising semantics, so in the future we are planning to remove support for complex numbers for this operation ( #560 ).
• For quantized types:
  • dequantize_op_quantize(minimum, lhs, rhs, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

تکثیر کردن

مفاهیم

Performs element-wise product of two tensors lhs and rhs and produces a result tensor. Depending on the element type, does the following:

• For booleans: logical AND.
• For integers: integer multiplication.
• For floats: multiplication from IEEE-754.
• For complex numbers: complex multiplication.
• For quantized types:
  • dequantize_op_quantize(multiply, lhs, rhs, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

نفی کردن

مفاهیم

Performs element-wise negation of operand tensor and produces a result tensor. Depending on the element type, does the following:

• For signed integers: integer negation.
• For unsigned integers: bitcast to signed integer, integer negation, bitcast back to unsigned integer.
• For floats: negate from IEEE-754.
• For complex numbers: complex negation.
• For quantized types: dequantize_op_quantize(negate, operand, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

// Negation operation with integer Tensors
// %operand: [0, -2]
%result = "stablehlo.negate"(%operand) : (tensor<2xi32>) -> tensor<2xi32>
// %result: [0, 2]

// Negation operation with with complex tensors
// %operand: (2.5, 0.0)
%result = "stablehlo.negate"(%operand) : (tensor<1xcomplex<f32>>) -> tensor<1xcomplex<f32>>
// %result: [-2.5, -0.0]

نمونه های بیشتر

نه

مفاهیم

Performs element-wise NOT of tensor operand and produces a result tensor. Depending on the element type, does the following:

• For booleans: logical NOT.
• For integers: bitwise NOT.

استدلال ها

خروجی ها

محدودیت ها

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

مثال ها

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

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

نمونه های بیشتر

optimization_barrier

مفاهیم

Ensures that the operations that produce the operand are executed before any operations that depend on the result and prevents compiler transformations from moving operations across the barrier. Other than that, the operation is an identity, ie result = operand .

استدلال ها

خروجی ها

محدودیت ها

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

مثال ها

// %operand0: 0.0
// %operand1: 1.0
%result0, %result1 = "stablehlo.optimization_barrier"(%operand0, %operand1) : (tensor<f32>, tensor<f32>) -> (tensor<f32>, tensor<f32>)
// %result0: 0.0
// %result1: 1.0

نمونه های بیشتر

یا

مفاهیم

Performs element-wise OR of two tensors lhs and rhs and produces a result tensor. Depending on the element type, does the following:

• For booleans: logical OR.
• For integers: bitwise OR.

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

// Bitwise operation with with integer tensors
// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.or"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[5, 6], [7, 12]]

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

نمونه های بیشتر

غذا خوردن

مفاهیم

Writes inputs to the outfeed and produces a result token.

Semantics of outfeed_config is implementation-defined.

ورودی ها

خروجی ها

مثال ها

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

نمونه های بیشتر

پد

مفاهیم

Expands operand by padding around the tensor as well as between the elements of the tensor with the given padding_value .

edge_padding_low and edge_padding_high specify the amount of padding added at the low-end (next to index 0) and the high-end (next to the highest index) of each dimension respectively. The amount of padding can be negative, where the absolute value of negative padding indicates the number of elements to remove from the specified dimension.

interior_padding specifies the amount of padding added between any two elements in each dimension which may not be negative. Interior padding occurs before edge padding such that negative edge padding will remove elements from the interior-padded operand.

More formally, result[result_index] is defined as:

• operand[operand_index] if result_index = edge_padding_low + operand_index * (interior_padding + 1) .
• padding_value otherwise.

ورودی ها

خروجی ها

محدودیت ها

• (C1) element_type(operand) = element_type(padding_value) = element_type(result) .
• (C2) size(edge_padding_low) = size(edge_padding_high) = size(interior_padding) = rank(operand) .
• (C3) 0 <= interior_padding .
• (C4) shape(result) = shape(operand) + edge_padding_low + max(shape(operand) - 1, 0) * interior_padding + edge_padding_high .

مثال ها

// %operand: [
//            [1, 2, 3],
//            [4, 5, 6]
//           ]
// %padding_value: 0
%result = "stablehlo.pad"(%operand, %padding_value) {
  edge_padding_low = array<i64: 0, 1>,
  edge_padding_high = array<i64: 2, 1>,
  interior_padding = array<i64: 1, 2>
} : (tensor<2x3xi32>, tensor<i32>) -> tensor<5x9xi32>
// %result: [
//           [0, 1, 0, 0, 2, 0, 0, 3, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0],
//           [0, 4, 0, 0, 5, 0, 0, 6, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0]
//          ]

نمونه های بیشتر

partition_id

مفاهیم

Produces partition_id of the current process.

خروجی ها

مثال ها

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

نمونه های بیشتر

popcnt

مفاهیم

Performs element-wise count of the number of bits set in the operand tensor and produces a result tensor.

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

قدرت

مفاهیم

Performs element-wise exponentiation of lhs tensor by rhs tensor and produces a result tensor. Depending on the element type, does the following:

• For integers: integer exponentiation.
• For floats: pow from IEEE-754.
• For complex numbers: complex exponentiation.
• For quantized types: dequantize_op_quantize(power, lhs, rhs, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

// %lhs: [-2.0, -0.0, -36.0, 5.0, 3.0, 10000.0]
// %rhs: [2.0, 2.0, 1.1, 2.0, -1.0, 10.0]
%result = "stablehlo.power"(%lhs, %rhs) : (tensor<6xf64>, tensor<6xf64>) -> tensor<6xf64>
// %result: [4.0, 0.0, -nan, 25.0, 0.333333343, inf]

نمونه های بیشتر

واقعی

مفاهیم

Extracts the real part, element-wise, from the operand and produces a result tensor. More formally, for each element x : real(x) = is_complex(x) ? real_part(x) : x .

ورودی ها

خروجی ها

محدودیت ها

• (C1) shape(result) = shape(operand) .
• (C2) element_type(result) is defined as:
  • complex_element_type(element_type(operand)) if is_complex(operand) .
  • element_type(operand) otherwise.

مثال ها

// %operand: [(1.0, 2.0), (3.0, 4.0)]
%result = "stablehlo.real"(%operand) : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
// %result: [1.0, 3.0]

نمونه های بیشتر

recv

مفاهیم

Receives data from a channel with channel_id and produces results .

If is_host_transfer is true , then the operation transfers data from the host. Otherwise, it transfers data from another device. What this means is implementation-defined. This flag duplicates the information provided in channel_type , so in the future we are planning to only keep one of them ( #666 ).

results consist of payload values which come first and a token which comes last. In the future, we are planning to split the payload and the token into two separate outputs to improve clarity ( #670 ).

ورودی ها

خروجی ها

محدودیت ها

• (C1) channel_type is defined as:
  • HOST_TO_DEVICE if is_host_transfer = true ,
  • DEVICE_TO_DEVICE otherwise.
• (C2) 0 < size(results) .
• (C3) is_empty(result[:-1]) or is_tensor(type(results[:-1])) .
• (C4) is_token(type(results[-1])) .

مثال ها

%results0, %results1 = "stablehlo.recv"(%token) {
  channel_handle = #stablehlo.channel_handle<handle = 1, type = 3>,
  is_host_transfer = true
} : (!stablehlo.token) -> (tensor<2x2xi64>, !stablehlo.token)

نمونه های بیشتر

كاهش دادن

مفاهیم

Applies a reduction function body to inputs and init_values along the dimensions and produces results tensors.

The order of reductions is implementation-defined, which means that body and init_values must form a monoid to guarantee that the operation produces the same results for all inputs on all implementations. However, this condition doesn't hold for many popular reductions. Eg floating-point addition for body and zero for init_values don't actually form a monoid because floating-point addition is not associative.

More formally, results...[j0, ..., jR-1] = reduce(input_slices_converted) where:

• input_slices = inputs...[j0, ..., :, ..., jR-1] , where : are inserted at dimensions .
• input_slices_converted = to_destination_type(input_slices..., type(func_inputs(body)[:len(func_inputs(body))//2])...) .
• init_values_converted = to_destination_type(init_values..., type(func_inputs(body)[len(func_inputs(body))//2:])...) .
• reduce(input_slices_converted) = exec(schedule) for some binary tree schedule where:
  • exec(node) = body(exec(node.left), exec(node.right)) .
  • exec(leaf) = leaf.value .
• schedule is an implementation-defined full binary tree whose in-order traversal consists of:
  • input_slices_converted...[index] values, for all index in index_space(input_slices_converted) in the ascending lexicographic order of index .
  • Interspersed with an implementation-defined amount of init_values_converted at implementation-defined positions.

ورودی ها

خروجی ها

محدودیت ها

• (C1) same(shape(inputs...)) .
• (C2) element_type(inputs...) = element_type(init_values...) .
• (C3) 0 < size(inputs) = size(init_values) = size(results) = N .
• (C4) 0 <= dimensions < rank(inputs[0]) .
• (C5) is_unique(dimensions) .
• (C6) body has type (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>) where is_promotable(element_type(inputs[i]), Ei) .
• (C7) shape(results...) = shape(inputs...) except that the dimension sizes of inputs... corresponding to dimensions are not included.
• (C8) element_type(results[i]) = Ei for all i in [0,N) .

مثال ها

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

نمونه های بیشتر

reduce_precision

مفاهیم

Performs element-wise conversion of operand to another floating-point type that uses exponent_bits and mantissa_bits and back to the original floating-point type and produces an output tensor.

رسمی تر:

• The mantissa bits of the original value are updated to round the original value to the nearest value representable with mantissa_bits using roundToIntegralTiesToEven semantics.
• Then, if mantissa_bits are smaller than the number of mantissa bits of the original value, the mantissa bits are truncated to mantissa_bits .
• Then, if the exponent bits of the intermediate result don't fit into the range provided by exponent_bits , the intermediate result overflows to infinity using the original sign or underflows to zero using the original sign.
• For quantized types, performs dequantize_op_quantize( lambda operand: reduce_precision(operand, exponent_bits, mantissa_bits), operand, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

reduce_scatter

مفاهیم

Within each process group in the StableHLO process grid, performs reduction, using computations , over the values of the operand tensor from each process, splits the reduction result along scatter_dimension into parts, and scatters the split parts between the processes to produce the result .

The operation splits the StableHLO process grid into process_groups which is defined as follows:

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

Afterwards, within each process_group :

• reduced_value = all_reduce(operand, replica_groups, channel_id, use_global_device_ids, computation) .
• parts@sender = split(reduced_value@sender, dim(process_groups, 1), scatter_dimension) .
• result@receiver = parts@sender[receiver_index] for all sender in process_group , where receiver_index = process_group.index(receiver) .

ورودی ها

خروجی ها

محدودیت ها

• (C1) dim(operand, scatter_dimension) % dim(process_groups, 1) = 0 .
• (C2) 0 <= scatter_dimension < rank(operand) .
• (C3) is_unique(replica_groups) .
• (C4) size(replica_groups) is defined as:
  • num_replicas if cross_replica is used.
  • num_replicas if cross_replica_and_partition is used.
  • num_processes if flattened_ids is used.
• (C5) 0 <= replica_groups < size(replica_groups) .
• (C6) If use_global_device_ids = true , then channel_id > 0 .
• (C7) computation has type (tensor<E>, tensor<E>) -> (tensor<E>) where is_promotable(element_type(operand), E) .
• (C8) shape(result) = shape(operand) except:
  • dim(result, scatter_dimension) = dim(operand, scatter_dimension) / dim(process_groups, 1) .
• (C9) element_type(result) = E .

مثال ها

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

نمونه های بیشتر

reduce_window

مفاهیم

Applies a reduction function body to windows of inputs and init_values and produces results .

The following diagram shows how elements in results... are computed from inputs... using a concrete example.

More formally, results...[result_index] = reduce(windows, init_values, axes(inputs...), body) (see reduce ) where:

• padded_inputs = pad(inputs..., init_values..., padding[:, 0], padding[:, 1], base_dilations - 1) .
• window_start = result_index * window_strides .
• window_end = window_start + (window_dimensions - 1) * window_dilations + 1 .
• windows = slice(padded_inputs..., window_start, window_end, window_dilations) .

ورودی ها

خروجی ها

محدودیت ها

• (C1) 0 < size(inputs) = size(init_values) = size(results) = N .
• (C2) same(shape(inputs...)) .
• (C3) element_type(inputs...) = element_type(init_values...) .
• (C4) size(window_dimensions) = rank(inputs[0]) .
• (C5) 0 < window_dimensions .
• (C6) size(window_strides) = rank(inputs[0]) .
• (C7) 0 < window_strides .
• (C8) size(base_dilations) = rank(inputs[0]) .
• (C9) 0 < base_dilations .
• (C10) size(window_dilations) = rank(inputs[0]) .
• (C11) 0 < window_dilations .
• (C12) shape(padding) = [rank(inputs[0]), 2] .
• (C13) body has type (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>) where is_promotable(element_type(inputs[i]), Ei) .
• (C14) same(shape(results...)) .
• (C15) shape(results[0]) = num_windows where:
  • dilated_input_shape = shape(inputs[0]) = 0 ? 0 : (shape(inputs[0]) - 1) * base_dilations + 1 .
  • padded_input_shape = padding[:, 0] + dilated_input_shape + padding[:, 1] .
  • dilated_window_shape = (window_dimensions - 1) * window_dilations + 1 .
  • is_empty_window = padded_input_shape = 0 || dilated_window_shape > padded_input_shape .
  • num_windows = is_empty_window ? 0 : floor((padded_input_shape - dilated_window_shape) / window_strides) + 1 .
• (C16) element_type(results[i]) = Ei for all i in [0,N) .

مثال ها

// %input = [[1, 2], [3, 4], [5, 6]]
// %init_value = 0
%result = "stablehlo.reduce_window"(%input, %init_value) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  window_dimensions = array<i64: 2, 1>,
  window_strides = array<i64: 4, 1>,
  base_dilations = array<i64: 2, 1>,
  window_dilations = array<i64: 3, 1>,
  padding = dense<[[2, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<3x2xi64>, tensor<i64>) -> tensor<2x2xi64>
// %result = [[0, 0], [3, 4]]

نمونه های بیشتر

باقی مانده

مفاهیم

Performs element-wise remainder of dividend lhs and divisor rhs tensors and produces a result tensor.

More formally, the sign of the result is taken from the dividend, and the absolute value of the result is always less than the divisor's absolute value. The remainder is calculated as lhs - d * rhs , where d is given by:

• For integers: stablehlo.divide(lhs, rhs) .
• For floats: division(lhs, rhs) from IEEE-754 with rounding attribute roundTowardZero .
• For complex numbers: TBD ( #997 ).
• For quantized types:
  • dequantize_op_quantize(remainder, lhs, rhs, type(result)) .

For floating-point element types, this operation is in contrast with the remainder operation from IEEE-754 specification where d is an integral value nearest to the exact value of lhs/rhs with ties to even.

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

// %lhs: [17, -17, 17, -17]
// %rhs: [3, 3, -3, -3]
%result = "stablehlo.remainder"(%lhs, %rhs) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi64>
// %result: [2, -2, 2, -2]

نمونه های بیشتر

replica_id

مفاهیم

Produces replica_id of the current process.

خروجی ها

مثال ها

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

نمونه های بیشتر

تغییر شکل دادن

مفاهیم

Performs reshape of operand tensor to a result tensor. Conceptually, it amounts to keeping the same canonical representation but potentially changing the shape, eg from tensor<2x3xf32> to tensor<3x2xf32> or tensor<6xf32> .

More formally, result[result_index] = operand[operand_index] where result_index and operand_index have the same position in the lexicographic ordering of index_space(result) and index_space(operand) .

ورودی ها

خروجی ها

محدودیت ها

• (C1) element_type(result) is given by:
  • element_type(operand) , if !is_per_axis_quantized(operand) .
  • element_type(operand) except that quantization_dimension(operand) and quantization_dimension(result) may differ, otherwise.
• (C2) size(operand) = size(result) .
• (C3) If is_per_axis_quantized(operand) :
  • reduce(dims(operand, [0, 1, ..., quantization_dimension(operand) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) = reduce(dims(result, [0, 1, ..., quantization_dimension(result) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) .
  • dim(operand, quantization_dimension(operand)) = dim(result, quantization_dimension(result)) .
  • reduce(dims(operand, [quantization_dimension(operand) + 1, ..., rank(operand) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) = reduce(dims(result, [quantization_dimension(result) + 1, ..., rank(result) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) .

مثال ها

// %operand: [[1, 2, 3], [4, 5, 6]]
%result = "stablehlo.reshape"(%operand) : (tensor<2x3xi32>) -> tensor<3x2xi32>
// %result: [[1, 2], [3, 4], [5, 6]]

نمونه های بیشتر

معکوس

مفاهیم

Reverses the order of elements in the operand along the specified dimensions and produces a result tensor. More formally, result[result_index] = operand[operand_index] where:

• operand_index[d] = dim(result, d) - result_index[d] - 1 if d in dimensions .
• operand_index[d] = result_index[d] otherwise.

ورودی ها

خروجی ها

محدودیت ها

• (C1) type(operand) = type(result) .
• (C2) is_unique(dimensions) .
• (C3) 0 <= dimensions < rank(result) .

مثال ها

// %operand = [[1, 2], [3, 4], [5, 6]]
%result = "stablehlo.reverse"(%operand) {
  dimensions = array<i64: 1>
} : (tensor<3x2xi32>) -> tensor<3x2xi32>
// %result: [[2, 1], [4, 3], [6, 5]]

نمونه های بیشتر

rng

مفاهیم

Generates random numbers using the rng_distribution algorithm and produces a result tensor of a given shape shape .

If rng_distribution = UNIFORM , then the random numbers are generated following the uniform distribution over the interval [a, b) . If a >= b , the behavior is undefined.

If rng_distribution = NORMAL , then the random numbers are generated following the normal distribution with mean = a and standard deviation = b . If b < 0 , the behavior is undefined.

The exact way how random numbers are generated is implementation-defined. For example, they may or may not be deterministic, and they may or may not use hidden state.

In conversations with many stakeholders, this op has come up as effectively deprecated, so in the future we are planning to explore removing it ( #597 ).

ورودی ها

خروجی ها

محدودیت ها

• (C1) element_type(a) = element_type(b) = element_type(result) .
• (C2) If rng_distribution = NORMAL , then is_float(a) .
• (C3) shape(result) = shape .

مثال ها

// %a = 0
// %b = 2
// %shape = [3, 3]
%result = "stablehlo.rng"(%a, %b, %shape) {
  rng_distribution = #stablehlo<rng_distribution UNIFORM>
} : (tensor<i32>, tensor<i32>, tensor<2xi64>) -> tensor<3x3xi32>
// %result: [
//           [1, 0, 1],
//           [1, 1, 1],
//           [0, 0, 0]
//          ]

rng_bit_generator

مفاهیم

Returns an output filled with uniform random bits and an updated output state output_state using the pseudorandom number generator algorithm rng_algorithm given an initial state initial_state . The output is guaranteed to be deterministic function of initial_state , but it is not guaranteed to be deterministic between implementations.

rng_algorithm is one of the following:

• DEFAULT : Implementation-defined algorithm.
• THREE_FRY : Implementation-defined variant of the Threefry algorithm.*
• PHILOX : Implementation-defined variant of the Philox algorithm.*

* See: Salmon et al. SC 2011. Parallel random numbers: as easy as 1, 2, 3.

ورودی ها

خروجی ها

محدودیت ها

• (C1) type(initial_state) = type(output_state) .
• (C2) size(initial_state) is defined as:
  • implementation-defined if rng_algorithm = DEFAULT .
  • 2 if rng_algorithm = THREE_FRY .
  • 2 or 3 if rng_algorithm = PHILOX .

مثال ها

// %initial_state: [1, 2]
%output_state, %output = "stablehlo.rng_bit_generator"(%initial_state) {
  rng_algorithm = #stablehlo<rng_algorithm THREE_FRY>
} : (tensor<2xui64>) -> (tensor<2xui64>, tensor<2x2xui64>)
// %output_state: [1, 6]
// %output: [
//           [9236835810183407956, 16087790271692313299],
//           [18212823393184779219, 2658481902456610144]
//          ]

round_nearest_afz

مفاهیم

Performs element-wise rounding towards the nearest integer, breaking ties away from zero, on the operand tensor and produces a result tensor. Implements the roundToIntegralTiesToAway operation from the IEEE-754 specification. For quantized types, performs dequantize_op_quantize(round_nearest_afz, operand, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

// %operand = [-2.5, 0.4, 0.5, 0.6, 2.5]
%result = "stablehlo.round_nearest_afz"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// %result: [-3.0, 0.0, 1.0, 1.0, 3.0]

نمونه های بیشتر

round_nearest_even

مفاهیم

Performs element-wise rounding towards the nearest integer, breaking ties towards the even integer, on the operand tensor and produces a result tensor. Implements the roundToIntegralTiesToEven operation from the IEEE-754 specification. For quantized types, performs dequantize_op_quantize(round_nearest_even, operand, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

// %operand = [-2.5, 0.4, 0.5, 0.6, 2.5]
%result = "stablehlo.round_nearest_even"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// %result: [-2.0, 0.0, 0.0, 1.0, 2.0]

نمونه های بیشتر

rsqrt

مفاهیم

Performs element-wise reciprocal square root operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: rSqrt from IEEE-754.
• For complex numbers: complex reciprocal square root.
• For quantized types: dequantize_op_quantize(rsqrt, operand, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

// %operand: [[1.0, 4.0], [9.0, 25.0]]
%result = "stablehlo.rsqrt"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[1.0, 0.5], [0.33333343, 0.2]]

نمونه های بیشتر

پراکنده کردن

مفاهیم

Produces results tensors which are equal to inputs tensors except that several slices specified by scatter_indices are updated with the values updates using update_computation .

The following diagram shows how elements in updates... map on elements in results... using a concrete example. The diagram picks a few example updates... indices and explains in detail which results... indices they correspond to.

More formally, for all update_index in index_space(updates[0]) :

• update_scatter_dims = [d for d in axes(updates[0]) and d not in update_window_dims] .
• update_scatter_index = update_index[update_scatter_dims...] .
• start_index is defined as:
  • scatter_indices[si0, ..., :, ..., siN] where si are individual elements in update_scatter_index and : is inserted at the index_vector_dim index, if index_vector_dim < rank(scatter_indices) .
  • [scatter_indices[update_scatter_index]] otherwise.
• For d_input in axes(inputs[0]) ,
  • full_start_index[d_input] = start_index[d_start] if d_input = scatter_dims_to_operand_dims[d_start] .
  • full_start_index[d_input] = 0 otherwise.
• For d_input in axes(inputs[0]) ,
  • full_batching_index[d_input] = update_scatter_index[d_start - (d_start < index_vector_dim ? 0 : 1)] if d_input = input_batching_dims[i_batching] and d_start = scatter_indices_batching_dims[i_batching] .
  • full_batching_index[d_input] = 0 otherwise.
• update_window_index = update_index[update_window_dims...] .
• full_window_index = [wi0, ..., 0, ..., wiN] where wi are individual elements in update_window_index , and 0 is inserted at indices from inserted_window_dims and input_batching_dims .
• result_index = full_start_index + full_batching_index + full_window_index .

Given that, results = exec(schedule, inputs) , where:

• schedule is an implementation-defined permutation of index_space(updates[0]) .
• exec([update_index, ...], results) = exec([...], updated_results) where:
  • If result_index is in bounds for shape(results...)
  • updates_converted = to_destination_type( updates...[update_index], type(func_inputs(update_computation) [len(func_inputs(update_computation))//2:])... )
  • updated_values = update_computation(results...[result_index], updates_converted)
  • updated_results is a copy of results with results...[result_index] set to updated_values... .
  • در غیر این صورت
  • updated_results = results .
• exec([], results) = results .

If indices_are_sorted is true then the implementation can assume that scatter_indices are sorted with respect to scatter_dims_to_operand_dims , otherwise the behavior is undefined. More formally, for all i1 < i2 from indices(result) , full_start_index(i1) <= full_start_index(i2) .

If unique_indices is true then the implementation can assume that all result_index indices being scattered to are unique. If unique_indices is true but the indices being scattered to are not unique then the behavior is undefined.

ورودی ها

خروجی ها

محدودیت ها

• (C1) same(shape(inputs...)) .
• (C2) `rank(inputs[0]) = size(update_window_dims) + size(inserted_window_dims)
  • size(input_batching_dims)`.
• (C3) same(shape(updates...)) .
• (C4) shape(updates[0]) = combine(update_scatter_dim_sizes, update_window_dim_sizes) where:
  • update_scatter_dim_sizes = shape(scatter_indices) except that the dimension size of scatter_indices corresponding to index_vector_dim is not included.
  • update_window_dim_sizes <= shape(inputs[0]) except that the dimension sizes in inputs[0] corresponding to inserted_window_dims and input_batching_dims are not included.
  • combine puts update_scatter_dim_sizes at axes corresponding to update_scatter_dims and update_window_dim_sizes at axes corresponding to update_window_dims .
• (C5) 0 < size(inputs) = size(updates) = N .
• (C6) element_type(updates...) = element_type(inputs...) .
• (C7) is_unique(update_window_dims) and is_sorted(update_window_dims) .
• (C8) 0 <= update_window_dims < rank(updates[0]) .
• (C9) is_unique(concatenate(inserted_window_dims, input_batching_dims))
• (C10) is_sorted(inserted_window_dims) .
• (C11) 0 <= inserted_window_dims < rank(inputs[0]) .
• (C12) is_sorted(input_batching_dims) .
• (C13) 0 <= input_batching_dims < rank(inputs[0])) .
• (C14) is_unique(scatter_indices_batching_dims) .
• (C15) 0 <= scatter_indices_batching_dims < rank(scatter_indices) .
• (C16) index_vector_dim not in scatter_indices_batching_dims .
• (C17) size(input_batching_dims) == size(scatter_indices_batching_dims) .
• (C18) dim(inputs[0], input_batching_dims...) = dim(scatter_indices, scatter_indices_batching_dims...) .
• (C19) size(scatter_dims_to_operand_dims) = index_vector_dim < rank(scatter_indices) ? dim(scatter_indices, index_vector_dim) : 1 .
• (C20) is_unique(concatenate(scatter_dims_to_operand_dims, input_batching_dims)) .
• (C21) 0 <= scatter_dims_to_operand_dims < rank(inputs[0]) .
• (C22) 0 <= index_vector_dim <= rank(scatter_indices) .
• (C23) update_computation has type (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>) , where is_promotable(element_type(inputs[i]), Ei) .
• (C24) shape(inputs...) = shape(results...) .
• (C25) element_type(results[i]) = Ei for all i in [0,N) .

مثال ها

// %input: [
//          [
//           [[1, 2], [3, 4], [5, 6], [7, 8]],
//           [[9, 10],[11, 12], [13, 14], [15, 16]],
//           [[17, 18], [19, 20], [21, 22], [23, 24]]
//          ],
//          [
//           [[25, 26], [27, 28], [29, 30], [31, 32]],
//           [[33, 34], [35, 36], [37, 38], [39, 40]],
//           [[41, 42], [43, 44], [45, 46], [47, 48]]
//          ]
//         ]
// %scatter_indices: [
//                    [
//                     [[0, 0], [1, 0], [2, 1]],
//                     [[0, 1], [1, 1], [0, 9]]
//                    ],
//                    [
//                     [[0, 0], [2, 1], [2, 2]],
//                     [[1, 2], [0, 1], [1, 0]]
//                    ]
//                   ]
// %update: [
//           [
//            [[1, 1], [1, 1], [1, 1]],
//            [[1, 1], [1, 1], [1, 1]]
//           ],
//           [
//            [[1, 1], [1, 1], [1, 1]],
//            [[1, 1], [1, 1], [1, 1]]
//           ]
//          ]
%result = "stablehlo.scatter"(%input, %scatter_indices, %update) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  scatter_dimension_numbers = #stablehlo.scatter<
    update_window_dims = [3, 4],
    inserted_window_dims = [1],
    input_batching_dims = [0],
    scatter_indices_batching_dims = [1],
    scatter_dims_to_operand_dims = [2, 1],
    index_vector_dim = 3>,
  indices_are_sorted = false,
  unique_indices = false
} : (tensor<2x3x4x2xi64>, tensor<2x2x3x2xi64>, tensor<2x2x3x2x2xi64>) -> tensor<2x3x4x2xi64>
// %result: [
//           [
//            [[3, 4], [6, 7], [6, 7], [7, 8]],
//            [[9, 10],[11, 12], [15, 16], [17, 18]],
//            [[17, 18], [19, 20], [22, 23], [24, 25]]
//           ],
//           [
//            [[25, 26], [28, 29], [30, 31], [31, 32]],
//            [[35, 36], [38, 39], [38, 39], [39, 40]],
//            [[41, 42], [44, 45], [46, 47], [47, 48]]
//           ]
//          ]

نمونه های بیشتر

انتخاب کنید

مفاهیم

Produces a result tensor where each element is selected from on_true or on_false tensor based on the value of the corresponding element of pred . More formally, result[result_index] = pred_element ? on_true[result_index] : on_false[result_index] , where pred_element = rank(pred) = 0 ? pred[] : pred[result_index] . For quantized types, performs dequantize_select_quantize(pred, on_true, on_false, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

• (C1) rank(pred) = 0 or shape(pred) = shape(on_true) .
• (C2) baseline_type(on_true) = baseline_type(on_false) = baseline_type(result) .

مثال ها

// %pred: [[false, true], [true, false]]
// %on_true: [[1, 2], [3, 4]]
// %on_false: [[5, 6], [7, 8]]
%result = "stablehlo.select"(%pred, %on_true, %on_false) : (tensor<2x2xi1>, tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[5, 2], [3, 8]]

نمونه های بیشتر

select_and_scatter

مفاهیم

Scatters the values from the source tensor using scatter based on the outcome of reduce_window of the input tensor using select and produces a result tensor.

The following diagram shows how elements in result are computed from operand and source using a concrete example.

رسمی تر:

• selected_values = reduce_window_without_init(...) with the following inputs:

  • inputs = [operand].
  • window_dimensions , window_strides , and padding which are used as is.
  • base_dilations = windows_dilations = 1 .
  • body is defined as:

  def body(arg0: tensor<E>, arg1: tensor<E>) -> tensor<E>:
    return select(arg0, arg1) ? arg0 : arg1;
  

  where E = element_type(operand) , and reduce_window_without_init works exactly like reduce_window , except that the schedule of the underlying reduce (see reduce ) doesn't include init values. It is currently unspecified what happens if the corresponding window doesn't have values ( #731 ).

• result[result_index] = reduce([source_values], [init_value], [0], scatter) where:

  • source_values = [source[source_index] for source_index in source_indices] .
  • selected_index(source_index) = operand_index if selected_values[source_index] has the operand element from operand_index .
  • source_indices = [source_index for source_index in indices(source) if selected_index(source_index) = result_index] .

ورودی ها

خروجی ها

محدودیت ها

• (C1) element_type(operand) = element_type(source) .
• (C2) shape(source) = num_windows where:
  • padded_operand_shape = padding[:, 0] + shape(operand) + padding[:, 1] .
  • is_empty_window = padded_operand_shape = 0 || window_dimensions > padded_operand_shape .
  • num_windows = is_empty_window ? 0 : floor((padded_operand_shape - window_dimensions) / window_strides) + 1 .
• (C3) element_type(init_value) = element_type(operand) .
• (C4) size(window_dimensions) = rank(operand) .
• (C5) 0 < window_dimensions .
• (C6) size(window_strides) = rank(operand) .
• (C7) 0 < window_strides .
• (C8) shape(padding) = [rank(operand), 2] .
• (C9) select has type (tensor<E>, tensor<E>) -> tensor<i1> where E = element_type(operand) .
• (C10) scatter has type (tensor<E>, tensor<E>) -> tensor<E> where is_promotable(element_type(operand), E) .
• (C11) shape(operand) = shape(result) .
• (C12) element_type(result) = E .

مثال ها

// %operand: [[1, 5], [2, 5], [3, 6], [4, 4]]
// %source: [[5, 6], [7, 8]]
// %init_value: 0
%result = "stablehlo.select_and_scatter"(%operand, %source, %init_value) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.compare"(%arg0, %arg1) {
      comparison_direction = #stablehlo<comparison_direction GE>
    } : (tensor<i64>, tensor<i64>) -> tensor<i1>
    "stablehlo.return"(%0) : (tensor<i1>) -> ()
}, {
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  window_dimensions = array<i64: 3, 1>,
  window_strides = array<i64: 2, 1>,
  padding = dense<[[0, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<4x2xi64>, tensor<2x2xi64>, tensor<i64>) -> tensor<4x2xi64>
// %result: [[0, 0], [0, 0], [5, 14], [7, 0]]

نمونه های بیشتر

ارسال

مفاهیم

Sends inputs to a channel channel_id and produces a result token.

If is_host_transfer is true , then the operation transfers data to the host. Otherwise, it transfers data to another device. What this means is implementation-defined. This flag duplicates the information provided in channel_type , so in the future we are planning to only keep one of them ( #666 ).

ورودی ها

خروجی ها

محدودیت ها

• (C1) channel_type is defined as:
  • DEVICE_TO_HOST if is_host_transfer = true ,
  • DEVICE_TO_DEVICE otherwise.

مثال ها

%result = "stablehlo.send"(%operand, %token) {
  channel_handle = #stablehlo.channel_handle<handle = 1, type = 2>,
  is_host_transfer = true
} : (tensor<2x2xi64>, !stablehlo.token) -> !stablehlo.token

نمونه های بیشتر

shift_left

مفاهیم

Performs element-wise left-shift operation on the lhs tensor by rhs number of bits and produces a result tensor.

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

// %lhs: [-1, 0, 1]
// %rhs: [1, 2, 3]
%result = "stablehlo.shift_left"(%lhs, %rhs): (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
// %result: [-2, 0, 8]

نمونه های بیشتر

shift_right_arithmetic

مفاهیم

Performs element-wise arithmetic right-shift operation on the lhs tensor by rhs number of bits and produces a result tensor.

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

// %lhs: [-1, 0, 8]
// %rhs: [1, 2, 3]
%result = "stablehlo.shift_right_arithmetic"(%lhs, %rhs): (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
// %result: [-1, 0, 1]

نمونه های بیشتر

shift_right_logical

مفاهیم

Performs element-wise logical right-shift operation on the lhs tensor by rhs number of bits and produces a result tensor.

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

// %lhs: [-1, 0, 8]
// %rhs: [1, 2, 3]
%result = "stablehlo.shift_right_logical"(%lhs, %rhs): (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
// %result: [9223372036854775807, 0, 1]

نمونه های بیشتر

امضا کردن

مفاهیم

Returns the sign of the operand element-wise and produces a result tensor. More formally, for each element x , the semantics can be expressed using Python syntax as follows:

def sign(x):
  if is_integer(x):
    if compare(x, 0, LT, SIGNED): return -1
    if compare(x, 0, EQ, SIGNED): return 0
    return 1
  elif is_float(x):
    if is_nan(x): return NaN
    if compare(x, -0.0, EQ, FLOAT): return -0.0
    if compare(x, +0.0, EQ, FLOAT): return +0.0
    if compare(x, 0.0, LT, FLOAT): return -1.0
    return 1.0
  elif is_complex(x):
    if is_nan(real(x)) or is_nan(imag(x)): return (NaN, NaN)
    if compare(x, (0.0, 0.0), EQ, FLOAT): return (0.0, 0.0)
    return divide(x, convert(abs(x), type(x)))

For quantized types, performs dequantize_op_quantize(sign, operand, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

// Logical values: +NaN, -1.0, -0.0, +0.0, 1.0
// operand: [0x7FFFFFFFFFFFFFFF, -1.0, -0.0, 0.0, 1.0]
%result = "stablehlo.sign"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// Logical values: +NaN, -1.0, -0.0, +0.0, 1.0
// %result: [0x7FFFFFFFFFFFFFFF, -1.0, -0.0, 0.0, 1.0]

نمونه های بیشتر

سینوسی

مفاهیم

Performs element-wise sine operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: sin from IEEE-754.
• For complex numbers: complex sine.
• For quantized types: dequantize_op_quantize(sine, operand, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

تکه

مفاهیم

Extracts a slice from the operand using statically-computed starting indices and produces a result tensor. start_indices contain the starting indices of the slice for each dimension, limit_indices contain the ending indices (exclusive) for the slice for each dimension, and strides contain the strides for each dimension.

More formally, result[result_index] = operand[operand_index] where operand_index = start_indices + result_index * strides .

ورودی ها

خروجی ها

محدودیت ها

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

مثال ها

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

نمونه های بیشتر

مرتب سازی

مفاهیم

Sorts 1-dimensional slices of inputs along the dimension dimension together, according to a comparator and produces results .

Unlike similar inputs in other operations, dimension allows negative values, with the semantics described below. In the future, this may be disallowed for consistency reasons ( #1377 ).

If is_stable is true, then the sorting is stable, that is, relative order of elements considered to be equal by the comparator is preserved. For the case where there is a single input, two elements e1 and e2 are considered to be equal by the comparator if and only if comparator(e1, e2) = comparator(e2, e1) = false . See the formalization below for how this generalizes to multiple inputs.

More formally, for all result_index in index_space(results[0]) :

• adjusted_dimension = dimension >= 0 ? dimension : rank(inputs[0]) + dimension .
• result_slice = [ri0, ..., :, ..., riR-1] where riN are individual elements in result_index , and : is inserted at adjusted_dimension .
• inputs_together = (inputs[0]..., ..., inputs[N-1]...) .
• results_together[result_slice] = sort(inputs_together[result_slice], comparator_together) .
• where sort sorts a 1-dimensional slice in non-descending order expecting that comparator_together returns true if the left-hand side argument is less than the right-hand second argument.

• def comparator_together(lhs_together, rhs_together):
    args = []
    for (lhs_el, rhs_el) in zip(lhs_together, rhs_together):
      args.append(lhs_el)
      args.append(rhs_el)
    return comparator(*args)
  

• (results[0]..., ..., results[N-1]...) = results_together .

ورودی ها