StableHLO 사양

StableHLO는 머신러닝 (ML) 모델의 고급 작업 (HLO)을 위한 작업 집합입니다. StableHLO는 다양한 ML 프레임워크와 ML 컴파일러 간의 이동성 레이어로 작동합니다. StableHLO 프로그램을 생성하는 ML 프레임워크는 StableHLO 프로그램을 사용하는 ML 컴파일러와 호환됩니다.

Google의 목표는 다양한 ML 프레임워크 (예: TensorFlow, JAX, PyTorch)와 ML 컴파일러 (예: XLA, IREE) 간의 상호 운용성을 높여 ML 개발을 간소화하고 가속화하는 것입니다. 이를 위해 이 문서에서는 StableHLO 프로그래밍 언어의 사양을 제공합니다.

이 사양에는 세 가지 주요 섹션이 포함되어 있습니다. 먼저 프로그램 섹션에서는 StableHLO 함수로 구성되며 StableHLO 연산으로 구성된 StableHLO 프로그램의 구조를 설명합니다. 이 구조 내에서 연산 섹션은 개별 연산의 시맨틱을 지정합니다. 실행 섹션은 프로그램 내에서 함께 실행되는 이러한 모든 작업에 대한 의미 체계를 제공합니다. 마지막으로 표기법 섹션에서는 사양 전체에서 사용되는 표기법을 설명합니다.

이전 StableHLO 출시의 사양을 보려면 관심 있는 태그된 출시 버전에서 저장소를 여세요. 예를 들어 StableHLO v0.19.0 사양을 들 수 있습니다. StableHLO의 각 마이너 버전 버프에서 발생한 변경사항을 보려면 VhloDialect.td의 버전 로그를 참고하세요.

프로그램

Program ::= {Func}

StableHLO 프로그램은 임의 개수의 StableHLO 함수로 구성됩니다. 다음은 입력 3개(%image, %weights, %bias)와 출력 1개가 있는 @main 함수가 포함된 프로그램 예입니다. 함수의 본문에는 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인 텐서 유형입니다. 크기가 2와 3인 0차원과 1차원이라는 두 개의 차원(즉, 두 축)이 있습니다. 순위는 2입니다.

도형은 부분적으로 또는 완전히 알 수 없을 수 있습니다 (동적). 예를 들어 tensor<?x2xf64>는 부분적으로 알 수 없고 tensor<?x?xf64>는 완전히 알 수 없습니다. 동적 측정기준 크기는 ?를 사용하여 표현됩니다. 도형은 순위를 지정할 수 없습니다.

앞으로는 예를 들어 레이아웃(#629)과 희소성을 포함하도록 차원 크기 및 요소 유형 이상으로 텐서 유형을 확장할 계획입니다(#1078).

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

양자화된 요소 유형은 표현된 유형의 부동 소수점 값에 해당하는 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)

유형, 값, 양자화 텐서 유형에 0점이 하나만 있거나 여러 개 있을 수 있는지 여부 등 QuantizationZeroPoint의 의미 체계에 대한 논의가 계속되고 있습니다. 이 논의의 결과에 따라 0점에 관한 사양은 향후 변경될 수 있습니다. (#1405)

진행 중인 또 다른 논의는 이러한 값과 정량화된 텐서의 값에 제약 조건을 적용해야 하는지 결정하기 위한 QuantizationStorageMin 및 QuantizationStorageMax의 시맨틱스와 관련이 있습니다(#1406).

마지막으로, 알 수 없는 크기를 나타내는 방법 (#1407)과 마찬가지로 알 수 없는 크기를 나타내는 방법을 살펴볼 계획입니다.

양자화 텐서 유형은 양자화 요소가 있는 텐서를 나타냅니다. 이러한 텐서는 요소에 일반 요소 유형 대신 정량화된 요소 유형이 있다는 점을 제외하고 일반 텐서와 정확히 동일합니다.

정규화된 텐서에서 정규화는 텐서별일 수 있습니다. 즉, 전체 텐서에 하나의 scale 및 zero_point가 있거나 축별일 수 있습니다. 즉, 특정 크기 quantization_dimension의 슬라이스당 하나씩 여러 개의 scales 및 zero_points가 있을 수 있습니다. 더 공식적으로 축당 양자화를 사용하는 텐서 t에는 quantization_dimension의 dim(t, quantization_dimension) 슬라이스(t[:, ..., 0, ..., :], t[:, ..., 1, ..., :] 등)가 있습니다. i번째 슬라이스의 모든 요소는 scales[i] 및 zero_points[i]를 양자화 매개변수로 사용합니다. 정규화된 텐서 유형에는 다음과 같은 제약 조건이 있습니다.

• 텐서별 양자화의 경우:
  • 추가 제약 조건이 없습니다.
• 축당 양자화의 경우:
  • (C12) quantization_dimension < rank(self).
  • (C13) dim(self, quantization_dimension) = size(scales)입니다.

TokenType ::= 'token'

토큰 유형은 토큰을 나타냅니다. 즉, 일부 작업에서 생성되고 소비되는 불투명한 값입니다. 토큰은 실행 섹션에 설명된 대로 작업에 실행 순서를 적용하는 데 사용됩니다.

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

튜플 유형은 튜플, 즉 이종 목록을 나타냅니다. 튜플은 HLO와의 호환성을 위해만 존재하는 기존 기능입니다. HLO에서 튜플은 가변 입력과 출력을 나타내는 데 사용됩니다. StableHLO에서는 가변 입력과 출력이 기본적으로 지원되며 StableHLO에서 튜플을 사용하는 유일한 목적은 HLO ABI를 포괄적으로 나타내는 것입니다. 예를 들어 T, tuple<T>, tuple<tuple<T>>는 특정 구현에 따라 실질적으로 다를 수 있습니다. 향후 StableHLO에서 튜플 유형을 삭제할 수 있는 HLO ABI를 변경할 계획입니다(#598).

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

요소 유형은 텐서 유형의 요소를 나타냅니다. 많은 프로그래밍 언어와 달리 이러한 유형은 StableHLO의 최우선 클래스가 아닙니다. 즉, StableHLO 프로그램은 이러한 유형의 값을 직접 표현할 수 없습니다. 따라서 T 유형의 스칼라 값을 tensor<T> 유형의 0차원 텐서 값으로 표현하는 것이 관례입니다.

• 불리언 유형은 불리언 값 true 및 false을 나타냅니다.
• 정수 유형은 부호 있는(si) 유형 또는 부호 없는(ui) 유형일 수 있으며 지원되는 비트 너비(2, 4, 8, 16, 32 또는 64) 중 하나를 가질 수 있습니다. 부호 있는 siN 유형은 -2^(N-1)과 2^(N-1)-1 사이의 정수 값(양 끝값 포함)을 나타내고 부호 없는 uiN 유형은 0와 2^N-1 사이의 정수 값(양 끝값 포함)을 나타냅니다.
• 부동 소수점 유형은 다음 중 하나일 수 있습니다.
  • IEEE-754 규칙에 따른 f8E3M4, f8E4M3, f8E5M2 8비트 부동 소수점 숫자입니다.
  • 딥 러닝을 위한 FP8 형식에 설명된 FP8 형식의 E4M3 및 E5M2 인코딩에 각각 해당하는 f8E4M3FN 및 f8E5M2 유형
  • 심층신경망용 8비트 수치 형식에 설명된 FP8 형식의 E4M3 및 E5M2 인코딩에 해당하는 f8E4M3FNUZ 및 f8E5M2FNUZ 유형
  • 하이브리드 8비트 부동 소수점(HFP8) 심층신경망 학습 및 추론에 설명된 FP8 형식의 E4M3 인코딩에 해당하는 f8E4M3B11FNUZ 유형입니다.
  • BFloat16: Cloud TPU에서 고성능을 얻는 비결에 설명된 bfloat16 형식에 해당하는 bf16 유형입니다.
  • IEEE 754 표준에 설명된 binary16('절반 정밀도'), binary32('단일 정밀도') 및 binary64('배정밀도') 형식에 각각 해당하는 f16, f32, f64 유형입니다.
  • tf32 유형은 TensorFloat32 형식에 해당하며 StableHLO에서 제한적으로 지원됩니다.
  • OCP Microscaling Formats Specification에 설명된 f4E2M1FN, f6E2M3FN, f6E3M2FN, f8E8M0FNU MX (마이크로 확장) 유형
• 복소수 유형은 동일한 요소 유형의 실부와 허수를 가진 복합 값을 나타냅니다. 지원되는 복합 유형은 complex<f32> (두 부분 모두 f32 유형) 및 complex<f64>(두 부분 모두 f64 유형)입니다.

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

함수 유형은 이름이 지정된 함수와 익명 함수를 모두 나타냅니다. 입력 유형 (-> 왼쪽의 유형 목록)과 출력 유형(-> 오른쪽의 유형 목록)이 있습니다. 많은 프로그래밍 언어에서 함수 유형은 퍼스트 클래스이지만 StableHLO에서는 그렇지 않습니다.

StringType ::= 'string'

문자열 유형은 바이트 시퀀스를 나타냅니다. 많은 프로그래밍 언어와 달리 문자열 유형은 StableHLO의 첫 번째 클래스가 아니며 프로그램 요소의 정적 메타데이터를 지정하는 데만 사용됩니다.

운영

StableHLO 작업 (작업이라고도 함)은 머신러닝 모델에서 상위 수준의 비공개 작업을 나타냅니다. 위에서 설명한 대로 StableHLO 문법은 MLIR에서 많은 영향을 받았습니다. MLIR는 반드시 가장 인체공학적인 대안은 아니지만 ML 프레임워크와 ML 컴파일러 간의 상호 운용성을 높이려는 StableHLO의 목표에 가장 적합합니다.

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

StableHLO 작업(작업이라고도 함)에는 이름, 입력/출력, 서명이 있습니다. 이름은 stablehlo. 접두사와 지원되는 작업 중 하나를 고유하게 식별하는 기호로 구성됩니다. 지원되는 모든 작업의 포괄적인 목록은 아래를 참고하세요.

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

작업은 입력을 소비하고 출력을 생성합니다. 입력은 입력 값(실행 중에 계산됨), 입력 함수(StableHLO에서 함수가 최상위 값이 아니므로 정적으로 제공됨), 입력 속성(정적으로 제공됨)으로 분류됩니다. 연산에서 소비하고 생성하는 입력과 출력의 종류는 니모닉에 따라 다릅니다. 예를 들어 add 작업은 입력 값 2개를 사용하고 출력 값 1개를 생성합니다. 이에 비해 select_and_scatter 작업은 입력 값 3개, 입력 함수 2개, 입력 속성 3개를 사용합니다.

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

입력 함수(익명 함수라고도 함)는 1) 식별자가 없고(따라서 '익명'이라는 이름이 붙음) 2) 출력 유형을 선언하지 않는다는 점을 제외하고는 이름이 지정된 함수와 매우 유사합니다(출력 유형은 함수 내의 return 연산자에서 추론됨).

입력 함수의 구문에는 MLIR와의 호환성을 위해 현재 사용되지 않는 부분(위의 Unused 프로덕션 참고)이 포함되어 있습니다. MLIR에는 점프 작업을 통해 함께 연결된 여러 작업 '블록'을 포함할 수 있는 보다 일반적인 '리전' 개념이 있습니다. 이러한 블록에는 Unused 프로덕션에 해당하는 ID가 있으므로 서로 구분할 수 있습니다. StableHLO에는 점프 연산자가 없으므로 MLIR 문법의 해당 부분은 사용되지 않지만 여전히 존재합니다.

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

입력 속성에는 이름과 지원되는 상수 중 하나인 값이 있습니다. 프로그램 요소의 정적 메타데이터를 지정하는 기본 방법입니다. 예를 들어 concatenate 연산자는 dimension 속성을 사용하여 입력 값이 연결되는 측정기준을 지정합니다. 마찬가지로 slice 작업은 start_indices 및 limit_indices와 같은 여러 속성을 사용하여 입력 값을 슬라이스하는 데 사용되는 경계를 지정합니다.

현재 실제 StableHLO 프로그램에는 이 문서에 설명되지 않은 속성이 포함되어 있는 경우가 있습니다. 향후 이러한 속성을 StableHLO opset에 통합하거나 StableHLO 프로그램에 표시되지 않도록 할 계획입니다. 그동안 다음과 같은 속성 목록을 확인하세요.

• layout(#629)
• mhlo.frontend_attributes (#628).
• mhlo.sharding(#619)
• output_operand_aliases (#740).
• 위치 메타데이터 (#594)

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

연산자 서명은 모든 입력 값의 유형(-> 왼쪽의 유형 목록)과 모든 출력 값의 유형(-> 오른쪽의 유형 목록)으로 구성됩니다. 엄밀히 말하면 입력 유형은 중복되고 출력 유형도 거의 항상 중복됩니다(대부분의 StableHLO 연산자의 경우 출력 유형을 입력에서 추론할 수 있기 때문). 그럼에도 불구하고 op 서명은 MLIR와의 호환성을 위해 의도적으로 StableHLO 문법의 일부입니다.

다음은 mnemonik이 select_and_scatter인 연산의 예입니다. 이 연산자는 3개의 입력 값(%operand, %source, %init_value), 2개의 입력 함수, 3개의 입력 속성(window_dimensions, window_strides, padding)을 사용합니다. 연산자의 서명에는 입력 값의 유형만 포함되고 인라인으로 제공되는 입력 함수 및 속성의 유형은 포함되지 않습니다.

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

상수

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

StableHLO 상수에는 StableHLO 값을 함께 나타내는 리터럴과 유형이 있습니다. 일반적으로 유형은 모호하지 않은 경우를 제외하고 상수 문법의 일부입니다 (예: 불리언 상수는 유형이 i1임이 명확하지만 정수 상수는 여러 유형이 가능함).

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

부울 상수는 불리언 값 true 및 false를 나타냅니다. 불리언 상수의 유형은 i1입니다.

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

정수 상수는 십진수 또는 16진수 표기법을 사용하는 문자열을 통해 정수 값을 나타냅니다. 다른 기수(예: 바이너리 또는 옥텟)는 지원되지 않습니다. 정수 상수에는 다음과 같은 제약이 있습니다.

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

부동 소수점 상수는 십진수 또는 과학적 표기법을 사용하는 문자열을 통해 부동 소수점 값을 나타냅니다. 또한 16진수 표기법을 사용하여 상응하는 유형의 부동 소수점 형식으로 기본 비트를 직접 지정할 수 있습니다. 부동 소수점 상수에는 다음과 같은 제약이 있습니다.

• (C1) 16진수 이외의 표기법이 사용된 경우 is_wellformed(float_literal, float_type).
• (C2) 16진수 표기법이 사용되는 경우 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 텐서를 생성합니다. 요소 유형에 따라 다음 작업을 실행합니다.

• 부호 있는 정수의 경우: 정수 모듈러스
• 부동 소수점 수: IEEE-754의 abs
• 복소수: 복소수
• 정규화된 유형: dequantize_op_quantize(abs, operand, type(result)).

입력

출력

제약조건

• (C1) shape(result) = shape(operand).
• (C2) baseline_element_type(result)는 다음과 같이 정의됩니다.
  • is_complex(operand)인 경우 complex_element_type(element_type(operand))
  • 그 외에는 baseline_element_type(operand)입니다.

예

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

 예시 더보기

추가

시맨틱

두 텐서 lhs 및 rhs의 요소별 추가를 수행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 불리언: 논리합(OR)
• 정수의 경우 정수 덧셈
• 부동 소수점의 경우: IEEE-754의 addition
• 복소수의 경우: 복소수 덧셈
• 정규화된 유형: dequantize_op_quantize(add, lhs, rhs, type(result)).

입력

출력

제약조건

• 작업에서 정규화되지 않은 텐서를 사용하는 경우:
  • (C1) type(lhs) = type(rhs) = type(result).
• 작업에서 정규화된 텐서를 사용하는 경우:
  • (C2) is_quantized(lhs) and is_quantized(rhs) and is_quantized(result).
  • (C3) storage_type(lhs) = storage_type(rhs) = storage_type(result).
  • (C4) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C5) (is_per_axis_quantized(lhs) or is_per_axis_quantized(rhs)) = is_per_axis_quantized(result).
  • (C6) is_per_axis_quantized(lhs)이면 quantization_dimension(lhs) = quantization_dimension(result).
  • (C7) is_per_axis_quantized(rhs)이면 quantization_dimension(rhs) = quantization_dimension(result)입니다.

예

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

예시 더보기

after_all

시맨틱

inputs를 생성하는 작업이 result에 종속된 작업 전에 실행되도록 합니다. 이 작업을 실행해도 아무 일도 일어나지 않습니다. 이 작업은 result에서 inputs로의 데이터 종속 항목을 설정하기 위한 용도로만 존재합니다.

입력

출력

예

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

 예시 더보기

all_gather

시맨틱

StableHLO 프로세스 그리드의 각 프로세스 그룹 내에서 all_gather_dim을 따라 각 프로세스의 operands 텐서 값을 연결하고 results 텐서를 생성합니다.

이 작업은 StableHLO 프로세스 그리드를 다음과 같이 정의된 process_groups로 분할합니다.

• channel_id <= 0 and use_global_device_ids = false인 경우 cross_replica(replica_groups)입니다.
• channel_id > 0 and use_global_device_ids = false인 경우 cross_replica_and_partition(replica_groups)
• channel_id > 0 and use_global_device_ids = true인 경우 flattened_ids(replica_groups)입니다.

그런 다음 각 process_group 내에서 다음을 실행합니다.

• process_group의 모든 receiver에 대해 operands...@receiver = [operand@sender for sender in process_group]
• process_group의 모든 process에 대해 results...@process = concatenate(operands...@process, all_gather_dim)

입력

출력

제약조건

• (C1) 0 <= all_gather_dim < rank(operands...).
• (C2) is_unique(replica_groups).
• (C3) size(replica_groups)는 다음과 같이 정의됩니다.
  • cross_replica가 사용된 경우 num_replicas
  • cross_replica_and_partition가 사용된 경우 num_replicas
  • flattened_ids가 사용되는 경우 num_processes
• (C4) 0 <= replica_groups < size(replica_groups).
• (C5) use_global_device_ids = true인 경우 channel_id > 0
• (C6) type(results...) = type(operands...) 예외:
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1).

예

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

예시 더보기

all_reduce

시맨틱

StableHLO 프로세스 그리드의 각 프로세스 그룹 내에서 각 프로세스의 operands 텐서 값에 감소 함수 computation를 적용하고 results 텐서를 생성합니다.

이 작업은 StableHLO 프로세스 그리드를 다음과 같이 정의된 process_groups로 분할합니다.

• channel_id <= 0 and use_global_device_ids = false인 경우 cross_replica(replica_groups)입니다.
• channel_id > 0 and use_global_device_ids = false인 경우 cross_replica_and_partition(replica_groups)
• channel_id > 0 and use_global_device_ids = true인 경우 flattened_ids(replica_groups)입니다.

그런 다음 각 process_group 내에서 다음을 실행합니다.

• 일부 바이너리 트리의 경우 results...@process[result_index] = exec(schedule) schedule 여기서:
  • exec(node) = computation(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value
• schedule는 순서대로 탐색이 to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0]))인 구현 정의 바이너리 트리입니다.

입력

출력

제약조건

• (C1) is_unique(replica_groups).
• (C2) size(replica_groups)은 다음과 같이 정의됩니다.
  • cross_replica가 사용된 경우 num_replicas
  • cross_replica_and_partition가 사용된 경우 num_replicas
  • flattened_ids가 사용된 경우 num_processes
• (C3) 0 <= replica_groups < size(replica_groups).
• (C4) use_global_device_ids = true이면 channel_id > 0.
• (C5) computation의 유형은 (tensor<E>, tensor<E>) -> (tensor<E>)입니다. 여기서 is_promotable(element_type(operand), E)입니다.
• (C6) shape(results...) = shape(operands...).
• (C7) element_type(results...) = E.

예

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

예시 더보기

all_to_all

시맨틱

StableHLO 프로세스 그리드의 각 프로세스 그룹 내에서 split_dimension을 따라 operands 텐서의 값을 부분으로 분할하고, 분할된 부분을 프로세스 간에 흩어 뿌리고, concat_dimension을 따라 흩어진 부분을 연결하여 results 텐서를 생성합니다. 이 작업은 StableHLO 프로세스 그리드를 다음과 같이 정의된 process_groups로 분할합니다.

• channel_id <= 0인 경우 cross_replica(replica_groups)
• channel_id > 0인 경우 cross_partition(replica_groups)

그런 다음 각 process_group 내에서 다음을 실행합니다.

• process_group의 모든 sender에 대한 split_parts...@sender = split(operands...@sender, split_count, split_dimension)
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group], 여기서 receiver_index = process_group.index(receiver).
• results...@process = concatenate(scattered_parts...@process, concat_dimension).

입력

출력

제약조건

• (C1) 0 <= split_dimension < rank(operands...).
• (C2) dim(operands..., split_dimension) % split_count = 0.
• (C3) 0 <= concat_dimension < rank(operands...).
• (C4) 0 < split_count.
• (C5) is_unique(replica_groups).
• (C6) size(replica_groups)는 다음과 같이 정의됩니다.
  • cross_replica가 사용된 경우 num_replicas
  • cross_partition가 사용된 경우 num_partitions
• (C7) 0 <= replica_groups < size(replica_groups).
• (C8) dim(replica_groups, 1) = split_count.
• (C9) split_dimension != concat_dimension인 경우 다음을 제외하고는 type(results...) = type(operands...)입니다.
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count.
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count.

예

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

 예시 더보기

및

시맨틱

두 텐서 lhs 및 rhs의 요소별 AND를 실행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 불리언: 논리곱
• 정수의 경우: 비트별 AND

입력

출력

제약조건

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

예

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

예시 더보기

atan2

시맨틱

lhs 및 rhs 텐서에서 요소별 atan2 연산을 실행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 부동 소수점의 경우: IEEE-754의 atan2
• 복소수의 경우 복소수 atan2입니다.
• 정규화된 유형: dequantize_op_quantize(atan2, lhs, rhs, type(result)).

입력

출력

제약조건

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

예

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

 예시 더보기

batch_norm_grad

시맨틱

grad_output에서 역전파되는 여러 개의 batch_norm_training 입력의 기울기를 계산하고 grad_operand, grad_scale, grad_offset 텐서를 생성합니다. 더 엄밀히 말해 이 작업은 다음과 같이 Python 문법을 사용하여 기존 StableHLO 작업의 분해로 표현할 수 있습니다.

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

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

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

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

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

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

  return grad_operand, grad_scale, grad_offset

양자화 유형의 경우 dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, mean, variance, grad_output: batch_norm_grad(operand, scale, mean, variance, grad_output, epsilon, feature_index), operand, scale, mean, variance, grad_output, type(grad_operand), type(grad_scale), type(feature_index))를 실행합니다.

입력

출력

제약조건

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, mean, variance, grad_output, grad_operand, grad_scale 및 grad_offset의 baseline_element_type가 동일합니다.
• (C3) operand, grad_output, grad_operand의 모양이 동일합니다.
• (C4) scale, mean, variance, grad_scale 및 grad_offset의 모양이 동일합니다.
• (C5) size(scale) = dim(operand, feature_index).

예

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

batch_norm_inference

시맨틱

feature_index 차원을 제외한 모든 차원에서 operand 텐서를 정규화하고 result 텐서를 생성합니다. 더 엄밀히 말해 이 작업은 다음과 같이 Python 문법을 사용하여 기존 StableHLO 작업의 분해로 표현할 수 있습니다.

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

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

양자화 유형의 경우 dequantize_op_quantize(lambda operand, scale, offset, mean, variance: batch_norm_inference(operand, scale, offset, mean, variance, epsilon, feature_index), operand, scale, offset, mean, variance, type(result))를 실행합니다.

입력

출력

제약조건

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, mean, variance, result의 baseline_element_type가 동일합니다.
• (C3) size(scale) = dim(operand, feature_index).
• (C4) size(offset) = dim(operand, feature_index).
• (C5) size(mean) = dim(operand, feature_index).
• (C6) size(variance) = dim(operand, feature_index).
• (C7) baseline_type(operand) = baseline_type(result).

예

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

batch_norm_training

시맨틱

feature_index 측정기준을 제외한 모든 측정기준에서 평균과 분산을 계산하고 operand 텐서를 정규화하여 output, batch_mean, batch_var 텐서를 생성합니다. 더 공식적으로는 이 작업을 다음과 같이 Python 문법을 사용하여 기존 StableHLO 작업의 분해로 표현할 수 있습니다.

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

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

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

정규화된 유형의 경우 dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, offset: batch_norm_training(operand, scale, offset, epsilon, feature_index), operand, scale, offset, type(output), type(batch_mean), type(batch_var))를 실행합니다.

입력

출력

제약조건

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, batch_mean, batch_var, output는 동일한 baseline_element_type를 갖습니다.
• (C3) size(scale) = dim(operand, feature_index).
• (C4) size(offset) = dim(operand, feature_index).
• (C5) size(batch_mean) = dim(operand, feature_index).
• (C6) size(batch_var) = dim(operand, feature_index).
• (C7) baseline_type(output) = baseline_type(operand).

예

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

bitcast_convert

시맨틱

operand 텐서에 비트 전송 작업을 실행하고 전체 operand 텐서의 비트가 result 텐서의 유형을 사용하여 재해석되는 result 텐서를 생성합니다.

더 엄밀히 말해 E = element_type(operand), E' = element_type(result), R = rank(operand)가 주어지면 다음과 같습니다.

• num_bits(E') < num_bits(E)이면 bits(result[i0, ..., iR-1, :]) = bits(operand[i0, ..., iR-1]).
• num_bits(E') > num_bits(E)이면 bits(result[i0, ..., iR-2]) = bits(operand[i0, ..., iR-2, :])입니다.
• num_bits(E') = num_bits(E)이면 bits(result[i0, ..., iR-1]) = bits(operand[i0, ..., iR-1])입니다.

bits는 지정된 값의 메모리 내 표현을 반환하며, 텐서의 정확한 표현은 구현을 통해 정의되고 요소 유형의 정확한 표현도 구현에서 정의되므로 동작이 구현으로 정의됩니다.

입력

출력

제약조건

• (C1) E = is_quantized(operand) ? storage_type(operand) : element_type(operand), E' = is_quantized(result) ? storage_type(result) : element_type(result), R = rank(operand)가 주어집니다.
  • num_bits(E') = num_bits(E)이면 shape(result) = shape(operand).
  • num_bits(E') < num_bits(E)인 경우:
  • rank(result) = R + 1.
  • 모든 0 <= i < R의 dim(result, i) = dim(operand, i)
  • dim(result, R) * num_bits(E') = num_bits(E).
  • num_bits(E') > num_bits(E)인 경우:
  • rank(result) = R - 1.
  • 모든 0 <= i < R의 dim(result, i) = dim(operand, i)
  • dim(operand, R - 1) * num_bits(E) = num_bits(E').
• (C2) is_complex(operand) or is_complex(result)인 경우 is_complex(operand) and is_complex(result)입니다.

예

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

 예시 더보기

broadcast_in_dim

시맨틱

operand 텐서의 데이터를 복제하여 입력 텐서의 크기 또는 순위를 확장하고 result 텐서를 생성합니다. 더 엄밀히 말해 axes(operand)의 모든 d에 대해 result[result_index] = operand[operand_index]입니다.

• dim(operand, d) = 1인 경우 operand_index[d] = 0
• 그 밖의 경우에는 operand_index[d] = result_index[broadcast_dimensions[d]]입니다.

입력

출력

제약조건

• (C1) element_type(result)는 다음과 같이 주어집니다.
  • !is_per_axis_quantized(operand)인 경우 element_type(operand)
  • element_type(operand). 단, quantization_dimension(operand), scales(operand), zero_points(operand)는 quantization_dimension(result), scales(result), zero_points(result) 응답과 다를 수 있습니다.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) axes(operand)의 모든 d에 대해 다음을 실행합니다.
  • dim(operand, d) = 1 또는
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) is_per_axis_quantized(result)인 경우:
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • dim(operand, quantization_dimension(operand)) = 1이면 scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result)))입니다.

예

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

예시 더보기

케이스

시맨틱

index 값에 따라 branches에서 정확히 하나의 함수를 실행한 결과를 생성합니다. 더 엄밀히 말하면 result = selected_branch()이며 여기서

• 0 <= index < size(branches)인 경우 selected_branch = branches[index]입니다.
• 그 밖의 경우에는 selected_branch = branches[-1]입니다.

입력

출력

제약조건

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

예

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

 예시 더보기

cbrt

시맨틱

operand 텐서에서 요소별 세제곱근 연산을 실행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 부동 소수점의 경우: IEEE-754의 rootn(x, 3)
• 복소수의 경우: 복소 세제곱근
• 양자화 유형의 경우: dequantize_op_quantize(cbrt, operand, type(result))

입력

출력

제약조건

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

예

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

 예시 더보기

ceil

시맨틱

operand 텐서의 요소별 천정을 실행하고 result 텐서를 생성합니다. IEEE-754 사양의 roundToIntegralTowardPositive 작업을 구현합니다. 정규화된 유형의 경우 dequantize_op_quantize(ceil, operand, type(result))를 실행합니다.

입력

출력

제약조건

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

예

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

예시 더보기

콜레스키

시맨틱

일련의 행렬의 콜레스키 분해를 계산합니다.

더 엄밀히 말해 index_space(result)의 모든 i에 대해 result[i0, ..., iR-3, :, :]는 a[i0, ..., iR-3, :, :]의 콜레스키 분해이며, 이 분해는 하삼각형(lower가 true인 경우) 또는 상삼각형(lower가 false인 경우) 행렬의 형태를 취합니다. 반대 삼각형의 출력 값(즉, 엄격한 상위 삼각형 또는 엄격한 하위 삼각형)은 구현에 의해 정의됩니다.

입력 행렬이 허미티안 양의 정규 행렬이 아닌 i가 있는 경우 동작은 정의되지 않습니다.

정규화된 유형의 경우 dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result))를 실행합니다.

입력

출력

제약조건

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

예

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

고정하다

시맨틱

operand 텐서의 모든 요소를 최솟값과 최댓값 사이로 고정하고 result 텐서를 생성합니다. 더 엄밀히 말해 result[result_index] = minimum(maximum(operand[result_index], min_element), max_element)이며 여기서 min_element = rank(min) = 0 ? min[] : min[result_index], max_element = rank(max) = 0 ? max[] : max[result_index]입니다. 정규화된 유형의 경우 dequantize_op_quantize(clamp, min, operand, max, type(result))를 실행합니다.

복소수에 정렬을 적용할 때는 놀라운 의미 체계가 수반되므로 향후 이 연산에서 복소수에 대한 지원을 중단할 계획입니다. (#560)

입력

출력

제약조건

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

예

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

예시 더보기

collective_broadcast

시맨틱

StableHLO 프로세스 그리드의 각 프로세스 그룹 내에서 operand 텐서의 값을 소스 프로세스에서 타겟 프로세스로 전송하고 result 텐서를 생성합니다.

이 작업은 StableHLO 프로세스 그리드를 다음과 같이 정의된 process_groups로 분할합니다.

• channel_id <= 0인 경우 cross_replica(replica_groups)
• channel_id > 0인 경우 cross_partition(replica_groups)

그 후 result@process는 다음과 같이 주어집니다.

• 프로세스가 process_groups[i]에 있는 i가 있는 경우 operand@process_groups[i, 0]입니다.
• 그 밖의 경우에는 broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result))입니다.

입력

출력

제약조건

• (C1) is_unique(replica_groups).
• (C2) 0 <= replica_groups < N, 여기서 N는 다음과 같이 정의됩니다.
  • cross_replica가 사용된 경우 num_replicas
  • cross_partition가 사용된 경우 num_partitions
• (C3) type(result) = type(operand).

예

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

collective_permute

시맨틱

StableHLO 프로세스 그리드의 각 프로세스 그룹 내에서 소스 프로세스에서 타겟 프로세스로 operand 텐서의 값을 전송하고 result 텐서를 생성합니다.

이 작업은 StableHLO 프로세스 그리드를 다음과 같이 정의된 process_groups로 분할합니다.

• channel_id <= 0인 경우 cross_replica(source_target_pairs)
• channel_id > 0인 경우 cross_partition(source_target_pairs)입니다.

그 후 result@process는 다음과 같이 주어집니다.

• process_groups[i, 1] = process와 같은 i가 있는 경우 operand@process_groups[i, 0]입니다.
• 그 밖의 경우에는 broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result))입니다.

입력

출력

제약조건

• (C1) dim(source_target_pairs, 1) = 2.
• (C2) is_unique(source_target_pairs[:, 0]).
• (C3) is_unique(source_target_pairs[:, 1]).
• (C4) 0 <= source_target_pairs < N, 여기서 N는 다음과 같이 정의됩니다.
  • cross_replica가 사용된 경우 num_replicas
  • cross_partition가 사용된 경우 num_partitions
• (C5) type(result) = type(operand).

예

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

예시 더보기

compare

시맨틱

comparison_direction 및 compare_type에 따라 lhs 및 rhs 텐서의 요소별 비교를 실행하고 result 텐서를 생성합니다.

comparison_direction 및 compare_type 값은 다음과 같은 시맨틱을 갖습니다.

부울 및 정수 요소 유형의 경우:

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

compare_type = FLOAT가 있는 부동 소수점 요소 유형의 경우 이 작업은 다음 IEEE-754 연산을 구현합니다.

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

compare_type = TOTALORDER가 있는 부동 소수점 요소 유형의 경우 작업은 IEEE-754의 totalOrder 및 compareQuietEqual 연산 조합을 사용합니다.

복잡한 요소 유형의 경우 제공된 comparison_direction 및 compare_type를 사용하여 (real, imag) 쌍의 사전식 비교가 수행됩니다. 복소수에 정렬을 적용할 때는 놀라운 시맨틱스가 필요하기 때문에 앞으로 comparison_direction가 GE, GT, LE 또는 LT일 때 복소수에 대한 지원을 삭제할 계획입니다. (#560)

정규화된 유형의 경우 dequantize_compare(lhs, rhs, comparison_direction)를 실행합니다.

입력

출력

제약조건

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs).
• (C2) shape(lhs) = shape(rhs) = shape(result).
• (C3) compare_type는 다음과 같이 정의됩니다.
  • is_signed_integer(element_type(lhs))인 경우 SIGNED
  • is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs))인 경우 UNSIGNED
  • FLOAT 또는 TOTALORDER(is_float(element_type(lhs))인 경우)
  • is_complex(element_type(lhs))인 경우 FLOAT입니다.

예

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

 예시 더보기

복잡함

시맨틱

실수와 허수 값 쌍인 lhs 및 rhs에서 복소수 값으로 요소별 변환을 실행하고 result 텐서를 생성합니다.

입력

출력

제약조건

• (C1) type(lhs) = type(rhs).
• (C2) shape(result) = shape(lhs).
• (C3) element_type(result)의 유형은 complex<E>입니다. 여기서 E = element_type(lhs)입니다.

예

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

 예시 더보기

복합

시맨틱

다른 StableHLO 작업으로 구성된 작업을 캡슐화하여 inputs 및 composite_attributes를 사용하고 results를 생성합니다. 작업의 의미는 decomposition 속성으로 구현됩니다. composite 연산자는 프로그램 시맨틱을 변경하지 않고도 분해로 대체할 수 있습니다. 분해를 인라인 처리해도 동일한 연산 시맨틱을 제공하지 않는 경우 custom_call를 사용하는 것이 좋습니다.

version 필드(기본값: 0)는 합성의 시맨틱이 변경될 때를 나타내는 데 사용됩니다.

입력

출력

제약조건

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

예

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

 예시 더보기

이어붙이기

시맨틱

지정된 인수와 동일한 순서로 dimension 차원을 따라 inputs를 연결하고 result 텐서를 생성합니다. 더 공식적으로 result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1]이며, 각 항목의 의미는 다음과 같습니다.

1. id = d0 + ... + dk-1 + kd.
2. d은 dimension와 같고 d0, ...은 inputs의 d번째 크기입니다.

입력

출력

제약조건

• (C1) same(element_type(inputs...)).
• (C2) dim(inputs..., dimension)를 제외한 same(shape(inputs...))
• (C3) 0 < size(inputs).
• (C4) 0 <= dimension < rank(inputs[0]).
• (C5) element_type(result) = element_type(inputs[0]).
• (C6) 다음을 제외한 shape(result) = shape(inputs[0]):
  • dim(result, dimension) = dim(inputs[0], dimension) + ....

예

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

 예시 더보기

상수

시맨틱

상수 value에서 output 텐서를 생성합니다.

입력

출력

제약조건

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

예

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

 예시 더보기

전환

시맨틱

operand 텐서에서 한 요소 유형에서 다른 요소 유형으로 요소별 변환을 실행하고 result 텐서를 생성합니다.

boolean-to-any-supported-type 변환의 경우 값 false은 0으로 변환되고 값 true은 1로 변환됩니다. any-supported-type-to-boolean 변환의 경우 0 값은 false로 변환되고 0이 아닌 값은 true로 변환됩니다. 복잡한 유형에서 이 작업이 어떻게 작동하는지 아래를 참고하세요.

integer-to-integer, integer-to-floating-point 또는 floating-point-to-floating-point 관련 변환의 경우 소스 값이 대상 유형에 정확하게 표시될 수 있는 경우 결과 값은 해당 표현입니다. 그 외의 경우에는 동작은 미정입니다. (#180)

floating-point-to-integer 변환의 경우 소수 부분이 잘립니다. 잘린 값을 대상 유형으로 표현할 수 없는 경우 동작은 미정입니다(#180).

복소수-복소수 변환은 실수와 허수 부분을 변환하는 부동 소수점-부동 소수점 변환과 동일한 동작을 따릅니다.

complex-to-any-other-type 및 any-other-type-to-complex 변환의 경우 각각 소스 허수 값이 무시되거나 대상 허수 값이 0으로 설정됩니다. 실수 부분의 변환은 부동 소수점 변환을 따릅니다.

원칙적으로 이 연산은 역양자화 (양자화 텐서에서 일반 텐서로 변환), 양자화 (일반 텐서에서 양자화 텐서로 변환) 및 재양자화 (양자화 텐서 간 변환)를 표현할 수 있지만, 현재 첫 번째 사용 사례에서는 uniform_dequantize, 두 번째 및 세 번째 사용 사례에서는 uniform_quantize를 위한 전용 연산이 있습니다. 향후 이 두 연산은 convert로 병합될 수 있습니다(#1576).

입력

출력

제약조건

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

예

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

 예시 더보기

컨볼루션

시맨틱

lhs의 창과 rhs의 슬라이스 간의 내적을 계산하고 result를 생성합니다. 다음 다이어그램은 구체적인 예를 사용하여 result의 요소가 lhs 및 rhs에서 계산되는 방식을 보여줍니다.

좀 더 공식적으로 lhs의 기간을 표현할 수 있도록 lhs 관점에서 입력을 다음과 같이 재구성해 보세요.

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

이 프레임 조정에는 다음과 같은 도우미 함수가 사용됩니다.

• lhs_shape(n, hw, c) = permute([n] + hw + [c], [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]).
• result_shape(n1, hw, c1) = permute([n1] + hw + [c1], [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]).
• j[d] = i[permutation[d]]인 경우 permute([j0, j1, ..., jR-1], permutation) = [i0, i1, ..., iR-1]입니다.

feature_group_count = 1 및 batch_group_count = 1인 경우 index_space(dim(result, output_spatial_dimensions...))의 모든 output_spatial_index에 대해 다음과 같이 result[result_shape(:, output_spatial_index, :)] = dot_product:

• padding_value = constant(0, element_type(lhs)).
• padded_lhs = pad(lhs, padding_value, lhs_padding[:, 0], lhs_padding[:, 1], lhs_base_dilations - 1).
• lhs_window_start = lhs_shape(0, output_spatial_index, 0) * lhs_window_strides.
• lhs_window = slice(padded_lhs, lhs_window_start, lhs_window_start + lhs_window_dimensions, lhs_window_dilations).
• reversed_lhs_window = reverse(lhs_window, [input_spatial_dimensions[dim] for dim in range(size(window_reversal)) if window_reversal[dim] = true]). 이 기능은 사용되지 않는 것으로 보여 향후 삭제될 예정입니다. (#1181)
• dot_product = dot_general(reversed_lhs_window, rhs, lhs_batching_dimensions=[], lhs_contracting_dimensions=input_spatial_dimensions + [input_feature_dimension], rhs_batching_dimensions=[], rhs_contracting_dimensions=kernel_spatial_dimensions + [kernel_input_feature_dimension]).

feature_group_count > 1인 경우:

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

batch_group_count > 1인 경우:

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

정규화된 유형의 경우 dequantize_op_quantize( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs, type(result))를 실행합니다.

하이브리드 양자화된 유형의 경우 hybrid_dequantize_then_op( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs)를 실행합니다.

입력

출력

제약조건

• (C1) N = rank(lhs) = rank(rhs).
• (C2) size(window_strides) = N - 2.
• (C3) 0 < window_strides.
• (C4) shape(padding) = [N - 2, 2].
• (C5) size(lhs_dilation) = N - 2.
• (C6) 0 < lhs_dilation.
• (C7) size(rhs_dilation) = N - 2.
• (C8) 0 < rhs_dilation.
• (C9) size(window_reversal) = N - 2.
• (C10) dim(lhs, input_batch_dimension) % batch_group_count = 0.
• (C11) dim(lhs, input_feature_dimension) % feature_group_count = 0입니다.
• (C12) size(input_spatial_dimensions) = N - 2.
• (C13) input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]가 주어집니다.
  • is_unique(input_dimensions).
  • 0 <= input_dimensions < N.
• (C14) dim(rhs, kernel_input_feature_dimension) = dim(lhs, input_feature_dimension) / feature_group_count입니다.
• (C15) dim(rhs, kernel_output_feature_dimension) % batch_group_count = 0.
• (C16) dim(rhs, kernel_output_feature_dimension) % feature_group_count = 0.
• (C17) size(kernel_spatial_dimensions) = N - 2.
• (C18) kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]가 주어집니다.
  • is_unique(kernel_dimensions).
  • 0 <= kernel_dimensions < N.
• (C19) size(output_spatial_dimensions) = N - 2.
• (C20) output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]가 주어집니다.
  • is_unique(output_dimensions).
  • 0 <= output_dimensions < N.
• (C21) 0 < feature_group_count입니다.
• (C22) 0 < batch_group_count.
• (C23) feature_group_count = 1 or batch_group_count = 1.
• (C24) size(precision_config) = 2.
• (C25) dim(result, result_dim)는 다음과 같이 정의됩니다.
  • result_dim = output_batch_dimension인 경우 dim(lhs, input_batch_dimension) / batch_group_count입니다.
  • result_dim = output_feature_dimension인 경우 dim(rhs, kernel_output_feature_dimension)
  • 그렇지 않으면 num_windows입니다. 여기서:
  • output_spatial_dimensions[spatial_dim] = result_dim.
  • lhs_dim = input_spatial_dimensions[spatial_dim].
  • rhs_dim = kernel_spatial_dimensions[spatial_dim].
  • dilated_input_shape[lhs_dim] = dim(lhs, lhs_dim) = 0 ? 0 : (dim(lhs, lhs_dim) - 1) * lhs_dilation[spatial_dim] + 1.
  • padded_input_shape[lhs_dim] = padding[spatial_dim, 0] + dilated_input_shape[lhs_dim] + padding[spatial_dim, 1].
  • dilated_window_shape[lhs_dim] = dim(rhs, rhs_dim) = 0 ? 0 : (dim(rhs, rhs_dim) - 1) * rhs_dilation[spatial_dim] + 1.
  • is_empty_window[lhs_dim] = padded_input_shape[lhs_dim] = 0 || dilated_window_shape[lhs_dim] > padded_input_shape[lhs_dim].
  • num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1.
• (C26) rank(result) = N.
• 작업에서 정규화되지 않은 텐서를 사용하는 경우:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• 연산이 양자화 텐서를 사용하는 경우:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) is_per_axis_quantized(rhs)이면 quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) is_per_axis_quantized(result)이면 quantization_dimension(result) = output_feature_dimension입니다.
  • is_quantized(lhs)인 경우:
  • (C31) storage_type(lhs) = storage_type(rhs)입니다.
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)입니다.
  • (C33) is_per_tensor_quantized(rhs)이면 is_per_tensor_quantized(result)입니다.
  • !is_quantized(lhs)인 경우:
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

예

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

 예시 더보기

코사인

시맨틱

operand 텐서에서 요소별로 코사인 연산을 실행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음 작업을 실행합니다.

• 부동 소수점의 경우: IEEE-754의 cos
• 복소수의 경우: 복소 코사인
• 정규화된 유형: dequantize_op_quantize(cosine, operand, type(result)).

입력

출력

제약조건

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

예

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

예시 더보기

count_leading_zeros

시맨틱

operand 텐서의 선행 0 비트 수를 요소별로 계산하고 result 텐서를 생성합니다.

입력

출력

제약조건

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

예

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

 예시 더보기

custom_call

시맨틱

inputs 및 called_computations을 사용하여 results을 생성하는 구현 정의 작업 call_target_name를 캡슐화합니다. has_side_effect, backend_config, api_version는 구현 정의 메타데이터를 추가로 제공하는 데 사용할 수 있습니다.

현재 이 작업에는 XLA 컴파일러의 상응하는 작업의 유기적 진화를 반영하는 상당히 비조직적인 메타데이터 모음이 포함되어 있습니다. 향후 이 메타데이터를 통합할 계획입니다. (#741)

입력

출력

예

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

나누기

시맨틱

배당금 lhs 텐서와 제수 rhs 텐서의 요소별 나눗셈을 실행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음 작업을 실행합니다.

• 정수의 경우: 소수 부분이 삭제된 대수적 몫을 생성하는 정수 나눗셈입니다.
• 부동 소수점의 경우: IEEE-754의 division
• 복소수의 경우: 복소수 나눗셈
• 정규화된 유형의 경우:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)).

입력

출력

제약조건

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

예

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

예시 더보기

dot_general

시맨틱

lhs 슬라이스와 rhs 슬라이스 간의 내적을 계산하고 result 텐서를 생성합니다.

더 엄밀히 말하면 result[result_index] = dot_product이며, 여기서

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

정규화된 유형의 경우 dequantize_op_quantize( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs, type(result))를 실행합니다.

하이브리드 양자화된 유형의 경우 hybrid_dequantize_then_op( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs)를 실행합니다.

precision_config는 가속기 백엔드에서 계산의 속도와 정확성 간의 균형을 제어합니다. 다음 중 하나일 수 있습니다. 현재 이러한 enum 값의 시맨틱은 미정이지만 #755에서 이 문제를 해결할 계획입니다.

• DEFAULT: 가장 빠른 계산이지만 원래 숫자에 가장 근접하지 않습니다.
• HIGH: 계산 속도는 느리지만 원래 숫자에 더 가깝습니다.
• HIGHEST: 가장 느린 계산이지만 원래 숫자에 가장 근접한 근사치입니다.

DotAlgorithm는 도트 연산을 구현하는 데 사용되는 알고리즘의 기본 속성을 정의하며, 이는 정밀도도 정의합니다. 알고리즘 속성 필드가 설정된 경우 precision_config는 DEFAULT이어야 합니다. DotAlgorithms에는 기본 매개변수가 구현되어 정의되므로 기본값이 없습니다. 따라서 모든 점 알고리즘 필드를 None로 설정하여 빈 점 알고리즘을 지정할 수 있으며, 이 경우 대신 precision_config 값이 사용됩니다.

DotAlgorithm 필드에는 다음이 포함됩니다.

• lhs_precision_type 및 rhs_precision_type: 작업의 좌항과 우항이 반올림되는 정밀도입니다. 정밀도 유형은 입력과 출력의 저장소 유형과 무관합니다.
• accumulation_type 누적에 사용되는 정밀도입니다.
• lhs_component_count, rhs_component_count, num_primitive_operations는 LHS 또는 RHS를 여러 구성요소로 분해하고 이러한 값에 여러 '원시' 내적 연산을 실행하는 알고리즘을 실행할 때 적용됩니다. 일반적으로 더 높은 정밀도를 에뮬레이션하기 위함입니다(예: 더 높은 정밀도 계산을 위해 bfloat16 인공지능 데이터 유형 활용: bf16_6x tf32_3x 등). 분해가 없는 알고리즘의 경우 이러한 값을 1로 설정해야 합니다.
• allow_imprecise_accumulation: 일부 단계에서 더 낮은 정밀도의 누적이 허용되는지 지정합니다 (예: CUBLASLT_MATMUL_DESC_FAST_ACCUM).

DotAlgorithm 속성 예시:

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


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


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

어떤 조합이 지원되는지는 구현에서 결정합니다. 일반적으로 각 알고리즘이 StableHLO의 소비자에 의해 각 가속기 유형에서 지원되는 것은 아닙니다. 특정 알고리즘이 지원되지 않는 경우 대체 알고리즘으로 대체하는 대신 오류가 발생해야 합니다. StableHLO 확인은 최선의 확인을 제공하여 모든 하드웨어에서 지원되지 않는 것으로 알려진 알고리즘을 방지합니다.

지원되는 일부 알고리즘 값은 xla_data.proto > Algorithm을 참고하세요. 티켓 #2483은 백엔드에서 지원하는 알고리즘에 대해 중앙 집중식 문서를 생성하려는 계획을 나타냅니다.

입력

출력

제약조건

• (C1) size(lhs_batching_dimensions) = size(rhs_batching_dimensions).
• (C2) size(lhs_contracting_dimensions) = size(rhs_contracting_dimensions).
• (C3) is_unique(lhs_batching_dimensions + lhs_contracting_dimensions).
• (C4) is_unique(rhs_batching_dimensions + rhs_contracting_dimensions).
• (C5) 0 <= lhs_batching_dimensions < rank(lhs).
• (C6) 0 <= lhs_contracting_dimensions < rank(lhs).
• (C7) 0 <= rhs_batching_dimensions < rank(rhs).
• (C8) 0 <= rhs_contracting_dimensions < rank(rhs).
• (C9) dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...).
• (C10) dim(lhs, lhs_contracting_dimensions...) = dim(rhs, rhs_contracting_dimensions...).
• (C11) size(precision_config) = 2입니다.
• (C12) shape(result) = dim(lhs, lhs_batching_dimensions) + dim(lhs, lhs_result_dimensions) + dim(rhs, rhs_result_dimensions).
• 작업이 비 양자화 텐서를 사용하는 경우:
  • (C13) element_type(lhs) = element_type(rhs).
• 작업에서 정규화된 텐서를 사용하는 경우:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C15) zero_points(rhs) = 0.
  • (C16) is_per_axis_quantized(rhs)이면 quantization_dimension(rhs)가 rhs_contracting_dimensions에 없습니다.
  • is_quantized(lhs)인 경우:
  • (C17) storage_type(lhs) = storage_type(rhs)).
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C19) is_per_tensor_quantized(rhs)이면 is_per_tensor_quantized(result)입니다.
  • !is_quantized(lhs)인 경우:
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result)입니다.
• !is_empty_algorithm(lhs_precision_type, rhs_precision_type, accumulation_type, lhs_component_count, rhs_component_count, num_primitive_operations allow_imprecise_accumulation)인 경우:
  • (C21) precision_config... = DEFAULT입니다.
  • (C22) 0 < lhs_component_count입니다.
  • (C23) 0 < rhs_component_count입니다.
  • (C24) 0 < num_primitive_operations.

예

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

 예시 더보기

dynamic_broadcast_in_dim

시맨틱

이 연산은 기능적으로 broadcast_in_dim 연산과 동일하지만 결과 셰이프는 output_dimensions를 통해 동적으로 지정됩니다.

이 연산은 선택적 속성 known_expanding_dimensions, known_nonexpanding_dimensions도 허용하여 크기 조절 동작에 관한 정적 지식을 표현합니다. 지정하지 않으면 모든 측정기준이 확장될 수 있다고 가정됩니다.

입력

출력

제약조건

• (C1) element_type(result)는 다음과 같이 주어집니다.
  • !is_per_axis_quantized(operand)인 경우 element_type(operand)
  • element_type(operand). 단, quantization_dimension(operand), scales(operand), zero_points(operand)는 quantization_dimension(result), scales(result), zero_points(result) 응답과 다를 수 있습니다.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) axes(operand)의 모든 d에 대해 다음을 실행합니다.
  • dim(operand, d) = 1 또는
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) is_per_axis_quantized(result)인 경우:
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • dim(operand, quantization_dimension(operand)) = 1이면 scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result)))입니다.
• (C7) size(output_dimensions) = rank(result).
• (C8) is_unique(known_expanding_dimensions + known_nonexpanding_dimensions).
• (C9) 0 <= known_expanding_dimensions < rank(operand).
• (C10) 0 <= known_nonexpanding_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_nonexpanding_dimensions = array<i64: 1>
} : (tensor<1x3xi64>, tensor<3xi64>) -> tensor<2x3x2xi64>
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

 예시 더보기

dynamic_conv

시맨틱

이 연산은 컨볼루션 작업과 기능적으로 동일하지만 패딩은 padding를 통해 동적으로 지정됩니다.

입력

출력

제약조건

• (C1) N = rank(lhs) = rank(rhs).
• (C2) size(window_strides) = N - 2.
• (C3) 0 < window_strides.
• (C4) shape(padding) = [N - 2, 2].
• (C5) size(lhs_dilation) = N - 2.
• (C6) 0 < lhs_dilation.
• (C7) size(rhs_dilation) = N - 2.
• (C8) 0 < rhs_dilation.
• (C9) size(window_reversal) = N - 2.
• (C10) dim(lhs, input_batch_dimension) % batch_group_count = 0.
• (C11) dim(lhs, input_feature_dimension) % feature_group_count = 0입니다.
• (C12) size(input_spatial_dimensions) = N - 2.
• (C13) input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]가 주어집니다.
  • is_unique(input_dimensions).
  • 0 <= input_dimensions < N.
• (C14) dim(rhs, kernel_input_feature_dimension) = dim(lhs, input_feature_dimension) / feature_group_count입니다.
• (C15) dim(rhs, kernel_output_feature_dimension) % batch_group_count = 0.
• (C16) dim(rhs, kernel_output_feature_dimension) % feature_group_count = 0.
• (C17) size(kernel_spatial_dimensions) = N - 2.
• (C18) kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]가 주어집니다.
  • is_unique(kernel_dimensions).
  • 0 <= kernel_dimensions < N.
• (C19) size(output_spatial_dimensions) = N - 2.
• (C20) output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]가 주어집니다.
  • is_unique(output_dimensions).
  • 0 <= output_dimensions < N.
• (C21) 0 < feature_group_count입니다.
• (C22) 0 < batch_group_count.
• (C23) feature_group_count = 1 or batch_group_count = 1.
• (C24) size(precision_config) = 2.
• (C25) dim(result, result_dim)는 다음과 같이 정의됩니다.
  • result_dim = output_batch_dimension인 경우 dim(lhs, input_batch_dimension) / batch_group_count입니다.
  • result_dim = output_feature_dimension인 경우 dim(rhs, kernel_output_feature_dimension)
  • 그렇지 않으면 num_windows입니다. 여기서:
  • output_spatial_dimensions[spatial_dim] = result_dim.
  • lhs_dim = input_spatial_dimensions[spatial_dim].
  • rhs_dim = kernel_spatial_dimensions[spatial_dim].
  • dilated_input_shape[lhs_dim] = dim(lhs, lhs_dim) = 0 ? 0 : (dim(lhs, lhs_dim) - 1) * lhs_dilation[spatial_dim] + 1.
  • padded_input_shape[lhs_dim] = padding[spatial_dim, 0] + dilated_input_shape[lhs_dim] + padding[spatial_dim, 1].
  • dilated_window_shape[lhs_dim] = dim(rhs, rhs_dim) = 0 ? 0 : (dim(rhs, rhs_dim) - 1) * rhs_dilation[spatial_dim] + 1.
  • is_empty_window[lhs_dim] = padded_input_shape[lhs_dim] = 0 || dilated_window_shape[lhs_dim] > padded_input_shape[lhs_dim].
  • num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1.
• (C26) rank(result) = N.
• 작업에서 정규화되지 않은 텐서를 사용하는 경우:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• 연산이 양자화 텐서를 사용하는 경우:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) is_per_axis_quantized(rhs)이면 quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) is_per_axis_quantized(result)이면 quantization_dimension(result) = output_feature_dimension입니다.
  • is_quantized(lhs)인 경우:
  • (C31) storage_type(lhs) = storage_type(rhs)입니다.
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)입니다.
  • (C33) is_per_tensor_quantized(rhs)이면 is_per_tensor_quantized(result)입니다.
  • !is_quantized(lhs)인 경우:
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

예

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