مشخصات 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_type نوع عدد صحیح (C1-C3)، (C8)
storage_min ثابت عدد صحیح (C1)، (C3)، (C7)
storage_max ثابت عدد صحیح (C2)، (C3)، (C7)
expressed_type نوع ممیز شناور (C4)
quantization_dimension ثابت عدد صحیح اختیاری (C10-C12)
scales تعداد متغیر ثابت های ممیز شناور (C4-C6)، (C9)، (C10)، (C13)
zero_points تعداد متغیر ثابت های اعداد صحیح (C7-C9)

انواع عناصر کوانتیزه مقادیر صحیح یک نوع ذخیره‌سازی را در محدوده از storage_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 را نشان می دهند.
  • انواع ممیز شناور می تواند یکی از موارد زیر باشد:
  • انواع مختلط مقادیر پیچیده ای را نشان می دهند که یک قسمت واقعی و یک قسمت خیالی از همان نوع عنصر دارند. انواع پیچیده پشتیبانی شده 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)) .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand تانسور عدد صحیح علامت دار، ممیز شناور، یا نوع مختلط یا تانسور کوانتیزه هر تانسور (C1-C2)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور از نوع اعداد صحیح علامت دار یا ممیز شناور یا تانسور کوانتیزه هر تانسور (C1-C2)

محدودیت ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) lhs تانسور یا تانسور کوانتیزه (C1-C6)
(I2) rhs تانسور یا تانسور کوانتیزه (C1-C5)، (C7)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور یا تانسور کوانتیزه (C1-C7)

محدودیت ها

  • اگر عملیات از تانسورهای غیر کوانتیزه استفاده می کند:
    • (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 وجود دارد.

ورودی ها

برچسب نام تایپ کنید
(I1) inputs تعداد متغیر token

خروجی ها

نام تایپ کنید
result token

مثال ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand تانسور یا تانسور کوانتیزه در هر تانسور (C1)، (C6)
(I2) all_gather_dim ثابت از نوع si64 (C1)، (C6)
(I3) replica_groups ثابت تانسور دو بعدی از نوع si64 (C2-C4)
(I4) channel_id ثابت از نوع si64 (C5)
(I5) use_global_device_ids ثابت از نوع i1 (C5)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور یا تانسور کوانتیزه در هر تانسور (C6)

محدودیت ها

  • (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])) است.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand تانسور یا تانسور کوانتیزه در هر تانسور (C5)، (C6)
(I2) replica_groups تعداد متغیر ثابت های تانسور یک بعدی از نوع si64 (C1-C3)
(I3) channel_id ثابت از نوع si64 (C4)
(I4) use_global_device_ids ثابت از نوع i1 (C4)
(I5) computation تابع (C5)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور یا تانسور کوانتیزه در هر تانسور (C6-C7)

محدودیت ها

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

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

all_to_all

مفاهیم

all_to_all

در هر گروه فرآیندی در شبکه فرآیند 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) .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand تانسور یا تانسور کوانتیزه در هر تانسور (C1-C3)، (C9)
(I2) split_dimension ثابت از نوع si64 (C1)، (C2)، (C9)
(I3) concat_dimension ثابت از نوع si64 (C3)، (C9)
(I4) split_count ثابت از نوع si64 (C2)، (C4)، (C8)، (C9)
(I5) replica_groups ثابت تانسور دو بعدی از نوع si64 (C5-C8)
(I6) channel_id ثابت از نوع si64

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور یا تانسور کوانتیزه در هر تانسور (C9)

محدودیت ها

  • (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) به جز، اگر split_dimension != concat_dimension :
    • 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.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) lhs تانسور از نوع بولی یا عدد صحیح (C1)
(I2) rhs تانسور از نوع بولی یا عدد صحیح (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور از نوع بولی یا عدد صحیح (C1)

محدودیت ها

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

مثال ها

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

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

آتان 2

مفاهیم

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

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) lhs تانسور از نوع ممیز شناور یا مختلط یا تانسور کوانتیزه هر تانسور (C1)
(I2) rhs تانسور از نوع ممیز شناور یا مختلط یا تانسور کوانتیزه هر تانسور (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور از نوع ممیز شناور یا مختلط یا تانسور کوانتیزه هر تانسور (C1)

محدودیت ها

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

مثال ها

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

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

batch_norm_grad

مفاهیم

گرادیان های چندین ورودی batch_norm_training را که از grad_output پس انتشار می یابند محاسبه می کند و تانسورهای grad_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)) .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand تانسور از نوع ممیز شناور یا تانسور کوانتیزه هر تانسور (C1-C3)، (C5)
(I2) scale تانسور 1 بعدی از نوع کوانتیزه ممیز شناور یا هر تانسور (C2)، (C4)، (C5)
(I3) mean تانسور 1 بعدی از نوع کوانتیزه ممیز شناور یا هر تانسور (C2)، (C4)
(I4) variance تانسور 1 بعدی از نوع کوانتیزه ممیز شناور یا هر تانسور (C2)، (C4)
(I5) grad_output تانسور از نوع ممیز شناور یا تانسور کوانتیزه هر تانسور (C2)، (C3)
(I6) epsilon ثابت از نوع f32
(I7) feature_index ثابت از نوع si64 (C1)، (C5)

خروجی ها

نام تایپ کنید محدودیت ها
grad_operand تانسور از نوع ممیز شناور یا تانسور کوانتیزه هر تانسور (C2)، (C3)
grad_scale تانسور 1 بعدی از نوع کوانتیزه ممیز شناور یا هر تانسور (C2)، (C4)
grad_offset تانسور 1 بعدی از نوع کوانتیزه ممیز شناور یا هر تانسور (C2)، (C4)

محدودیت ها

  • (C1) 0 <= feature_index < rank(operand) .
  • (C2) 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)) .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand تانسور از نوع ممیز شناور یا تانسور کوانتیزه هر تانسور (C1-C7)
(I2) scale تانسور 1 بعدی از نوع کوانتیزه ممیز شناور یا هر تانسور (C2)، (C3)
(I3) offset تانسور 1 بعدی از نوع کوانتیزه ممیز شناور یا هر تانسور (C2)، (C4)
(I4) mean تانسور 1 بعدی از نوع کوانتیزه ممیز شناور یا هر تانسور (C5)
(I5) variance تانسور 1 بعدی از نوع کوانتیزه ممیز شناور یا هر تانسور (C2)، (C6)
(I6) epsilon ثابت از نوع f32
(I7) feature_index ثابت از نوع si64 (C1)، (C3-C6)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور از نوع ممیز شناور یا تانسور کوانتیزه هر تانسور (C2)، (C7)

محدودیت ها

  • (C1) 0 <= feature_index < rank(operand) .
  • (C2) 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)) انجام می دهد.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand تانسور از نوع ممیز شناور یا تانسور کوانتیزه هر تانسور (C1)
(I2) scale تانسور 1 بعدی ممیز شناور یا هر تانسور کوانتیزه شده است (C2)، (C3)
(I3) offset تانسور 1 بعدی ممیز شناور یا هر تانسور کوانتیزه شده است (C2)، (C4)
(I4) epsilon ثابت از نوع f32 (C1)، (C3-C6)
(I5) feature_index ثابت از نوع si64 (C1)، (C3-C6)

خروجی ها

نام تایپ کنید محدودیت ها
output تانسور از نوع ممیز شناور یا تانسور کوانتیزه هر تانسور (C7)
batch_mean تانسور 1 بعدی ممیز شناور یا هر تانسور کوانتیزه شده است (C2)، (C5)
batch_var تانسور 1 بعدی ممیز شناور یا هر تانسور کوانتیزه شده است (C2)، (C6)

محدودیت ها

  • (C1) 0 <= feature_index < rank(operand) .
  • (C2) 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 نمایش درون حافظه یک مقدار داده شده را برمی گرداند، و رفتار آن به صورت پیاده سازی تعریف شده است، زیرا نمایش دقیق تانسورها توسط پیاده سازی تعریف شده است، و نمایش دقیق انواع عناصر نیز به صورت پیاده سازی تعریف شده است.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand تانسور یا تانسور کوانتیزه (C1-C2)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور یا تانسور کوانتیزه (C1-C2)

محدودیت ها

  • (C1) با توجه به E = is_quantized(operand) ? storage_type(operand) : element_type(operand) , E' = is_quantized(result) ? storage_type(result) : element_type(result) ، و R = rank(operand) :
    • اگر 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]] در غیر این صورت.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand تانسور یا تانسور کوانتیزه (C1-C2)، (C5-C6)
(I2) broadcast_dimensions ثابت تانسور 1 بعدی از نوع si64 (C2-C6)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور یا تانسور کمکی (C1) ، (C3) ، (C5-C6)

محدودیت ها

  • (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] در غیر این صورت.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) index تانسور 0 بعدی از نوع si32
(i2) branches تعداد متغیر توابع (C1-C4)

خروجی ها

نام تایپ کنید محدودیت ها
results تعداد متغیرهای تانسور ، تنش یا توکن های کمکی (C4)

محدودیت ها

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

مثال ها

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

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

CBRT

مفاهیم

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

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) operand تانسور از نقطه شناور یا نوع پیچیده یا تانسور کمیته در هر تانسور (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور از نقطه شناور یا نوع پیچیده یا تانسور کمیته در هر تانسور (C1)

محدودیت ها

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

مثال ها

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

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

سقف

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) operand تانسور از نوع نقطه شناور یا تانسور کمیته در هر تانسور (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور از نوع نقطه شناور یا تانسور کمیته در هر تانسور (C1)

محدودیت ها

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

مثال ها

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

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

بی پروا

مفاهیم

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

به طور رسمی تر ، برای همه 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)) را انجام می دهد.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) a تانسور از نقطه شناور یا نوع پیچیده یا تانسور کمیته در هر تانسور (C1-C3)
(i2) lower ثابت تانسور 0 بعدی از نوع i1

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور از نقطه شناور یا نوع پیچیده یا تانسور کمیته در هر تانسور (C1)

محدودیت ها

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

مثال ها

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

گیره

مفاهیم

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

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) min تانسور یا تانسور هر تانسور (C1) ، (C3)
(i2) operand تانسور یا تانسور هر تانسور (C1-C4)
(i3) max تانسور یا تانسور هر تانسور (C2) ، (C3)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور یا تانسور هر تانسور (C4)

محدودیت ها

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

مثال ها

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

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

مجموعه_

مفاهیم

در هر گروه فرآیند در شبکه فرآیند 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)) در غیر این صورت.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) operand تانسور یا تانسور هر تانسور (C3)
(i2) replica_groups تعداد متغیر ثابت 1 بعدی تنشور نوع si64 (C1) ، (C2)
(i3) channel_id ثابت نوع si64

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور یا تانسور هر تانسور (C3)

محدودیت ها

  • (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)) در غیر این صورت.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) operand تانسور یا تانسور هر تانسور (C5)
(i2) source_target_pairs ثابت 2 بعدی تانسور نوع si64 (C1-C4)
(i3) channel_id ثابت نوع si64

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور یا تانسور هر تانسور (C1)

محدودیت ها

  • (C1) dim(source_target_pairs, 1) = 2 .
  • (C2) is_unique(source_target_pairs[:, 0]) .
  • (C3) is_unique(source_target_pairs[:, 1]) .
  • (C4) 0 <= source_target_pairs < N ، جایی که N به صورت تعریف شده است:
    • 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) .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) lhs تانسور یا تانسور هر تانسور (C1-C3)
(i2) rhs تانسور یا تانسور هر تانسور (C1-C2)
(i3) comparison_direction Enum of EQ ، NE ، GE ، GT ، LE و LT
(i4) compare_type Enum of FLOAT ، TOTALORDER ، SIGNED و UNSIGNED (C3)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور از نوع بولی (C2)

محدودیت ها

  • (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 تولید می کند.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) lhs تانسور نوع f32 یا f64 (C1-C3)
(i2) rhs تانسور نوع f32 یا f64 (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور از نوع پیچیده (C2) ، (C3)

محدودیت ها

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

مثال ها

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

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

کامپوزیت

مفاهیم

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

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

ورودی ها

برچسب نام تایپ کنید
(i1) inputs تعداد متغیرهای مقادیر
(i2) name ثابت string از نوع
(i3) composite_attributes فرهنگ لغت
(i4) decomposition ثابت string از نوع
(i5) version ثابت نوع si32

خروجی ها

نام تایپ کنید
results تعداد متغیرهای مقادیر

محدودیت ها

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

مثال ها

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

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

هماهنگ کردن

مفاهیم

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 .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) inputs تعداد تانسور یا تانسور کمیته در هر تانسور (C1-C6)
(i2) dimension ثابت نوع si64 (C2) ، (C4) ، (C6)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور یا تانسور هر تانسور (C5-C6)

محدودیت ها

  • (C1) same(element_type(inputs...)) .
  • (C2) same(shape(inputs...)) به جز dim(inputs..., dimension) .
  • (C3) 0 < size(inputs) .
  • (C4) 0 <= dimension < rank(inputs[0]) .
  • (c5) element_type(result) = element_type(inputs[0]) .
  • (c6) shape(result) = shape(inputs[0]) به جز:
    • 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 ثابت تولید می کند.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) value ثابت (C1)

خروجی ها

نام تایپ کنید محدودیت ها
output تانسور یا تانسور کمکی (C1)

محدودیت ها

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

مثال ها

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

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

تبدیل

مفاهیم

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

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

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

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

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

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

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) operand تنشور (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result تنشور (C1)

محدودیت ها

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

مثال ها

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

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

پیچیدگی

مفاهیم

محصولات 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) .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) lhs تانسور یا تانسور هر تانسور (C1) ، (C10-C11) ، (C14) (C25) ، (C27-C28) ، (C31-C32) ، (C34)
(i2) rhs تانسور یا تانسور کمکی (C1) ، (C14-C16) ، (C25) ، (C27-C29) ، (C31-C34)
(i3) window_strides ثابت تانسور 1 بعدی از نوع si64 (C2-C3) ، (C25)
(i4) padding ثابت 2 بعدی تانسور نوع si64 (C4) ، (C25)
(i5) lhs_dilation ثابت تانسور 1 بعدی از نوع si64 (C5-C6) ، (C25)
(i6) rhs_dilation ثابت تانسور 1 بعدی از نوع si64 (C7-C8) ، (C25)
(i7) window_reversal ثابت 1 بعدی تانسور نوع i1 (C9)
(i8) input_batch_dimension ثابت نوع si64 (C10) ، (C13) ، (C25)
(i9) input_feature_dimension ثابت نوع si64 (C11) ، (C13-C14)
(i10) input_spatial_dimensions ثابت تانسور 1 بعدی از نوع si64 (C12) ، (C13) ، (C25)
(i11) kernel_input_feature_dimension ثابت نوع si64 (C14) ، (C18)
(I12) kernel_output_feature_dimension ثابت نوع si64 (C15-C16) ، (C18) ، (C25) ، (C29)
(I13) kernel_spatial_dimensions ثابت تانسور 1 بعدی از نوع si64 (C17-C18) ، (C25)
(I14) output_batch_dimension ثابت نوع si64 (C20) ، (C25)
(I15) output_feature_dimension ثابت نوع si64 (C20) ، (C25) ، (C30)
(I16) output_spatial_dimensions ثابت تانسور 1 بعدی از نوع si64 (C19-C20) ، (C25)
(I17) feature_group_count ثابت نوع si64 (C11) ، (C14) ، (C16) ، (C21) ، (C23)
(i18) batch_group_count ثابت نوع si64 (C10) ، (C15) ، (C22) ، (C23) ، (C25)
(I19) precision_config تعداد متغیرهای عناصر DEFAULT ، HIGH و HIGHEST (C24)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور یا تانسور کمکی (C25-C28) ، (C30) ، (C32-34)

محدودیت ها

  • (C1) N = rank(lhs) = rank(rhs) .
  • (C2) size(window_strides) = N - 2 .
  • (C3) 0 < window_strides .
  • (C4) shape(padding) = [N - 2, 2] .
  • (C5) size(lhs_dilation) = N - 2 .
  • (C6) 0 < lhs_dilation .
  • (C7) size(rhs_dilation) = N - 2 .
  • (c8) 0 < rhs_dilation .
  • (C9) size(window_reversal) = N - 2 .
  • (C10) dim(lhs, input_batch_dimension) % batch_group_count = 0 .
  • (C11) dim(lhs, input_feature_dimension) % feature_group_count = 0 .
  • (C12) size(input_spatial_dimensions) = N - 2 .
  • (C13) با توجه به input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension] :
    • 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)) .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) operand تانسور از نقطه شناور یا نوع پیچیده یا تانسور کمیته در هر تانسور (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور از نقطه شناور یا نوع پیچیده یا تانسور کمیته در هر تانسور (C1)

محدودیت ها

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

مثال ها

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

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

count_leading_zeros

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) operand تانسور از نوع عدد صحیح (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور از نوع عدد صحیح (C1)

محدودیت ها

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

مثال ها

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

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

CUSTION_CALL

مفاهیم

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

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

ورودی ها

برچسب نام تایپ کنید
(i1) inputs تعداد متغیرهای مقادیر
(i2) call_target_name ثابت string از نوع
(i3) has_side_effect ثابت نوع i1
(i4) backend_config ثابت string از نوع
(i5) api_version ثابت نوع si32
(i6) called_computations تعداد متغیرهای ثابت رشته string

خروجی ها

نام تایپ کنید
results تعداد متغیرهای مقادیر

مثال ها

%results = "stablehlo.custom_call"(%input0) {
  call_target_name = "foo",
  has_side_effect = false,
  backend_config = "bar",
  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)) .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) lhs تانسور عدد صحیح ، نقطه شناور یا نوع پیچیده یا تانسور کمیته در هر تانسور (C1)
(i2) rhs تانسور عدد صحیح ، نقطه شناور یا نوع پیچیده یا تانسور کمیته در هر تانسور (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور عدد صحیح ، نقطه شناور یا نوع پیچیده یا تانسور کمیته در هر تانسور (C1)

محدودیت ها

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

مثال ها

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

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

dot_general

مفاهیم

محصولات 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 : کمترین محاسبه ، اما دقیق ترین تقریب به تعداد اصلی.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) lhs تانسور یا تانسور هر تانسور (C5-C6) ، (C9-C10) ، (C12-C14) ، (C17-C18) ، (C20)
(i2) rhs تانسور یا تانسور کمکی (C7-C10) ، (C12-C20)
(i3) lhs_batching_dimensions ثابت تانسور 1 بعدی از نوع si64 (c1) ، (c3) ، (c5) ، (c9) ، (c12)
(i4) rhs_batching_dimensions ثابت تانسور 1 بعدی از نوع si64 (c1) ، (c4) ، (c7) ، (c9)
(i5) lhs_contracting_dimensions ثابت تانسور 1 بعدی از نوع si64 (c2) ، (c3) ، (c6) ، (c10)
(i6) rhs_contracting_dimensions ثابت تانسور 1 بعدی از نوع si64 (C2) ، (C4) ، (C8) ، (C10) ، (C16)
(i7) precision_config تعداد متغیرهای عناصر DEFAULT ، HIGH و HIGHEST (C11)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور یا تانسور کمکی (C12) ، (C14) ، (C18-C20)

محدودیت ها

  • (C1) size(lhs_batching_dimensions) = size(rhs_batching_dimensions) .
  • (C2) size(lhs_contracting_dimensions) = size(rhs_contracting_dimensions) .
  • (C3) is_unique(lhs_batching_dimensions + lhs_contracting_dimensions) .
  • (C4) is_unique(rhs_batching_dimensions + rhs_contracting_dimensions) .
  • (c5) 0 <= lhs_batching_dimensions < rank(lhs) .
  • (C6) 0 <= lhs_contracting_dimensions < rank(lhs) .
  • (C7) 0 <= rhs_batching_dimensions < rank(rhs) .
  • (c8) 0 <= rhs_contracting_dimensions < rank(rhs) .
  • (C9) dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...) .
  • (C10) dim(lhs, lhs_contracting_dimensions...) = dim(rhs, rhs_contracting_dimensions...) .
  • (C11) size(precision_config) = 2 .
  • (C12) shape(result) = dim(lhs, lhs_batching_dimensions) + dim(lhs, lhs_result_dimensions) + dim(rhs, rhs_result_dimensions) .
  • اگر این عملیات از تانسور غیرقانونی استفاده می کند:
    • (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 برای بیان دانش استاتیک در مورد رفتار در حال گسترش ابعاد. اگر مشخص نشده باشد ، فرض بر این است که همه ابعاد در حال گسترش هستند.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) operand تانسور یا تانسور کمکی (C1-C2) ، (C5-C6) ، (C9)
(i2) output_dimensions تانسور 1 بعدی از نوع عدد صحیح (C7)
(i3) broadcast_dimensions تانسور ثابت 1 بعدی از نوع عدد صحیح (C2-C6)
(i4) known_expanding_dimensions تانسور ثابت 1 بعدی از نوع عدد صحیح (C8-C9)
(i5) known_non_expanding_dimensions تانسور ثابت 1 بعدی از نوع عدد صحیح (C8-C9)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور یا تانسور کمکی (C1) ، (C3) ، (C5-C7)

محدودیت ها

  • (C1) element_type(result) توسط:
    • 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 مشخص می شود.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(i1) lhs تانسور یا تانسور هر تانسور (C1) ، (C10-C11) ، (C14) (C25) ، (C26-C27) ، (C30-C31) ، (C33)
(i2) rhs تانسور یا تانسور کمکی (C1) ، (C14-C16) ، (C26-C28) ، (C30-C33)
(i3) padding تانسور 2 بعدی از نوع عدد صحیح (C4)
(i4) window_strides ثابت تانسور 1 بعدی از نوع si64 (C2-C3)
(i5) lhs_dilation ثابت تانسور 1 بعدی از نوع si64 (C5-C6)
(i6) rhs_dilation ثابت تانسور 1 بعدی از نوع si64 (C7-C8)
(i7) window_reversal ثابت 1 بعدی تانسور نوع i1 (C9)
(i8) input_batch_dimension ثابت نوع si64 (C10) ، (C13)
(i9) input_feature_dimension ثابت نوع si64 (C11) ، (C13-C14)
(i10) input_spatial_dimensions ثابت تانسور 1 بعدی از نوع si64 (C12) ، (C13)
(i11) kernel_input_feature_dimension ثابت نوع si64 (C14) ، (C18)
(I12) kernel_output_feature_dimension ثابت نوع si64 (C15-C16) ، (C18) ، (C28)
(I13) kernel_spatial_dimensions ثابت تانسور 1 بعدی از نوع si64 (C17-C18)
(I14) output_batch_dimension ثابت نوع si64 (C20)
(I15) output_feature_dimension ثابت نوع si64 (C20) ، (C29)
(I16) output_spatial_dimensions ثابت تانسور 1 بعدی از نوع si64 (C19-C20)
(I17) feature_group_count ثابت نوع si64 (C11) ، (C14) ، (C16) ، (C21) ، (C23)
(i18) batch_group_count ثابت نوع si64 (C10) ، (C15) ، (C22) ، (C23)
(I19) precision_config تعداد متغیرهای عناصر DEFAULT ، HIGH و HIGHEST (C24)

خروجی ها

نام تایپ کنید محدودیت ها
result تانسور یا تانسور کمکی (C25-C27) ، (C29) ، (C31-C33)

محدودیت ها

  • (C1) N = rank(lhs) = rank(rhs) .
  • (C2) size(window_strides) = N - 2 .
  • (C3) 0 < window_strides .
  • (C4) shape(padding) = [N - 2, 2] .
  • (C5) size(lhs_dilation) = N - 2 .
  • (C6) 0 < lhs_dilation .
  • (C7) size(rhs_dilation) = N - 2 .
  • (c8) 0 < rhs_dilation .
  • (C9) size(window_reversal) = N - 2 .
  • (C10) dim(lhs, input_batch_dimension) % batch_group_count = 0 .
  • (C11) dim(lhs, input_feature_dimension) % feature_group_count = 0 .
  • (C12) size(input_spatial_dimensions) = N - 2 .
  • (C13) با توجه به input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension] :
    • is_unique(input_dimensions) .
    • 0 <= input_dimensions < N .
  • (C14) dim(rhs, kernel_input_feature_dimension) = dim(lhs, input_feature_dimension) / feature_group_count .
  • (C15) dim(rhs, kernel_output_feature_dimension) % batch_group_count = 0 .
  • (C16) dim(rhs, kernel_output_feature_dimension) % feature_group_count = 0 .
  • (C17) size(kernel_spatial_dimensions) = N - 2 .
  • (C18) 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.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor or per-tensor quantized tensor (C1), (C7), (C10-C12), (C14)
(I2) start_indices tensor of integer type (C2), (C3), (C13)
(I3) slice_sizes 1-dimensional tensor of integer type (C8), (C11-C13)
(I4) offset_dims 1-dimensional tensor constant of type si64 (C1), (C4-C5), (C13)
(I5) collapsed_slice_dims 1-dimensional tensor constant of type si64 (C1), (C6-C8), (C13)
(I6) start_index_map 1-dimensional tensor constant of type si64 (C3), (C9), (C10)
(I7) index_vector_dim constant of type si64 (C2), (C3), (C13)
(I8) indices_are_sorted constant of type i1

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or per-tensor quantized tensor (C5), (C13-C14)

محدودیت ها

  • (C1) rank(operand) = size(offset_dims) + size(collapsed_slice_dims) .
  • (C2) 0 <= index_vector_dim <= rank(start_indices) .
  • (C3) size(start_index_map) = index_vector_dim < rank(start_indices) ? dim(start_indices, index_vector_dim) : 1 .
  • (C4) is_unique(offset_dims) and is_sorted(offset_dims) .
  • (C5) 0 <= offset_dims < rank(result) .
  • (C6) is_unique(collapsed_slice_dims) and is_sorted(collapsed_slice_dims) .
  • (C7) 0 <= collapsed_slice_dims < rank(operand) .
  • (C8) slice_sizes[collapsed_slice_dims...] <= 1 .
  • (C9) is_unique(start_index_map) .
  • (C10) 0 <= start_index_map < rank(operand) .
  • (C11) size(slice_sizes) = rank(operand) .
  • (C12) 0 <= slice_sizes <= shape(operand) .
  • (C13) shape(result) = combine(batch_dim_sizes, offset_dim_sizes) 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 .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) output_shape 1-dimensional tensor of integer type (C1), (C2)
(I2) iota_dimension si64 (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of integer, floating-point, or complex type or per-tensor quantized tensor (C2)

محدودیت ها

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

مثال ها

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

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

dynamic_pad

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor or per-tensor quantized tensor (C1), (C2), (C4)
(I2) padding_value 0-dimensional tensor or per-tensor quantized tensor (C1)
(I3) edge_padding_low 1-dimensional tensor of integer type (C1), (C4)
(I4) edge_padding_high 1-dimensional tensor of integer type (C1), (C4)
(I5) interior_padding 1-dimensional tensor of integer type (C2-C4)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or per-tensor quantized tensor (C3-C6)

محدودیت ها

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

مثال ها

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

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

dynamic_reshape

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor or quantized tensor (C1-C3)
(I2) output_shape 1-dimensional tensor of integer type (C4)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or quantized tensor (C1-C4)

محدودیت ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor or per-tensor quantized tensor (C1), (C2), (C4)
(I2) start_indices variadic number of 0-dimensional tensors of integer type (C2), (C3)
(I3) slice_sizes 1-dimensional tensor constant of type si64 (C2), (C4), (C5)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or per-tensor quantized tensor (C1), (C5)

محدودیت ها

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

مثال ها

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

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

dynamic_update_slice

مفاهیم

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.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor or per-tensor quantized tensor (C1-C4), (C6)
(I2) update tensor or per-tensor quantized tensor (C2), (C3), (C6)
(I3) start_indices variadic number of 0-dimensional tensors of integer type (C4), (C5)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

نمایی

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of floating-point or complex type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of floating-point or complex type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

exponential_minus_one

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of floating-point or complex type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of floating-point or complex type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

fft

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of floating-point or complex type (C1), (C2), (C4), (C5)
(I2) fft_type enum of FFT , IFFT , RFFT , and IRFFT (C2), (C5)
(I3) fft_length 1-dimensional tensor constant of type si64 (C1), (C3), (C4)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of floating-point or complex type (C2), (C4), (C5)

محدودیت ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of floating-point type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of floating-point type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

جمع آوری

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor or per-tensor quantized tensor (C1), (C8), (C11), (C17), (C19-C21), (C23)
(I2) start_indices tensor of integer type (C2-C3), (C14), (C17), (C22)
(I3) offset_dims 1-dimensional tensor constant of type si64 (C1), (C4-C5), (C22)
(I4) collapsed_slice_dims 1-dimensional tensor constant of type si64 (C1), (C6-C9), (C22)
(I5) operand_batching_dims 1-dimensional tensor constant of type si64 (C1), (C6), (C10-C12), (C16-C18), (C22)
(I6) start_indices_batching_dims 1-dimensional tensor constant of type si64 (C13-C17)
(I7) start_index_map 1-dimensional tensor constant of type si64 (C3), (C18-C19)
(I8) index_vector_dim constant of type si64 (C2-C3), (C15), (C22)
(I9) slice_sizes 1-dimensional tensor constant of type si64 (C9), (C12), (C20-C22)
(I10) indices_are_sorted constant of type i1

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or per-tensor quantized tensor (C5), (C22-C23)

محدودیت ها

  • (C1) rank(operand) = size(offset_dims) + size(collapsed_slice_dims) + size(operand_batching_dims) .
  • (C2) 0 <= index_vector_dim <= rank(start_indices) .
  • (C3) size(start_index_map) = index_vector_dim < rank(start_indices) ? dim(start_indices, index_vector_dim) : 1 .
  • (C4) is_unique(offset_dims) and is_sorted(offset_dims) .
  • (C5) 0 <= offset_dims < rank(result) .
  • (C6) is_unique(concatenate(collapsed_slice_dims, operand_batching_dims))
  • (C7) is_sorted(collapsed_slice_dims) .
  • (C8) 0 <= collapsed_slice_dims < rank(operand) .
  • (C9) slice_sizes[collapsed_slice_dims...] <= 1 .
  • (C10) is_sorted(operand_batching_dims) .
  • (C11) 0 <= operand_batching_dims < rank(operand) .
  • (C12) slice_sizes[operand_batching_dims...] <= 1 .
  • (C13) is_unique(start_indices_batching_dims) .
  • (C14) 0 <= start_indices_batching_dims < rank(start_indices) .
  • (C15) index_vector_dim not in start_indices_batching_dims .
  • (C16) size(operand_batching_dims) == size(start_indices_batching_dims) .
  • (C17) dim(operand, operand_batching_dims...) = dim(start_indices, start_indices_batching_dims...) .
  • (C18) is_unique(concatenate(start_index_map, operand_batching_dims)) .
  • (C19) 0 <= start_index_map < rank(operand) .
  • (C20) size(slice_sizes) = rank(operand) .
  • (C21) 0 <= slice_sizes <= shape(operand) .
  • (C22) shape(result) = combine(batch_dim_sizes, offset_dim_sizes) 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.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor or quantized tensor (C1)
(I2) dimension constant of type si64 (C1)

خروجی ها

نام تایپ کنید
result 0-dimensional tensor of type si32

محدودیت ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand چندتایی (C1), (C2)
(I2) index constant of type si32 (C1), (C2)

خروجی ها

نام تایپ کنید محدودیت ها
result any supported type (C2)

محدودیت ها

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

مثال ها

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

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

اگر

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) pred 0-dimensional tensor of type i1
(I2) true_branch تابع (C1-C3)
(I3) false_branch تابع (C1), (C2)

خروجی ها

نام تایپ کنید محدودیت ها
results variadic number of tensors, quantized tensors or tokens (C3)

محدودیت ها

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

مثال ها

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

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

imag

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of floating-point or complex type (C1), (C2)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of floating-point type (C1), (C2)

محدودیت ها

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

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

infeed

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید
(I1) token token
(I2) infeed_config constant of type string

خروجی ها

نام تایپ کنید محدودیت ها
results variadic number of tensors, quantized tensors or tokens (C1-C3)

محدودیت ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) iota_dimension si64 (C1)

خروجی ها

نام تایپ کنید محدودیت ها
output tensor of integer, floating-point, or complex type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

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

is_finite

مفاهیم

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 .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) x tensor of floating-point type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
y tensor of boolean type (C1)

محدودیت ها

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

مثال ها

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

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

ورود به سیستم

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of floating-point or complex type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of floating-point or complex type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

log_plus_one

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of floating-point or complex type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of floating-point or complex type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

logistic

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of floating-point or complex type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of floating-point or complex type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

نقشه

مفاهیم

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

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) inputs variadic number of tensors or per-tensor quantized tensors (C1-C4)
(I2) dimensions 1-dimensional tensor constant of type si64 (C3)
(I3) computation تابع (C4)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or per-tensor quantized tensor (C1), (C4)

محدودیت ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) lhs tensor or per-tensor quantized tensor (C1)
(I2) rhs tensor or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

کمترین

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) lhs tensor or per-tensor quantized tensor (C1)
(I2) rhs tensor or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

تکثیر کردن

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) lhs tensor or per-tensor quantized tensor (C1)
(I2) rhs tensor or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

negate

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of integer, floating-point, or complex type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of integer, floating-point, or complex type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

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

نه

مفاهیم

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.

استدلال ها

نام تایپ کنید محدودیت ها
operand tensor of boolean or integer type (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of boolean or integer type (C1)

محدودیت ها

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

مثال ها

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

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

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

optimization_barrier

مفاهیم

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 .

استدلال ها

نام تایپ کنید محدودیت ها
operand variadic number of tensors, per-tensor quantized tensors or tokens (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result variadic number of tensors, per-tensor quantized tensors or tokens (C1)

محدودیت ها

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

مثال ها

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

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

یا

مفاهیم

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.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) lhs tensor of integer or boolean type (C1)
(I2) rhs tensor of integer or boolean type (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of integer or boolean type (C1)

محدودیت ها

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

مثال ها

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

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

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

outfeed

مفاهیم

Writes inputs to the outfeed and produces a result token.

Semantics of outfeed_config is implementation-defined.

ورودی ها

برچسب نام تایپ کنید
(I1) inputs variadic number of tensors or quantized tensors
(I2) token token
(I3) outfeed_config constant of type string

خروجی ها

نام تایپ کنید
result token

مثال ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor or per-tensor quantized tensor (C1), (C2), (C4)
(I2) padding_value 0-dimensional tensor or per-tensor quantized tensor (C1)
(I3) edge_padding_low 1-dimensional tensor constant of type si64 (C1), (C4)
(I4) edge_padding_high 1-dimensional tensor constant of type si64 (C1), (C4)
(I5) interior_padding 1-dimensional tensor constant of type si64 (C2-C4)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or per-tensor quantized tensor (C3-C6)

محدودیت ها

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

مثال ها

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

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

partition_id

مفاهیم

Produces partition_id of the current process.

خروجی ها

نام تایپ کنید
result 0-dimensional tensor of type ui32

مثال ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of integer type (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of integer type (C1)

محدودیت ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) lhs tensor of integer, floating-point, or complex type or per-tensor quantized tensor (C1)
(I2) rhs tensor of integer, floating-point, or complex type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of integer, floating-point, or complex type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

واقعی

مفاهیم

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 .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of floating-point or complex type (C1), (C2)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of floating-point type (C1), (C2)

محدودیت ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) token token (C4)
(I2) channel_id constant of type si64
(I3) channel_type enum of DEVICE_TO_DEVICE and HOST_TO_DEVICE (C1)
(I4) is_host_transfer constant of type i1 (C1)

خروجی ها

نام تایپ کنید محدودیت ها
results variadic number of tensors, quantized tensors or tokens (C2-C4)

محدودیت ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) inputs variadic number of tensors or per-tensor quantized tensors (C1-C4), (C6), (C7)
(I2) init_values variadic number of 0-dimensional tensors or per-tensor quantized tensors (C2), (C3)
(I3) dimensions 1-dimensional tensor constant of type si64 (C4), (C5), (C7)
(I4) body تابع (C6)

خروجی ها

نام تایپ کنید محدودیت ها
results variadic number of tensors or per-tensor quantized tensors (C3), (C7), (C8)

محدودیت ها

  • (C1) same(shape(inputs...)) .
  • (C2) element_type(inputs...) = element_type(init_values...) .
  • (C3) 0 < size(inputs) = size(init_values) = size(results) = N .
  • (C4) 0 <= dimensions < rank(inputs[0]) .
  • (C5) is_unique(dimensions) .
  • (C6) body 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.

More formally:

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of floating-point type or per-tensor quantized tensor (C1)
(I2) exponent_bits constant of type si32 (C2)
(I3) mantissa_bits constant of type si32 (C3)

خروجی ها

نام تایپ کنید محدودیت ها
output tensor of floating-point type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

reduce_scatter

مفاهیم

reduce_scatter

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor or per-tensor quantized tensor (C1), (C2), (C7), (C8)
(I2) scatter_dimension constant of type si64 (C1), (C2), (C8)
(I3) replica_groups 2-dimensional tensor constant of type si64 (C3-C5)
(I4) channel_id constant of type si64 (C6)
(I5) use_global_device_ids constant of type i1 (C6)
(I6) computation تابع (C7)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or per-tensor quantized tensor (C8-C9)

محدودیت ها

  • (C1) dim(operand, scatter_dimension) % dim(process_groups, 1) = 0 .
  • (C2) 0 <= scatter_dimension < rank(operand) .
  • (C3) is_unique(replica_groups) .
  • (C4) size(replica_groups) 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.

reduce_window

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) inputs variadic number of tensors or per-tensor quantized tensors (C1-C4), (C6), (C8), (C10), (C12), (C13), (C15)
(I2) init_values variadic number of 0-dimensional tensors or per-tensor quantized tensors (C1), (C13)
(I3) window_dimensions 1-dimensional tensor constant of type si64 (C4), (C5), (C15)
(I4) window_strides 1-dimensional tensor constant of type si64 (C6), (C7), (C15)
(I5) base_dilations 1-dimensional tensor constant of type si64 (C8), (C9), (C15)
(I6) window_dilations 1-dimensional tensor constant of type si64 (C10), (C11), (C15)
(I7) padding 2-dimensional tensor constant of type si64 (C12), (C15)
(I8) body تابع (C13)

خروجی ها

نام تایپ کنید محدودیت ها
results variadic number of tensors or per-tensor quantized tensors (C1), (C14-C16)

محدودیت ها

  • (C1) 0 < size(inputs) = size(init_values) = size(results) = N .
  • (C2) same(shape(inputs...)) .
  • (C3) element_type(inputs...) = element_type(init_values...) .
  • (C4) size(window_dimensions) = rank(inputs[0]) .
  • (C5) 0 < window_dimensions .
  • (C6) size(window_strides) = rank(inputs[0]) .
  • (C7) 0 < window_strides .
  • (C8) size(base_dilations) = rank(inputs[0]) .
  • (C9) 0 < base_dilations .
  • (C10) size(window_dilations) = rank(inputs[0]) .
  • (C11) 0 < window_dilations .
  • (C12) shape(padding) = [rank(inputs[0]), 2] .
  • (C13) body 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.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) lhs tensor of integer, floating-point or complex type or per-tensor quantized tensor (C1)
(I2) rhs tensor of integer, floating-point or complex type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of integer, floating-point or complex type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

replica_id

مفاهیم

Produces replica_id of the current process.

خروجی ها

نام تایپ کنید
result 0-dimensional tensor of type ui32

مثال ها

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

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

reshape

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor or quantized tensor (C1-C3)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or quantized tensor (C1-C3)

محدودیت ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor or per-tensor quantized tensor (C1), (C3)
(I2) dimensions 1-dimensional tensor constant of type si64 (C2), (C3)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or per-tensor quantized tensor (C1), (C3)

محدودیت ها

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

مثال ها

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

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

rng

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) a 0-dimensional tensor of integer, boolean, or floating-point type (C1), (C2)
(I2) b 0-dimensional tensor of integer, boolean, or floating-point type (C1), (C2)
(I3) shape 1-dimensional tensor constant of type si64 (C3)
(I4) rng_distribution enum of UNIFORM and NORMAL (C2)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of integer, boolean, or floating-point type (C1-C3)

محدودیت ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) rng_algorithm enum of DEFAULT , THREE_FRY , and PHILOX (C2)
(I2) initial_state 1-dimensional tensor of type ui64 (C1), (C2)

خروجی ها

نام تایپ کنید محدودیت ها
output_state 1-dimensional tensor of type ui64 (C1)
output tensor of integer or floating-point type

محدودیت ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of floating-point type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of floating-point type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

round_nearest_even

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of floating-point type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of floating-point type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

rsqrt

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of floating-point or complex type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of floating-point or complex type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

پراکنده کردن

مفاهیم

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.

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) inputs variadic number of tensors or per-tensor quantized tensors (C1), (C2), (C4-C6), (C11), (C13), (C18), (C21), (C23-C24)
(I2) scatter_indices tensor of integer type (C4), (C15), (C19), (C22)
(I3) updates variadic number of tensors or per-tensor quantized tensors (C3-C6), (C8)
(I4) update_window_dims 1-dimensional tensor constant of type si64 (C2), (C4), (C7-C8)
(I5) inserted_window_dims 1-dimensional tensor constant of type si64 (C2), (C4), (C9-C11)
(I6) input_batching_dims 1-dimensional tensor constant of type si64 (C2), (C4), (C9), (C12-13), (C17-18), (C20)
(I7) scatter_indices_batching_dims 1-dimensional tensor constant of type si64 (C14-C18)
(I8) scatter_dims_to_operand_dims 1-dimensional tensor constant of type si64 (C19-C21)
(I9) index_vector_dim constant of type si64 (C4), (C16), (C19), (C22)
(I10) indices_are_sorted constant of type i1
(I11) unique_indices constant of type i1
(I12) update_computation تابع (C23)

خروجی ها

نام تایپ کنید محدودیت ها
results variadic number of tensors or per-tensor quantized tensors (C24-C25)

محدودیت ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) pred tensor of type i1 (C1)
(I2) on_true tensor or per-tensor quantized tensor (C1-C2)
(I3) on_false tensor or per-tensor quantized tensor (C2)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or per-tensor quantized tensor (C2)

محدودیت ها

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

مثال ها

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

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

select_and_scatter

مفاهیم

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.

select_and_scatter

More formally:

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor or per-tensor quantized tensor (C1-C4), (C6), (C8-C11)
(I2) source tensor or per-tensor quantized tensor (C1), (C2)
(I3) init_value 0-dimensional tensor or per-tensor quantized tensor (C3)
(I4) window_dimensions 1-dimensional tensor constant of type si64 (C2), (C4), (C5)
(I5) window_strides 1-dimensional tensor constant of type si64 (C2), (C6), (C7)
(I6) padding 2-dimensional tensor constant of type si64 (C2), (C8)
(I7) select تابع (C9)
(I8) scatter تابع (C10)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or per-tensor quantized tensor (C11-C12)

محدودیت ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) inputs variadic number of tensors or quantized tensors
(I2) token token
(I3) channel_id constant of type si64
(I4) channel_type enum of DEVICE_TO_DEVICE and DEVICE_TO_HOST (C1)
(I5) is_host_transfer constant of type i1 (C1)

خروجی ها

نام تایپ کنید
result token

محدودیت ها

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) lhs tensor of integer type (C1)
(I2) rhs tensor of integer type (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of integer type (C1)

محدودیت ها

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

مثال ها

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

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

shift_right_arithmetic

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) lhs tensor of integer type (C1)
(I2) rhs tensor of integer type (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of integer type (C1)

محدودیت ها

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

مثال ها

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

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

shift_right_logical

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) lhs tensor of integer type (C1)
(I2) rhs tensor of integer type (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of integer type (C1)

محدودیت ها

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

مثال ها

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

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

امضا کردن

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of signed integer, floating-point, or complex type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of signed integer, floating-point, or complex type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

سینوسی

مفاهیم

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of floating-point or complex type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of floating-point or complex type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

تکه

مفاهیم

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 .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor or per-tensor quantized tensor (C1-C3), (C5)
(I2) start_indices 1-dimensional tensor constant of type si64 (C2), (C3), (C5)
(I3) limit_indices 1-dimensional tensor constant of type si64 (C2), (C3), (C5)
(I4) strides 1-dimensional tensor constant of type si64 (C2), (C4)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or per-tensor quantized tensor (C1), (C5)

محدودیت ها

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

مثال ها

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

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

مرتب سازی

مفاهیم

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 .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) inputs variadic number of tensors or per-tensor quantized tensors (C1-C5)
(I2) dimension constant of type si64 (C4)
(I3) is_stable constant of type i1
(I4) comparator تابع (C5)

خروجی ها

نام تایپ کنید محدودیت ها
results variadic number of tensors or per-tensor quantized tensors (C2), (C3)

محدودیت ها

  • (C1) 0 < size(inputs) .
  • (C2) type(inputs...) = type(results...) .
  • (C3) same(shape(inputs...) + shape(results...)) .
  • (C4) -R <= dimension < R , where R = rank(inputs[0]) .
  • (C5) comparator has type (tensor<E1>, tensor<E1>, ..., tensor<EN-1>, tensor<EN-1>) -> tensor<i1> , where Ei = element_type(inputs[i]) .

مثال ها

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

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

sqrt

مفاهیم

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

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of floating-point or complex type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of floating-point or complex type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

تفریق کردن

مفاهیم

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

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) lhs tensor of integer, floating-point, or complex type or per-tensor quantized tensor (C1)
(I2) rhs tensor of integer, floating-point, or complex type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of integer, floating-point, or complex type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

tanh

مفاهیم

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

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of floating-point or complex type or per-tensor quantized tensor (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of floating-point or complex type or per-tensor quantized tensor (C1)

محدودیت ها

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

مثال ها

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

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

جابجا کردن

مفاهیم

Permutes the dimensions of operand tensor using permutation and produces a result tensor. More formally, result[result_index] = operand[operand_index] where result_index[d] = operand_index[permutation[d]] .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor or quantized tensor (C1-C4)
(I2) permutation 1-dimensional tensor constant of type si64 (C2-C4)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor or quantized tensor (C1), (C3-C4)

محدودیت ها

  • (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) permutation is a permutation of range(rank(operand)) .
  • (C3) shape(result) = dim(operand, permutation...) .
  • (C4) If is_per_axis_quantized(result) , then quantization_dimension(operand) = permutation(quantization_dimension(result)) .

مثال ها

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

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

triangular_solve

مفاهیم

Solves batches of systems of linear equations with lower or upper triangular coefficient matrices.

More formally, given a and b , result[i0, ..., iR-3, :, :] is the solution to op(a[i0, ..., iR-3, :, :]) * x = b[i0, ..., iR-3, :, :] when left_side is true or x * op(a[i0, ..., iR-3, :, :]) = b[i0, ..., iR-3, :, :] when left_side is false , solving for the variable x where op(a) is determined by transpose_a , which can be one of the following:

  • NO_TRANSPOSE : Perform operation using a as-is.
  • TRANSPOSE : Perform operation on transpose of a .
  • ADJOINT : Perform operation on conjugate transpose of a .

Input data is read only from the lower triangle of a , if lower is true or upper triangle of a , otherwise. Output data is returned in the same triangle; the values in the other triangle are implementation-defined.

If unit_diagonal is true, then the implementation can assume that the diagonal elements of a are equal to 1, otherwise the behavior is undefined.

For quantized types, performs dequantize_op_quantize(lambda x, y: triangular_solve(x, y, left_side, lower, unit_diagonal, transpose_a), a, b, type(result)) .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) a tensor of floating-point or complex type or per-tensor quantized tensor (C1-C3)
(I2) b tensor of floating-point or complex type or per-tensor quantized tensor (C1-C4)
(I3) left_side constant of type i1 (C3)
(I4) lower constant of type i1
(I5) unit_diagonal constant of type i1
(I6) transpose_a enum of NO_TRANSPOSE , TRANSPOSE , and ADJOINT

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of floating-point or complex type or per-tensor quantized tensor (C1)

محدودیت ها

  • (C1) baseline_element_type(a) = baseline_element_type(b) .
  • (C2) 2 <= rank(a) = rank(b) = R .
  • (C3) The relationship between shape(a) and shape(b) is defined as follows:
    • shape(a)[:-3] = shape(b)[:-3] .
    • dim(a, -2) = dim(a, -1) = dim(b, left_side ? -2 : -1) .
  • (C4) baseline_type(b) = baseline_type(result) .

مثال ها

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

چندتایی

مفاهیم

Produces a result tuple from values val .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) val variadic number of values (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result چندتایی (C1)

محدودیت ها

  • (C1) result has type tuple<E0, ..., EN-1> where Ei = type(val[i]) .

مثال ها

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

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

uniform_dequantize

مفاهیم

Performs element-wise conversion of quantized tensor operand to a floating-point tensor result according to the quantization parameters defined by the operand type.

More formally, result = dequantize(operand) .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand quantized tensor (C1), (C2)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of floating-point type (C1), (C2)

محدودیت ها

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

مثال ها

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

uniform_quantize

مفاهیم

Performs element-wise conversion of floating-point tensor or quantized tensor operand to a quantized tensor result according to the quantization parameters defined by the result type.

به طور رسمی تر،

  • If is_float(operand) :
    • result = quantize(operand, type(result)) .
  • If is_quantized(operand) :
    • float_result = dequantize(operand) .
    • result = quantize(float_result, type(result)) .

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand tensor of floating-point or quantized type (C1), (C2)

خروجی ها

نام تایپ کنید محدودیت ها
result quantized tensor (C1), (C2)

محدودیت ها

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

مثال ها

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

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

در حالی که

مفاهیم

Produces the output from executing body function 0 or more times while the cond function outputs true . More formally, the semantics can be expressed using Python syntax as follows:

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

The behavior of an infinite loop is TBD ( #383 ).

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) operand variadic number of tensors, quantized tensors or tokens (C1-C3)
(I2) cond تابع (C1)
(I3) body تابع (C2)

خروجی ها

نام تایپ کنید محدودیت ها
results variadic number of tensors, quantized tensors or tokens (C3)

محدودیت ها

  • (C1) cond has type (T0, ..., TN-1) -> tensor<i1> , where Ti = type(operand[i]) .
  • (C2) body has type (T0, ..., TN-1) -> (T0, ..., TN-1) , where Ti = type(operand[i]) .
  • (C3) type(results...) = type(operand...) .

مثال ها

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

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

xor

مفاهیم

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

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

ورودی ها

برچسب نام تایپ کنید محدودیت ها
(I1) lhs tensor of boolean or integer type (C1)
(I2) rhs tensor of boolean or integer type (C1)

خروجی ها

نام تایپ کنید محدودیت ها
result tensor of boolean or integer type (C1)

محدودیت ها

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

مثال ها

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

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

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

Dialect Interop

At the moment, StableHLO programs in the wild sometimes contain operations that are not defined by StableHLO.

Module, Function, Call and Return

StableHLO uses upstream MLIR operations for ModuleOp, FuncOp, CallOp, and ReturnOp. This was done for better interop with existing MLIR machinery, as many useful passes are written targeting FuncOp and ModuleOp, and many compilation pipelines expect these ops to be present. Full compatibility guarantees are applied to these ops. If anything ever changes about these ops in an incompatible way (ie removal), StableHLO equivalents will be added to preserve compatibility.

CHLO

The CHLO opset contains higher level operations that decompose to StableHLO. Currently there are no compatibility guarantees for CHLO. For compatibility guarantees, the chlo-legalize-to-stablehlo pass must be used prior to serialization.

Shape Operations

It is a common use case in the community to use certain operations from core MLIR dialects in dynamic StableHLO programs to perform shape computations. Most commonly, these include shape dialect ops like shape_of or num_elements , tensor dialect ops like dim or from_elements , and the builtin index type.

The Dynamism RFC > O2 denotes these as out of scope, however some support for index types is included for interop purposes. There are no compatibility guarantees for these ops or types. The shape-legalize-to-stablehlo pass can be used to convert these operations to fully supported StableHLO ops.

Deprecated Operations

There are several StableHLO operations that were inherited from MHLO which are deprecated and on the way out of StableHLO. The full details on these removals can be found in the StableHLO v1.0 Cleanup #2283 . The tracker issue for these deprecations is #2340 .

These operations fall into a few categories:

  • "Not in HLO" category of StableHLO operations - they were initially part of the StableHLO opset but have been later deemed to not fit it well: broadcast , create_token , cross-replica-sum , dot , einsum , torch_index_select , unary_einsum ( #3 ) .
  • Unused ops - These operations may have been useful at some point, but the ops were either underdeveloped, or the pipelines using these ops have been refactored to not require them anymore. This includes map , tuple ( #598 ), get_tuple_element , rng , complex comparisons #560 , and convolution window_reversal ( #1181 ).

Some of these ops can be removed easily given that they can be expressed using existing ops ( broadcast , create_token , cross-replica-sum , dot , unary_einsum ) and will be removed after the existing compatibilty window passes (6 months). Others are still being explored for removal ( einsum , get_tuple_element , map , rng torch_index_select , tuple , complex comparisons, window_reversal ). Pending community feedback, these ops will either be removed, or added to the spec with full support. Until these ops futures are known, they are only guaranteed 6 months of compatibility.

اجرا

Sequential execution

A StableHLO program is executed by providing input values to the main function and computing output values. Output values of a function are computed by executing the graph of ops rooted in the corresponding return op.

The execution order is implementation-defined as long as it is aligned with dataflow, ie if ops are executed before their uses. In StableHLO, all side-effecting ops consume one token and produce one token (multiple tokens can be multiplexed into one token via after_all ), so the execution order of side effects is also aligned with dataflow. For example, in the below program there are two possible execution orders: %0%1%2return and %1%0%2return .

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

More formally, a StableHLO process is a combination of: 1) a StableHLO program, 2) operation statuses (not executed yet, already executed), and 3) intermediate values that the process is working on. The process starts with input values to the main function, progresses through the graph of ops updating operation statuses and intermediate values and finishes with output values. Further formalization is TBD ( #484 ).

Parallel execution

StableHLO programs can be executed in parallel, organized into a 2D process grid of num_replicas by num_partitions which both have type ui32 .

In the StableHLO process grid , num_replicas * num_partitions of StableHLO processes are executing at the same time. Each process has a unique process_id = (replica_id, partition_id) , where replica_id in replica_ids = range(num_replicas) and partition_id in partition_ids = range(num_partitions) which both have type ui32 .

The size of the process grid is known statically for every program (in the future, we are planning to make it an explicit part of StableHLO programs #650 ), and the position within the process grid is known statically for every process. Each process has access to its position within the process grid via the replica_id and partition_id ops.

Within the process grid, the programs can all be the same (in the "Single Program, Multiple Data" style), can all be different (in the "Multiple Program, Multiple Data" style) or something in between. In the future, we are planning to introduce support for other idioms of defining parallel StableHLO programs, including GSPMD ( #619 ).

Within the process grid, the processes are mostly independent from each other - they have separate operation statuses, separate input/intermediate/output values and most of the ops are executed separately between processes, with the exception of a small number of collective ops described below .

Given that execution of most of the ops is only using values from the same process, it is usually unambiguous to refer to these values by their names. However, when describing semantics of collective ops, that is insufficient, and that gives rise to the notation name@process_id to refer to the value name within a particular process. (From that perspective, unqualified name can be viewed as a shorthand for name@(replica_id(), partition_id()) ).

The execution order across processes is implementation-defined, except for the synchronization introduced by point-to-point communication and collective ops as described below.

Point-to-point communication

StableHLO processes can communicate with each other through StableHLO channels . A channel is represented by a positive id of type si64 . Through various ops, it is possible to send values to channels and receive them from channels.

Further formalization, eg where these channel ids are coming from, how processes programs become aware of them and what kind of synchronization is introduced by them, is TBD ( #484 ).

Streaming communication

Every StableHLO process has access to two streaming interfaces:

  • Infeed that can be read from.
  • Outfeed that can be written to.

Unlike channels, which are used to communicate between processes and therefore have processes at both of their ends, infeeds and outfeeds have their other end implementation-defined.

Further formalization, eg how streaming communication influences execution order and what kind of synchronization is introduced by it, is TBD ( #484 ).

Collective ops

There are six collective ops in StableHLO: all_gather , all_reduce , all_to_all , collective_broadcast , collective_permute , and reduce_scatter . All these ops split the processes in the StableHLO process grid into StableHLO process groups and execute a joint computation within each process group, independently from other process groups.

Within each process group, collective ops may introduce a synchronization barrier. Further formalization, eg elaborating on when exactly this synchronization happens, how exactly the processes arrive at this barrier, and what happens if they don't, is TBD ( #484 ).

If the process group involves cross-partition communication, ie there are processes in the process group whose partition ids are different, then execution of the collective op needs a channel, and the collective op must provide a positive channel_id of type si64 . Cross-replica communication doesn't need channels.

The computations performed by the collective ops are specific to individual ops and are described in individual op sections above. However, the strategies by which the process grid is split into process groups are shared between these ops and are described in this section. More formally, StableHLO supports the following four strategies.

cross_replica

Only cross-replica communications happen within each process group. This strategy takes replica_groups - a list of lists of replica ids - and computes a Cartesian product of replica_groups by partition_ids . replica_groups must have unique elements and cover all replica_ids . More formally, using Python syntax:

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

For example, for replica_groups = [[0, 1], [2, 3]] and num_partitions = 2 , cross_replica will produce [[(0, 0), (1, 0)], [(0, 1), (1, 1)], [(2, 0), (3, 0)], [(2, 1), (3, 1)]] .

cross_partition

Only cross-partition communications happen within each process group. This strategy takes partition_groups - a list of lists of partition ids - and computes a Cartesian product of partition_groups by replica_ids . partition_groups must have unique elements and cover all partition_ids . More formally, using Python syntax:

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

For example, for partition_groups = [[0, 1]] and num_replicas = 4 , cross_partition will produce [[(0, 0), (0, 1)], [(1, 0), (1, 1)], [(2, 0), (2, 1)], [(3, 0), (3, 1)]] .

cross_replica_and_partition

Both cross-replica and cross-partition communications may happen within each process group. This strategy takes replica_groups - a list of lists of replica ids - and computes Cartesian products of each replica_group by partition_ids . replica_groups must have unique elements and cover all replica_ids . More formally, using Python syntax:

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

For example, for replica_groups = [[0, 1], [2, 3]] and num_partitions = 2 , cross_replica_and_partition will produce [[(0, 0), (1, 0), (0, 1), (1, 1)], [(2, 0), (3, 0), (2, 1), (3, 1)]] .

flattened_ids

This strategy takes flattened_id_groups - a list of lists of "flattened" process ids in the form of replica_id * num_partitions + partition_id - and turns them into process ids. flattened_id_groups must have unique elements and cover all process_ids . More formally, using Python syntax:

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

For example, for flattened_id_groups = [[0, 1, 2, 3], [4, 5, 6, 7]] , num_replicas = 4 and num_partitions = 2 , flattened_ids will produce [[(0, 0), (0, 1), (1, 0), (1, 1)], [(2, 0), (2, 1), (3, 0), (3, 1)]] .

دقت

At the moment, StableHLO does not provide guarantees about numerical accuracy, but this may change in the future ( #1156 ).

Execution semantics of quantized operation

The interpretation of quantized StableHLO operations may vary depending on the hardware requirements and capabilities. For instance, some hardware may opt to interpret quantized operations using a "dequantize, perform floating-point operation, and finally quantize" strategy. Others may perform the entire computation with integer arithmetic. Consequently, the interpretation of quantized StableHLO operations is exclusively determined by the specific implementation. The interpretation of hybrid quantization ( #1575 ) should be based on the it's semantics as prescribed in the specification (via 1792 ).

خطاها

StableHLO programs are validated through an extensive set of constraints for individual ops, which rules out many classes of errors prior to run time. However, error conditions are still possible, eg through integer overflows, out-of-bounds accesses, etc. Unless explicitly called out, all these errors result in implementation-defined behavior, but this may change in the future ( #1157 ).

Floating-point exceptions

As an exception to this rule, floating-point exceptions in StableHLO programs have well-defined behavior. Operations which result in exceptions defined by the IEEE-754 standard (invalid operation, division-by-zero, overflow, underflow, or inexact exceptions) produce default results (as defined in the standard) and continue execution without raising the corresponding status flag; similar to raiseNoFlag exception handling from the standard. Exceptions for nonstandard operations (eg complex arithmetic and certain transcendental functions) are implementation-defined.

Shape mismatches

StableHLO supports dynamically-shaped tensors. However, shapes have to agree at runtime, otherwise the behavior is undefined. StableHLO does not explicitly provide an op that can assert that a tensor has a given shape at runtime. Generating correct code is the responsibility of the producer.

As a specific example, the below program is valid. However, at runtime, the exact shapes of %arg0 and %arg1 will have to be the same, otherwise the behavior of the program is undefined:

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

نشانه گذاری

For describing syntax, this document is using the modified ISO flavor of EBNF syntax ( ISO/IEC 14977:1996 , Wikipedia ), with two modifications: 1) rules are defined using ::= rather than = ,

2) concatenation is expressed using juxtaposition rather than , .

For describing semantics (ie within "Types", "Constants" and "Ops" sections), we are using formulas which are based on Python syntax extended with support for concisely expressing array operations as described below. This works well for small snippets of code, but in rare cases when larger snippets of code are needed, we use vanilla Python syntax which is always introduced explicitly.

فرمول ها

Let's explore how formulas work based on an example from the dot_general specification. One of the constraints for this operation looks as follows: dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...) .

The names used in this formula come from two sources: 1) global functions, ie dim , 2) member definitions of the corresponding program element, ie lhs , lhs_batching_dimensions , rhs and rhs_batching_dimensions inputs defined in the "Inputs" section of dot_general .

As mentioned above, the syntax of this formula is Python-based with some conciseness-oriented extensions. To make sense of the formula, let's transform it into vanilla Python syntax.

A) In these formulas, we are using = to represent equality, so the first step towards obtaining Python syntax is replacing = with == , as follows: dim(lhs, lhs_batching_dimensions...) == dim(rhs, rhs_batching_dimensions...) .

B) Also, these formulas support ellipses ( ... ) which turn scalar expressions into tensor expressions. In a nutshell, f(xs...) roughly means "for each scalar x in the tensor xs , compute a scalar f(x) and then return all these scalar results together as a tensor result". In vanilla Python syntax, our example formula turns into: [dim(lhs, dim1) for dim1 in lhs_batching_dimensions] == [dim(rhs, dim2) for dim2 in rhs_batching_dimensions] .

Thanks to ellipses, it is often possible to avoid working at the level of individual scalars. However, in some tricky cases, lower-level semi-informal syntax may be used like in the start_indices[bi0, ..., :, ..., biN] formula from the gather specification. In the service of conciseness, we don't provide an exact formalism for translating such syntax to vanilla Python, in hopes that it is still intuitively understandable on case-by-case basis. Please let us know if some specific formulas look opaque, and we'll try to improve them.

Also, you will notice that formulas use ellipses to expand all sorts of lists, including tensors, lists of tensors (which eg can arise from a variadic number of tensors), etc. This is another area where we don't provide an exact formalism (eg lists are not even part of the StableHLO type system) and instead rely on intuitive understandability.

C) The final noteworthy notational vehicle that we employ is implicit broadcasting. While the StableHLO opset doesn't support implicit broadcasting, the formulas do, also in the service of conciseness. In a nutshell, if a scalar is used in a context where a tensor is expected, the scalar is broadcasted to the expected shape.

To continue the dot_general example, here's another constraint: 0 <= lhs_batching_dimensions < rank(lhs) . As defined in the dot_general specification, lhs_batching_dimensions is a tensor, however both 0 and rank(lhs) are scalars. After we apply implicit broadcasting, the formula will become [0, ..., 0] <= lhs_batching_dimensions < [rank(lhs), ..., rank(lhs)] .

When applied to a particular dot_general operation, this formula will evaluate to a tensor of booleans. When formulas are used as constraints, the constraint holds if the formula evaluates to either true or to a tensor which only has true elements.

نام ها

In formulas, lexical scope includes: 1) global functions, 2) member definitions,

3) local definitions. The list of global functions is provided below. The list of element definitions depends on the program element that the notation is applied to:

  • For operations, member definitions include names introduced in "Inputs" and "Outputs" sections.
  • For everything else, member definitions include structural parts of the program element, named after the corresponding EBNF non-terminals. Most of the time, the names of these structural parts are obtained by converting the names of the non-terminals to snake case (eg IntegerLiteral => integer_literal ), but sometimes names get abbreviated in the process (eg QuantizationStorageType => storage_type ) in which case the names are introduced explicitly similarly to "Inputs" / "Outputs" sections in operation specifications.
  • Additionally, member definitions always include self to refer to the corresponding program element.

ارزش های

When formulas are evaluated, they work with the following types of values: 1) Value (actual values, eg dense<[[1, 2], [3, 4]]> : tensor<2x2xi32> ; they always know their types), 2) Placeholder (future values, eg lhs , rhs or result ; their actual values are not known yet, only their types are known), 3) Type (types as defined in the "Types" section), 4) Function (global functions as defined in the "Functions" section).

Depending on the context, names may be referring to different values. More specifically, the "Semantics" section for ops (and equivalents for other program elements) defines runtime logic, so all inputs are available as Value . In contrast, the "Constraints" section for ops (and equivalents) defines "compile-time" logic, ie something that is typically executed before runtime, so only constant inputs are available as Value and other inputs are available only as Placeholder .

نام ها In "Semantics" In "Constraints"
Global functions Function Function
Constant inputs Value Value
Non-constant inputs Value Placeholder
خروجی ها Value Placeholder
Local definitions Depends on the definition Depends on the definition

Let's consider an example transpose operation:

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

For this operation, permutation is a constant, so it's available as a Value in both semantics and constraints. In contrast, operand and result are available as a Value in semantics but only as a Placeholder in constraints.

کارکرد

Construction of types

There are no functions that can be used to construct types. Instead, we directly use type syntax because it's typically more concise. Eg (tensor<E>, tensor<E>) -> (tensor<E>) rather than function_type( [tensor_type([], E), tensor_type([], E)], [tensor_type([], E)]) .

Functions on types

  • element_type is defined on tensor types and quantized tensor types and returns, respectively, the TensorElementType or QuantizedTensorElementType part of the corresponding TensorType or QuantizedTensorType .
def element_type(x: Value | Placeholder | Type):
 if type(x) == TensorType:
    return tensor_element_type(x)
  if type(x) == QuantizedTensorType:
    return quantized_tensor_element_type(x)
  if type(x) is not Type:
    return element_type(type(x))
  • is_per_axis_quantized(x: Value | Placeholder | Type) -> Value is a shortcut for is_quantized(x) and quantization_dimension(x) is not None .

  • is_per_tensor_quantized(x: Value | Placeholder | Type) -> Value is a shortcut for is_quantized(x) and quantization_dimension(x) is None .

  • is_promotable(x: Type, y: Type) -> bool checks if type x can be promoted to type y . When x and y are QuantizedTensorElementType s, the promotion is applied only to the storage_type . This specific version of promotion is currently used in context of reduction computation (refer to RFC for more details).

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

  if is_same_type == False:
    return False

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

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

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

  return false
  • is_quantized(x: Value | Placeholder | Type) -> Value is a shortcut for is_quantized_tensor_element_type(x) .

  • is_type_name(x: Value | Placeholder | Type) -> Value . Available for all types. For example, is_float(x) returns true if x is a FloatType . If x is a value or placeholder, this function is a shortcut for is_type_name(type(x)) .

  • max_value(x: Type) -> Value returns the maximum value of an TensorElementType . If x is not an TensorElementType , returns None .

  • min_value(x: Type) -> Value returns the minimum possible value of an TensorElementType . If x is not an TensorElementType , returns None .

  • member_name(x: Value | Placeholder | Type) -> Any . Available for all member definitions member_name of all types. For example, tensor_element_type(x) returns the TensorElementType part of a corresponding TensorType . If x is a value or placeholder, this function is a shortcut for member_name(type(x)) . If x is not a type that has an appropriate member, or a value or a placeholder of such a type, returns None .

Construction of values

  • operation_name(*xs: Value | Type) -> Value . Available for all operations. For example, add(lhs, rhs) takes two tensor values lhs and rhs and returns the output of evaluating the add operation with these inputs. For some operations eg broadcast_in_dim , types of their outputs are "load-bearing", ie needed to evaluate an operation. In this case, the function takes these types as arguments.

Functions on values

  • All Python's operators and functions are available. Eg both subscription and slicing notations from Python are available to index into tensors, quantized tensors and tuples.

  • to_destination_type(x: Value, destination_type: Type) -> Value is defined on tensors and returns the converted value of x based on the type(x) and destination_type as follows:

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

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

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

  return convert(x, destination_type)

There is early discussion on merging convert , uniform_quantize and uniform_dequantize operations ( #1576 ). After the merge we do not need the above function and can use the operation name for convert instead.

  • is_nan(x: Value) -> Value is defined on tensors and returns true if all elements of x are NaN or false otherwise. If x is not a tensor, returns None .

  • is_sorted(x: Value) -> Value is defined on tensors and returns true if elements of x are sorted in ascending order with respect to the ascending lexicographical order of their indices or false otherwise. If x is not a tensor, returns None .

  • is_unique(x: Value) -> Value is defined on tensors and returns true if x doesn't have duplicate elements or false otherwise. If x is not a tensor, returns None .

  • member_name(x: Value) -> Any is defined for all member definitions member_name of all values. For example, real_part(x) returns the RealPart part of a corresponding ComplexConstant . If x is not a value that has an appropriate member, returns None .

  • same(x: Value) -> Value is defined on tensors and returns true if elements of x are all equal to each other or false otherwise. If the tensor doesn't have elements, that counts as "all equal to each other", ie the function returns true . If x is not a tensor, returns None .

  • split(x: Value, num_results: Value, axis: Value) -> Value is defined on tensors and returns num_results slices of x along the axis axis . If x is not a tensor or dim(x, axis) % num_results != 0 , returns None .

  • is_defined_in_parent_scope(x: Value) -> Value is defined on strings and returns true if x is the name of a function defined in the same scope as the parent function of the relevant op.

  • is_namespaced_op_name(x: Value) -> Value is defined on strings and returns true if x is a valid op name, that is it respects the following regular expression: [a-zA-Z][a-zA-Z0-9_]*([.][a-zA-Z0-9_$]+)+

Shape computations

  • axes(x: Value | Placeholder | Type) -> Value is a shortcut for range(rank(x)) .

  • dim(x: Value | Placeholder | Type, axis: Value) -> Value is a shortcut for shape(x)[axis] .

  • dims(x: Value | Placeholder | Type, axes: List) -> List is a shortcut for list(map(lambda axis: dim(x, axis), axes)) .

  • index_space(x: Value | Placeholder | Type) -> Value is defined on tensors and returns size(x) indices for the corresponding TensorType sorted in ascending lexicographical order, ie [0, ..., 0] , [0, ..., 1] , ..., shape(x) - 1 . If x is not a tensor type, a quantized tensor type, or a value or a placeholder of one of these types, returns None .

  • rank(x: Value | Placeholder | Type) -> Value is a shortcut for size(shape(x)) .

  • shape(x: Value | Placeholder | Type) -> Value is defined in the "Functions on types" section via member_name .

  • size(x: Value | Placeholder | Type) -> Value is a shortcut for reduce(lambda x, y: x * y, shape(x)) .

Quantization computations

  • def baseline_element_type(x: Value | Placeholder | Type) -> Type is a shortcut for element_type(baseline_type(x)) .

  • baseline_type is defined on tensor types and quantized tensor types and transforms them to a "baseline", ie a type with the same shape but with the quantization parameters of the element type reset to default values. This is used as a handy trick to compare both tensor and quantized tensor types uniformly, which is needed quite often. For quantized types, this enables comparing types ignoring the quantization parameters, that is, shape , storage_type , expressed_type , storage_min , storage_max , and quantization_dimension (for per-axis quantized type) must all match, but scales and zero points may differ.

def baseline_type(x: Value | Placeholder | Type) -> Type:
  if type(x) == TensorType:
    return x
  if type(x) == QuantizedTensorType:
    element_type = quantized_tensor_element_type(x)
    baseline_element_type = QuantizedTensorElementType(
      storage_type = storage_type(element_type),
      storage_min = storage_min(element_type),
      storage_max = storage_max(element_type),
      expressed_type = expressed_type(element_type),
      quantization_dimension = quantization_dimension(element_type),
      scales = [constant(1.0, expressed_type(element_type))] * dim(x, quantization_dimension(element_type)),
      zero_points = [constant(0, storage_type(element_type))] * dim(x, quantization_dimension(element_type)))
    return QuantizedTensorType(shape(x), baseline_element_type)
  if type(x) is not Type:
    return baseline_element_type(type(x))
  • dequantize is defined on quantized tensor types and turns them into floating-point tensor types. This happens via converting quantized elements which represent integer values of the storage type into corresponding floating-point values of the expressed type using the zero point and scale associated with the quantized element type.
def compute_zero_points(quantized_type, result_type):
  if is_per_tensor_quantized(quantized_type):
    return broadcast_in_dim(constant(zero_point(quantized_type), storage_type(quantized_type)), [], result_type)
  if is_per_axis_quantized(quantized_type):
    for i in index_space(result_type):
      d = quantization_dimension(quantized_type)
      zero_points[i] = zero_points(quantized_type)[i[d]]
    return zero_points

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

def dequantize(x: Value) -> Value:
  assert is_quantized(x)
  x_storage = bitcast_convert(x, storage_type(x))
  x_storage_sub = x_storage - compute_zero_points(type(x), type(x_storage))
  x_expressed_sub = convert(x_storage_sub, expressed_type(x))
  return x_expressed_sub * compute_scales(type(x), type(x_expressed_sub))
  • quantize is defined on floating-point tensor types and turns them into quantized tensor types. This happens via converting floating-point values of the expressed type into corresponding integer values of the storage type using the zero point and scale associated with the quantized element type.
def quantize(x: Value, result_type: Type) -> Value:
  assert is_float(x) and is_quantized(result_type)
  zero_points = compute_zero_points(result_type, TensorType(shape(x), storage_type(result_type)))
  converted_zero_points = convert(zero_points, expressed_type(result_type))
  converted_min = convert(storage_min(result_type), expressed_type(result_type))
  converted_max = convert(storage_max(result_type), expressed_type(result_type))

  x_scaled = x / compute_scales(result_type, type(x))
  x_scaled_add_zp = x_scaled + converted_zero_points
  x_clamped = clamp(converted_min, x_scaled_add_zp, converted_max)
  x_rounded = round_nearest_even(x_clamped)
  return convert(x_rounded, result_type)
  • dequantize_op_quantize is used to specify element-wise computations on quantized tensors. It dequantizes, ie turns quantized elements into their expressed types, then performs an operation, and then quantizes, ie turns the results back into their storage types. At the moment, this function only works for per-tensor quantization. Per-axis quantization is work in progress ( #1574 ).
def dequantize_op_quantize(op, *inputs_and_output_type):
  inputs = inputs_and_output_type[:-1]
  output_type = inputs_and_output_type[-1]

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

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

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

def dequantize_select_quantize(pred, on_true, on_false, output_type):
  float_on_true = dequantize(on_true)
  float_on_false = dequantize(on_false)
  float_result = select(pred, float_on_true, float_on_false)
  return quantize(float_result, output_type)
  • hybrid_dequantize_then_op is used to specify weight-only quantization for hybrid op which accepts lhs in floating-point and rhs in quantized types. It dequantizes quantized inputs into their expressed types and performs computation in float. Element type of float lhs tensor and expressed type of quantized rhs tensor should be identical.
def hybrid_dequantize_then_op(op, lhs, rhs):
  assert(is_float(lhs) and is_quantized(rhs) and element_type(lhs) == expressed_type(rhs))
  return op(lhs, dequantize(rhs))

Grid computations

  • cross_partition(replica_groups: Value) -> Value . See the "cross_replica" section above.

  • cross_replica(replica_groups: Value) -> Value . See the "cross_replica" section above.

  • cross_replica_and_partition(replica_groups: Value) -> Value . See the "cross_replica_and_partition" section above.

  • flattened_ids(replica_groups: Value) -> Value . See the "flattened_ids" section above.

پویایی

StableHLO values can have dynamic dimension sizes, eg tensor<?xi64> . However, StableHLO values cannot have a dynamic number of dimensions (unranked dynamism, eg tensor<*xi64> ). Operands and results are allowed to use dynamic dimension sizes, even if there are constraints on the sizes. Constraints will be verified statically if possible, otherwise they are deferred to runtime and mismatches will result in undefined behavior. برای نمونه به زیر مراجعه کنید.

Shape mismatches for unary elementwise operations

Consider the following toy program:

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

Such a program is unusual, because it is not common to know the shape of the result but not the shape of the input. Nonetheless, this is a valid StableHLO program. It is not possible to statically validate the abs operation in this program, because the exact shape of the operand is unknown. However, the shapes are certainly compatible, and this can be checked statically: ? could turn out to be 2 at runtime, and there would be no issue. با این حال، ? could also turn out to be some other integer, in which case the behavior is undefined.

Note that if a dimension size is dynamic in the result, there cannot be undefined behavior. Indeed, there is no "expected" size, so there cannot be a mismatch.

Shape mismatches for binary elementwise operations

Consider the following toy program:

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

When it comes to binary elementwise operations, the shapes of the inputs and the result must agree at runtime. At compile time, static dimensions must be equal, otherwise they merely need to be compatible. If any dimension is dynamic in the inputs, then there could be undefined behavior at runtime, because the dynamic size may not match the corresponding size in the other operand (be it static or dynamic). If all the inputs are static, then whether the result is dynamic or not does not matter: statically known dimensions will be checked statically, and dynamic dimensions do not impose any constraints.

Shape mismatches for ops that take their output shape as an operand

Consider the following toy program:

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

The values in the shape operand at runtime must match the shape of the result, otherwise the behavior is undefined. That is, at runtime %arg0 must have a value of dense<[3, 4]> : tensor<2xi32> . If the shape operand is constant, this can be verified statically. If the result shape is fully dynamic, then there cannot be a mismatch.