StableHLO 사양

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

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

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

프로그램

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 함수 (이름이 지정된 함수라고도 함)에는 식별자, 입력/출력, 본문이 있습니다. Google은 향후 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개의 차원 (즉, 2개의 축)이 있고 크기는 2와 3입니다. 2위입니다.

이는 크기 크기가 정적으로 알려진 정적 도형 지원을 정의합니다. 앞으로는 크기 크기를 부분적으로 또는 완전히 알 수 없는 동적 도형도 지원할 계획입니다(#8). 또한 레이아웃(#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-C4), (C9)
storage_min	정수 상수	(C2), (C4), (C8)
storage_max	정수 상수	(C3), (C4), (C8)
expressed_type	부동 소수점 유형	(C1), (C5)
quantization_dimension	선택적 정수 상수	(C11-C13)
scales	부동 소수점 상수의 가변 수	(C5-C7), (C10), (C11), (C13)
zero_points	정수 상수의 가변 수	(C8-C10)

양자화된 요소 유형은 표현된 유형의 부동 소수점 값에 해당하는 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) num_bits(storage_type) < num_bits(expressed_type).
• (C2) type(storage_min) = storage_type.
• (C3) type(storage_max) = storage_type.
• (C4) min_value(storage_type) <= storage_min < storage_max <= max_value(storage_type).
• (C5) type(scales...) = expressed_type.
• (C6) 0 < scales.
• (C7) is_finite(scales...).
• (C8) storage_min <= zero_points <= storage_max.
• (C9) type(zero_points...) = storage_type.
• (C10) size(scales) = size(zero_points).
• (C11) is_empty(quantization_dimension)인 경우 size(scales) = 1입니다.
• (C12) 0 <= quantization_dimension.

현재 QuantizationScale는 부동 소수점 상수이지만 승수와 이동으로 표현되는 정수 기반 배율이 많은 관심을 보이고 있습니다. 조만간 이 기능을 살펴볼 계획입니다. (#1404)

유형, 값, 양자화 텐서 유형에 0점만 1개 또는 잠재적으로 여러 개일 수 있는지 등 QuantizationZeroPoint의 시맨틱스에 관한 논의가 진행 중입니다. 이 설명의 결과에 따라 0포인트에 대한 사양은 향후 변경될 수 있습니다. (#1405)

이러한 값과 양자화 텐서 값에 제약 조건을 적용해야 하는지 결정하는 QuantizationStorageMin 및 QuantizationStorageMax의 시맨틱스에 관한 또 다른 논의가 진행 중입니다. (#1406)

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

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

양자화 텐서에서 양자화는 텐서별일 수 있습니다. 즉, 전체 텐서에 하나의 scale 및 zero_point를 보유하거나 축당으로 사용할 수 있습니다. 즉, 특정 차원 quantization_dimension의 슬라이스당 한 쌍의 scales와 zero_points를 보유할 수 있습니다. 좀 더 공식적으로 축당 양자화가 포함된 텐서 t에는 quantization_dimension(t[:, ..., 0, ..., :], t[:, ..., 1, ..., :] 등)의 dim(t, quantization_dimension) 슬라이스가 있습니다. 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 ::= '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 유형의 스칼라 값을 tensor<T> 유형의 0차원 텐서 값으로 표현하는 것이 관용적입니다.

• 부울 유형은 불리언 값 true 및 false를 나타냅니다.
• 정수 유형은 부호 있는 유형 (si) 또는 부호 없는 유형 (ui)일 수 있으며 지원되는 비트 너비 (4, 8, 16, 32 또는 64) 중 하나를 가집니다. 부호 있는 siN 유형은 -2^(N-1) 이상 2^(N-1)-1까지의 정수 값을 나타내고, 부호 없는 uiN 유형은 0부터 2^N-1까지의 정수 값을 나타냅니다.
• 부동 소수점 유형은 다음 중 하나입니다.
  • 딥 러닝을 위한 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 유형입니다.
• 복소수 유형은 동일한 요소 유형의 실수와 허수를 갖는 복합 값을 나타냅니다. 지원되는 복합 유형은 complex<f32> (두 부분 모두 f32 유형임) 및 complex<f64>(두 부분 모두 f64 유형임)입니다.

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

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

StringType ::= 'string'

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

운영

안정적인 HLO 작업 (작업이라고도 함)은 머신러닝 모델에서 상위 수준의 폐쇄형 작업 집합을 나타냅니다. 위에서 설명한 것처럼 StableHLO 구문은 MLIR에서 많은 영감을 받았습니다. 가장 인체공학적인 대안은 아니지만, ML 프레임워크와 ML 컴파일러 간의 상호 운용성을 개선한다는 StableHLO의 목표에 가장 적합합니다.

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

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

현재 야생 상태의 StableHLO 프로그램에 이 문서에서 설명되지 않는 작업이 포함되는 경우가 있습니다. 향후 이러한 작업을 StableHLO 명령에 포함시키거나 StableHLO 프로그램에 표시되지 않도록 할 계획입니다. 그동안 이러한 작업의 목록은 다음과 같습니다.

• builtin.module, func.func, func.call, func.return(#425).
• chlo 연산 (#602)
• StableHLO 작업의 'Not in HLO' 카테고리 - 처음에는 StableHLO 명령의 일부였지만 나중에는 적합하지 않은 것으로 간주되었습니다. broadcast, create_token, cross-replica-sum, dot, einsum, torch_index_select, unary_einsum(#3).
• StableHLO 작업의 'Dynamism' 카테고리 - MHLO에서 부트스트랩되었지만 아직 지정하지 않았습니다. compute_reshape_shape, cstr_reshapable, dynamic_broadcast_in_dim, dynamic_conv, dynamic_gather, dynamic_iota, dynamic_pad, dynamic_reshape, real_dynamic_slice, set_dimension_size(#8)
• arith, shape, tensor 작업을 포함한 셰이프 계산(#8)

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) 식별자가 없으므로 (따라서 '익명'이라는 이름) 출력 유형을 선언하지 않습니다 (출력 유형은 함수 내의 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 명령에 흡수하거나 StableHLO 프로그램에 표시되지 않도록 할 계획입니다. 그 동안 이러한 속성 목록은 다음과 같습니다.

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

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

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

다음은 니모닉이 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'

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

• (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 텐서에서 요소별 abs 연산을 실행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 부호 있는 정수: 정수 모듈러스.
• 부동 소수점 수: IEEE-754의 abs
• 복소수: 복소수입니다.
• 양자화 유형: dequantize_op_quantize(abs, operand, type(result))

입력

라벨	이름	유형	제약 조건
(I1)	operand	부호 있는 정수, 부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1-C2)

출력

이름	유형	제약 조건
result	부호 있는 정수, 부동 소수점 유형 또는 텐서당 양자화 텐서의 텐서	(C1-C2)

제약 조건

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

예시 더보기

add

시맨틱

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

• 불리언: 논리 OR입니다.
• 정수: 정수 덧셈.
• 부동 소수점 수: IEEE-754의 addition
• 복소수: 복소수 덧셈.
• 양자화 유형: dequantize_op_quantize(add, lhs, rhs, type(result))

입력

라벨	이름	유형	제약 조건
(I1)	lhs	텐서 또는 텐서당 양자화 텐서	(C1)
(I2)	rhs	텐서 또는 텐서당 양자화 텐서	(C1)

출력

이름	유형	제약 조건
result	텐서 또는 텐서당 양자화 텐서	(C1)

제약 조건

• (C1) baseline_type(lhs) = baseline_type(rhs) = baseline_type(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까지의 데이터 종속 항목을 설정하기 위해서만 존재합니다.

입력

라벨	이름	유형
(I1)	inputs	token의 변수수

출력

이름	유형
result	token

예

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

예시 더보기

all_gather

시맨틱

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

이 작업은 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에 대한 result@process = concatenate(operands@process, all_gather_dim)입니다.

입력

라벨	이름	유형	제약 조건
(I1)	operand	텐서 또는 텐서당 양자화 텐서	(C1), (C6)
(I2)	all_gather_dim	si64 유형의 상수	(C1), (C6)
(I3)	replica_groups	si64 유형의 2차원 텐서 상수	(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).
• (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(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로 분할합니다.

• 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 내에서 다음을 실행합니다.

• 일부 바이너리 트리 schedule의 경우 result@process[result_index] = exec(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 유형 1차원 텐서 상수의 가변 숫자	(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).
• (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에는 is_promotable(element_type(operand), E)인 (tensor<E>, tensor<E>) -> (tensor<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<i64>) -> tensor<i64>
// %result@(0, 0): [6, 8, 10, 12]
// %result@(1, 0): [6, 8, 10, 12]

예시 더보기

all_to_all

시맨틱

StableHLO 프로세스 그리드의 각 프로세스 그룹 내에서 operand 텐서 값을 split_dimension에 따라 여러 부분으로 분할하고, 프로세스 간에 분할 부분을 분산하고, concat_dimension를 따라 분산된 부분을 연결하고, result 텐서를 생성합니다.

이 작업은 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(operand@sender, split_count, split_dimension).
• 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 유형의 2차원 텐서 상수	(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).
• (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) type(result) = type(operand) 제외:
  • dim(result, split_dimension) = dim(operand, split_dimension) / split_count.
  • dim(result, concat_dimension) = dim(operand, concat_dimension) * split_count.

예

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

예시 더보기

및

시맨틱

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

• 불리언: 논리 AND
• 정수: 비트 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]]

atan2

시맨틱

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

• 부동 소수점 수: IEEE-754의 atan2
• 복소수: 복소수 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

시맨틱

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))를 실행합니다.

입력

라벨	이름	유형	제약 조건
(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]

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))를 실행합니다.

입력

라벨	이름	유형	제약 조건
(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]]
//          ]

batch_norm_training

시맨틱

feature_index 차원을 제외한 모든 차원의 평균과 분산을 계산하고 output, batch_mean, batch_var 텐서를 생성하는 operand 텐서를 정규화합니다. 보다 공식적으로 이 작업은 다음과 같이 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))를 실행합니다.

입력

라벨	이름	유형	제약 조건
(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.
  • 모든 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 텐서를 생성합니다. 좀 더 공식적으로 result[result_index] = operand[operand_index], 여기서 axes(operand)의 모든 d에 대해 다음을 실행합니다.

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

입력

라벨	이름	유형	제약 조건
(I1)	operand	텐서 또는 양자화 텐서	(C1-C2), (C5-C6)
(I2)	broadcast_dimensions	si64 유형의 1차원 텐서 상수	(C2-C6)

출력

이름	유형	제약 조건
result	텐서 또는 양자화 텐서	(C1), (C3), (C5-C6)

제약 조건

• (C1) element_type(result)는 다음과 같이 주어집니다.
  • !is_per_axis_quantized(operand)인 경우 element_type(operand)입니다.
  • quantization_dimension(operand), scales(operand), zero_points(operand)가 quantization_dimension(result), scales(result), zero_points(result) 응답과 다를 수 있다는 점을 제외하고 element_type(operand)
• (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]입니다.

입력

라벨	이름	유형	제약 조건
(I1)	index	si32 유형의 0차원 텐서
(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 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 부동 소수점 수: IEEE-754의 rootn(x, 3)
• 복소수: 복소수 세제곱근
• 양자화 유형: 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]

예시 더보기

ceil

시맨틱

operand 텐서의 요소별 ceil을 수행하고 result 텐서를 생성합니다. IEEE-754 사양의 roundToIntegralTowardPositive 작업을 구현합니다. 양자화 유형의 경우 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]

예시 더보기

콜레스키

시맨틱

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

좀 더 공식적으로 index_space(result)의 모든 i에 대해 result[i0, ..., iR-3, :, :]는 하부삼각형(lower가 true인 경우) 또는 상삼각형 (lower가 false인 경우) 행렬의 형태로 a[i0, ..., iR-3, :, :]의 콜레스키 분해입니다. 반대쪽 삼각형의 출력 값, 즉 엄격한 위쪽 삼각형 또는 엄격한 아래쪽 삼각형은 구현을 통해 정의됩니다.

입력 행렬이 헤르미티아 양의 정적 행렬이 아닌 i가 있으면 동작이 정의되지 않습니다.

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

입력

라벨	이름	유형	제약 조건
(I1)	a	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1-C3)
(I2)	lower	i1 유형의 0차원 텐서 상수

출력

이름	유형	제약 조건
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]

예시 더보기

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(0, element_type(result)), [], type(result)) 하지 않으면 사용할 수 없습니다.

입력

라벨	이름	유형	제약 조건
(I1)	operand	텐서	(C3)
(I2)	replica_groups	si64 유형 1차원 텐서 상수의 가변 숫자	(C1), (C2)
(I3)	channel_id	si64 유형의 상수

출력

이름	유형	제약 조건
result	텐서	(C3)

제약 조건

• (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은 다음과 같이 제공됩니다.

• operand@process_groups[i, 0]: process_groups[i, 1] = process와 같은 i가 있는 경우
• 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	si64 유형의 2차원 텐서 상수	(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는 다음과 같이 정의됩니다.
  • 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 연산의 조합을 사용합니다. 이 기능은 사용되지 않는 것으로 보이므로 향후 삭제될 예정입니다. (#584)

복잡한 요소 유형의 경우 제공된 comparison_direction 및 compare_type를 사용하여 (real, imag) 쌍의 사전식 비교가 실행됩니다. 복소수에 정렬을 적용하는 데는 놀라운 의미 체계가 포함되므로 향후 comparison_direction이 GE, GT, LE 또는 LT일 때 복소수에 관한 지원을 중단할 계획입니다. (#560)

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

입력

라벨	이름	유형	제약 조건
(I1)	lhs	텐서 또는 텐서당 양자화 텐서	(C1-C3)
(I2)	rhs	텐서 또는 텐서당 양자화 텐서	(C1-C2)
(I3)	comparison_direction	EQ, NE, GE, GT, LE, LT의 열거형
(I4)	compare_type	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입니다.
  • 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 텐서를 생성합니다.

입력

라벨	이름	유형	제약 조건
(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)에는 E = element_type(lhs)인 complex<E> 유형이 있습니다.

예

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

예시 더보기

concatenate

시맨틱

지정된 인수와 동일한 순서로 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번째 차원 크기입니다.

입력

라벨	이름	유형	제약 조건
(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]]

예시 더보기

상수

시맨틱

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

입력

라벨	이름	유형	제약 조건
(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 텐서를 생성합니다.

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)

입력

라벨	이름	유형	제약 조건
(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)]

예시 더보기

컨볼루션

시맨틱

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

보다 공식적으로 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인 경우 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))를 실행합니다.

입력

라벨	이름	유형	제약 조건
(I1)	lhs	텐서 또는 텐서당 양자화 텐서	(C1), (C10-C11), (C14) (C25), (C27-C30)
(I2)	rhs	텐서 또는 양자화 텐서	(C1), (C14-C16), (C25), (C27-C32)
(I3)	window_strides	si64 유형의 1차원 텐서 상수	(C2-C3), (C25)
(I4)	padding	si64 유형의 2차원 텐서 상수	(C4), (C25)
(I5)	lhs_dilation	si64 유형의 1차원 텐서 상수	(C5-C6), (C25)
(I6)	rhs_dilation	si64 유형의 1차원 텐서 상수	(C7-C8), (C25)
(I7)	window_reversal	i1 유형의 1차원 텐서 상수	(C9)
(I8)	input_batch_dimension	si64 유형의 상수	(C10), (C13), (C25)
(I9)	input_feature_dimension	si64 유형의 상수	(C11), (C13-C14)
(I10)	input_spatial_dimensions	si64 유형의 1차원 텐서 상수	(C12), (C13), (C25)
(I11)	kernel_input_feature_dimension	si64 유형의 상수	(C14), (C18)
(I12)	kernel_output_feature_dimension	si64 유형의 상수	(C15-C16), (C18), (C25), (C32)
(I13)	kernel_spatial_dimensions	si64 유형의 1차원 텐서 상수	(C17-C18), (C25)
(I14)	output_batch_dimension	si64 유형의 상수	(C20), (C25)
(I15)	output_feature_dimension	si64 유형의 상수	(C20), (C25), (C33)
(I16)	output_spatial_dimensions	si64 유형의 1차원 텐서 상수	(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의 다양한 enum 수입니다.	(C24)

출력

이름	유형	제약 조건
result	텐서 또는 양자화 텐서	(C25-C28), (C30-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) 주어진 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_tensor(lhs) and is_quantized_tensor(rhs) and is_quantized_tensor(result).
  • (C29) storage_type(lhs) = storage_type(rhs).
  • (C30) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C31) is_per_tensor_quantized(rhs)인 경우 is_per_tensor_quantized(result)입니다.
  • (C32) is_per_axis_quantized(rhs)인 경우 quantization_dimension(rhs) = kernel_output_feature_dimension입니다.
  • (C33) is_per_axis_quantized(result)인 경우 quantization_dimension(result) = output_feature_dimension입니다.

예

// %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 = dense<4> : tensor<2xi64>,
  padding = dense<0> : tensor<2x2xi64>,
  lhs_dilation = dense<2> : tensor<2xi64>,
  rhs_dilation = dense<1> : tensor<2xi64>,
  window_reversal = dense<false> : tensor<2xi1>,
  // 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]>,
  feature_group_count = 1 : i64,
  batch_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi32>, tensor<3x3x1x1xi32>) -> tensor<1x2x2x1xi32>
// %result: [[
//            [[10], [26]],
//            [[46], [62]]
//          ]]

코사인

시맨틱

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

• 부동 소수점 수: IEEE-754의 cos
• 복소수: 복소수 코사인입니다.
• 양자화 유형: 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 텐서의 선행 0비트 수를 요소별로 계산하고 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]]

예시 더보기

custom_call

시맨틱

inputs 및 called_computations를 가져와 results를 생성하는 구현 정의 작업 call_target_name를 캡슐화합니다. 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와 제수 rhs 텐서의 요소별 나눗셈을 수행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 정수: 소수 부분을 삭제하고 대수적 몫을 생성하는 정수 나눗셈.
• 부동 소수점 수: IEEE-754의 division
• 복소수: 복소수 나누기
• 양자화 유형:
  • 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

시맨틱

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))를 실행합니다.

이는 텐서당 양자화를 위한 의미 체계만 지정합니다. 축당 양자화 작업이 진행 중입니다. (#1574) 또한 향후에 하이브리드 양자화 지원을 추가할 수도 있습니다. (#1575)

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

• DEFAULT: 계산 속도는 가장 빠르지만 원래 숫자에 대한 근사치가 가장 정확합니다.
• HIGH: 계산은 느리지만 원래 숫자에 대한 근사치가 더 정확합니다.
• HIGHEST: 계산은 가장 느리지만 원래 숫자에 가장 정확합니다.

입력

라벨	이름	유형	제약 조건
(I1)	lhs	텐서 또는 텐서당 양자화 텐서	(C5-C6), (C9-C10), (C12-C16)
(I2)	rhs	텐서 또는 텐서당 양자화 텐서	(C7-C10), (C12)
(I3)	lhs_batching_dimensions	si64 유형의 1차원 텐서 상수	(C1), (C3), (C5), (C9), (C12)
(I4)	rhs_batching_dimensions	si64 유형의 1차원 텐서 상수	(C1), (C4), (C7), (C9)
(I5)	lhs_contracting_dimensions	si64 유형의 1차원 텐서 상수	(C2), (C3), (C6), (C10)
(I6)	rhs_contracting_dimensions	si64 유형의 1차원 텐서 상수	(C2), (C4), (C8), (C10)
(I7)	precision_config	DEFAULT, HIGH, HIGHEST의 다양한 enum 수입니다.	(C11)

출력

이름	유형	제약 조건
result	텐서 또는 텐서당 양자화 텐서	(C12), (C14), (C16)

제약 조건

• (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) and is_quantized(rhs) and is_quantized(result).
  • (C15) storage_type(lhs) = storage_type(rhs).
  • (C16) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C17) zero_points(rhs) = 0.

예

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

시맨틱

동적으로 계산된 시작 색인을 사용하여 operand에서 슬라이스를 추출하고 result 텐서를 생성합니다. start_indices에는 잠재적 조정 대상인 각 측정기준의 슬라이스 시작 색인이 포함되며 slice_sizes에는 각 크기의 슬라이스 크기가 포함됩니다. 보다 공식적으로 result[result_index] = operand[operand_index] 각 항목의 의미는 다음과 같습니다.

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

입력

라벨	이름	유형	제약 조건
(I1)	operand	텐서 또는 텐서당 양자화 텐서	(C1), (C2), (C4)
(I2)	start_indices	정수 유형 0차원 텐서의 가변 수	(C2), (C3)
(I3)	slice_sizes	si64 유형의 1차원 텐서 상수	(C2), (C4), (C5)

출력

이름	유형	제약 조건
result	텐서 또는 텐서당 양자화 텐서	(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 = dense<[2, 2]> : tensor<2xi64>
} : (tensor<4x4xi32>, tensor<i64>, tensor<i64>) -> tensor<2x2xi32>
// %result: [
//           [1, 1],
//           [1, 1]
//          ]

예시 더보기

dynamic_update_slice

시맨틱

start_indices에서 시작하는 슬라이스가 update의 값으로 업데이트된다는 점을 제외하고 operand 텐서와 동일한 result 텐서를 생성합니다. 더 공식적으로 result[result_index]는 다음과 같이 정의됩니다.

• 0 <= update_index < shape(update)인 경우 update[update_index]이며 각 항목의 의미는 다음과 같습니다.
  • adjusted_start_indices = clamp(0, start_indices, shape(operand) - shape(update)).
  • update_index = result_index - adjusted_start_indices.
• 그 밖의 경우에는 operand[result_index]입니다.

입력

라벨	이름	유형	제약 조건
(I1)	operand	텐서 또는 텐서당 양자화 텐서	(C1-C4), (C6)
(I2)	update	텐서 또는 텐서당 양자화 텐서	(C2), (C3), (C6)
(I3)	start_indices	정수 유형 0차원 텐서의 가변 수	(C4), (C5)

출력

이름	유형	제약 조건
result	텐서 또는 텐서당 양자화 텐서	(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]
//          ]

예시 더보기

지수

시맨틱

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

• 부동 소수점 수: IEEE-754의 exp
• 복소수: 복소 지수입니다.
• 양자화 유형: dequantize_op_quantize(exponential, operand, type(result)).

입력

라벨	이름	유형	제약 조건
(I1)	operand	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1)

출력

이름	유형	제약 조건
result	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(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

시맨틱

operand 텐서에서 요소별 지수 - 1 연산을 수행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 부동 소수점 수: IEEE-754의 expm1
• 복소수: 복소수 지수에서 1을 뺀 값입니다.
• 양자화 유형: dequantize_op_quantize(exponential_minus_one, operand, type(result)).

입력

라벨	이름	유형	제약 조건
(I1)	operand	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1)

출력

이름	유형	제약 조건
result	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(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

시맨틱

실제 및 복잡한 입력/출력에 대해 정방향 및 역 푸리에 변환을 수행합니다.

fft_type는 다음 중 하나입니다.

• FFT: 복잡에서 복잡하게 FFT를 전달합니다.
• IFFT: 복소수에서 복소수로의 역 FFT입니다.
• RFFT: 실수에서 복잡한 FFT로 전달합니다.
• IRFFT: 실수에서 복소로의 역 FFT입니다 (복소수를 받아 실수를 반환함).

좀 더 공식적으로 복잡한 유형의 1차원 텐서를 입력으로 사용하는 fft 함수를 통해 출력과 동일한 유형의 1차원 텐서를 생성하고 불연속 푸리에 변환을 계산합니다.

fft_type = FFT의 경우 result는 L = size(fft_length)인 일련의 L 계산의 최종 결과로 정의됩니다. 예를 들어 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]).

또한 동일한 유형 서명을 사용하고 fft의 역을 계산하는 ifft 함수가 지정된 경우 다음과 같습니다.

fft_type = IFFT의 경우 result은 fft_type = FFT의 계산의 역으로 정의됩니다. 예를 들어 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, ..., :]).

또한 부동 소수점 유형의 1차원 텐서를 사용하는 rfft 함수를 사용하면 동일한 부동 소수점 시맨틱스의 복합 유형의 1차원 텐서를 생성하고 다음과 같이 작동합니다.

• rfft(real_operand) = truncated_result 각 항목의 의미는 다음과 같습니다.
• complex_operand... = (real_operand..., 0.0).
• complex_result = fft(complex_operand).
• truncated_result = complex_result[:(rank(complex_result) / 2 + 1)].

(실제 피연산자에 대해 개별 푸리에 변환이 계산되면 결과의 첫 N/2 + 1 요소가 결과의 나머지 부분을 명확하게 정의하므로 rfft의 결과는 중복 요소가 계산되지 않도록 잘립니다.)

fft_type = RFFT의 경우 result는 L = size(fft_length)인 일련의 L 계산의 최종 결과로 정의됩니다. 예를 들어 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]).

마지막으로, 동일한 유형 서명을 가지고 있고 rfft의 역을 계산하는 irfft 함수가 지정된 경우 다음과 같습니다.

fft_type = IRFFT의 경우 result은 fft_type = RFFT의 계산의 역으로 정의됩니다. 예를 들어 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	부동 소수점 또는 복합 유형의 텐서	(C1), (C2), (C4), (C5)
(I2)	fft_type	FFT, IFFT, RFFT, IRFFT의 열거형	(C2), (C5)
(I3)	fft_length	si64 유형의 1차원 텐서 상수	(C1), (C3), (C4)

출력

이름	유형	제약 조건
result	부동 소수점 또는 복합 유형의 텐서	(C2), (C4), (C5)

제약 조건

• (C1) size(fft_length) <= rank(operand).
• (C2) operand 및 result 요소 유형 간의 관계가 다음과 같이 다릅니다.
  • fft_type = FFT, element_type(operand), element_type(result)의 복합 유형은 동일합니다.
  • fft_type = IFFT, element_type(operand), element_type(result)의 복합 유형은 동일합니다.
  • fft_type = RFFT인 경우 element_type(operand)는 부동 소수점 유형이고 element_type(result)는 동일한 부동 소수점 시맨틱스의 복합 유형입니다.
  • fft_type = IRFFT인 경우 element_type(operand)는 복합 유형이고 element_type(result)는 동일한 부동 소수점 의미 체계의 부동 소수점 유형입니다.
• (C3) 1 <= size(fft_length) <= 3.
• (C4) operand와 result 중에 부동 소수점 유형의 텐서 real가 있다면 shape(real)[-size(fft_length):] = fft_length입니다.
• (C5) shape(result) = shape(operand)(다음 제외):
  • fft_type = RFFT인 경우 dim(result, -1) = dim(operand, -1) = 0 ? 0 : dim(operand, -1) / 2 + 1입니다.
  • 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 = dense<4> : tensor<1xi64>
} : (tensor<4xcomplex<f32>>) -> tensor<4xcomplex<f32>>
// %result: [(1.0, 0.0), (1.0, 0.0), (1.0, 0.0), (1.0, 0.0)]

floor

시맨틱

operand 텐서의 요소별 하한선을 수행하고 result 텐서를 생성합니다. IEEE-754 사양의 roundToIntegralTowardNegative 작업을 구현합니다. 양자화 유형의 경우 dequantize_op_quantize(floor, 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.floor"(%operand) : (tensor<5xf32>) -> tensor<5xf32>
// %result: [-1.0, -1.0, 0.0, 0.0, 2.0]

예시 더보기

수집하다

시맨틱

start_indices에 지정된 오프셋에서 operand 텐서의 슬라이스를 수집하고 result 텐서를 생성합니다.

다음 다이어그램은 구체적인 예를 사용하여 result의 요소가 operand의 요소에 매핑되는 방식을 보여줍니다. 이 다이어그램은 몇 가지 result 색인 예를 선택하여 어떤 operand 색인에 해당하는지 자세히 설명합니다.

더 공식적으로 result[result_index] = operand[operand_index]이며, 각 항목의 의미는 다음과 같습니다.

• batch_dims = [d for d in axes(result) and d not in offset_dims].
• batch_index = result_index[batch_dims...].
• start_index는 다음과 같이 정의됩니다.
  • start_indices[bi0, ..., :, ..., biN], 여기서 bi는 batch_index의 개별 요소이고 :가 index_vector_dim < rank(start_indices)인 경우 index_vector_dim 색인에 삽입됩니다.
  • 그 밖의 경우에는 [start_indices[batch_index]]입니다.
• axes(operand)의 d_operand:
  • d_operand = start_index_map[d_start]인 경우 full_start_index[d_operand] = clamp(start_index[d_start], 0, dim(operand, d_operand) - slice_sizes[d_operand])입니다.
  • 그 밖의 경우에는 full_start_index[d_operand] = 0입니다.
• offset_index = result_index[offset_dims...].
• full_offset_index = [oi0, ..., 0, ..., oiN], 여기서 oi는 offset_index의 개별 요소이고 0는 collapsed_slice_dims의 색인에 삽입됩니다.
• operand_index = full_start_index + full_offset_index.

indices_are_sorted이 true이면 구현에서는 start_indices가 start_index_map를 기준으로 정렬되었다고 가정할 수 있습니다. 그렇지 않으면 동작이 정의되지 않습니다. 좀 더 공식적으로 indices(result)의 모든 i1 < i2에 대해 full_start_index(i1) <= full_start_index(i2)입니다.

입력

라벨	이름	유형	제약 조건
(I1)	operand	텐서 또는 텐서당 양자화 텐서	(C1), (C7), (C10-C12), (C14)
(I2)	start_indices	정수 유형의 텐서	(C2), (C3), (C13)
(I3)	offset_dims	si64 유형의 1차원 텐서 상수	(C1), (C4-C5), (C13)
(I4)	collapsed_slice_dims	si64 유형의 1차원 텐서 상수	(C1), (C6-C8), (C13)
(I5)	start_index_map	si64 유형의 1차원 텐서 상수	(C3), (C9), (C10)
(I6)	index_vector_dim	si64 유형의 상수	(C2), (C3), (C13)
(I7)	slice_sizes	si64 유형의 1차원 텐서 상수	(C8), (C11-C13)
(I8)	indices_are_sorted	i1 유형의 상수

출력

이름	유형	제약 조건
result	텐서 또는 텐서당 양자화 텐서	(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) 각 항목의 의미는 다음과 같습니다.
  • batch_dim_sizes = shape(start_indices). 단, index_vector_dim에 해당하는 start_indices의 크기 크기는 포함되지 않습니다.
  • offset_dim_sizes = shape(slice_sizes): collapsed_slice_dims에 해당하는 slice_sizes의 크기 크기가 포함되지 않은 경우를 제외하고
  • combine은 batch_dims에 해당하는 축에 batch_dim_sizes를 배치하고 offset_dims에 해당하는 축에 offset_dim_sizes를 배치합니다.
• (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]]
//                 ]
%result = "stablehlo.gather"(%operand, %start_indices) {
  dimension_numbers = #stablehlo.gather<
    offset_dims = [2, 3],
    collapsed_slice_dims = [0],
    start_index_map = [1, 0],
    index_vector_dim = 2>,
  slice_sizes = dense<[1, 2, 2]> : tensor<3xi64>,
  indices_are_sorted = false
} : (tensor<3x4x2xi32>, tensor<2x3x2xi64>) -> tensor<2x3x2x2xi32>
// %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]]
//            ]
//          ]

예시 더보기

get_dimension_size

시맨틱

operand의 지정된 dimension 크기를 생성합니다. 더 공식적으로는 result = dim(operand, dimension)입니다.

입력

라벨	이름	유형	제약 조건
(I1)	operand	텐서	(C1)
(I2)	dimension	si64 유형의 상수	(C1)

출력

이름	유형
result	si32 유형의 0차원 텐서

제약 조건

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

시맨틱

operand 튜플의 index 위치에 있는 요소를 추출하고 result를 생성합니다. 더 공식적으로 result = operand[index]입니다.

입력

라벨	이름	유형	제약 조건
(I1)	operand	tuple	(C1), (C2)
(I2)	index	si32 유형의 상수	(C1), (C2)

출력

이름	유형	제약 조건
result	지원되는 모든 유형	(C2)

제약 조건

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

예

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

예시 더보기

if

시맨틱

pred 값에 따라 true_branch 또는 false_branch에서 정확히 하나의 함수를 실행하여 출력을 생성합니다. 더 공식적으로 result = pred ? true_branch() : false_branch()입니다.

입력

라벨	이름	유형	제약 조건
(I1)	pred	i1 유형의 0차원 텐서
(I2)	true_branch	함수	(C1-C3)
(I3)	false_branch	함수	(C1), (C2)

출력

이름	유형	제약 조건
results	텐서, 양자화 텐서 또는 토큰의 가변 수	(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

예시 더보기

이미지

시맨틱

operand에서 요소별로 허수부를 추출하고 result 텐서를 생성합니다. 더 공식적으로 각 요소 x에 대해 imag(x) = is_complex(x) ? imaginary_part(x) : constant(0, element_type(result))입니다.

입력

라벨	이름	유형	제약 조건
(I1)	operand	부동 소수점 또는 복합 유형의 텐서	(C1), (C2)

출력

이름	유형	제약 조건
result	부동 소수점 유형의 텐서	(C1), (C2)

제약 조건

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

예

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

예시 더보기

인피드

시맨틱

인피드에서 데이터를 읽고 results를 생성합니다.

infeed_config의 시맨틱은 구현을 통해 정의됩니다.

results는 먼저 나오는 페이로드 값과 마지막에 나오는 토큰으로 구성됩니다. 앞으로는 명확성을 높이기 위해 페이로드와 토큰을 별도의 두 출력으로 분할할 계획입니다. (#670)

입력

라벨	이름	유형
(I1)	token	token
(I2)	infeed_config	string 유형의 상수

출력

이름	유형	제약 조건
results	텐서, 양자화 텐서 또는 토큰의 가변 수	(C1-C3)

제약 조건

• (C1) 0 < size(results).
• (C2) is_empty(result[:-1]) 또는 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]]

예시 더보기

아이오타

시맨틱

output 텐서에 iota_dimension 차원을 따라 0부터 오름차순으로 값을 채웁니다. 더 공식적으로

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

입력

라벨	이름	유형	제약 조건
(I1)	iota_dimension	si64	(C1)

출력

이름	유형	제약 조건
output	정수의 텐서, 부동 소수점 또는 복합 유형 또는 텐서당 양자화 텐서	(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

시맨틱

x의 값이 유한한지 (즉, +Inf, -Inf, NaN이 아님) 요소별 검사를 실행하고 y 텐서를 생성합니다. IEEE-754 사양의 isFinite 작업을 구현합니다. 양자화 유형의 경우 결과는 항상 true입니다.

입력

라벨	이름	유형	제약 조건
(I1)	x	부동 소수점 유형 또는 텐서당 양자화 텐서의 텐서	(C1)

출력

이름	유형	제약 조건
y	불리언 유형의 텐서	(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]

예시 더보기

로그

시맨틱

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

• 부동 소수점 수: IEEE-754의 log
• 복소수: 복소수 로그입니다.
• 양자화 유형: dequantize_op_quantize(log, operand, type(result))

입력

라벨	이름	유형	제약 조건
(I1)	operand	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1)

출력

이름	유형	제약 조건
result	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(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

시맨틱

operand 텐서에서 요소별 로그와 하나의 연산을 실행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 부동 소수점 수: IEEE-754의 logp1
• 복소수: 복소수 로그 +1
• 양자화 유형: dequantize_op_quantize(log_plus_one, operand, type(result)).

입력

라벨	이름	유형	제약 조건
(I1)	operand	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1)

출력

이름	유형	제약 조건
result	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(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]

예시 더보기

로지스틱

시맨틱

operand 텐서에서 요소별 로지스틱 연산을 수행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 부동 소수점 수: IEEE-754의 division(1, addition(1, exp(-x)))
• 복소수: 복잡한 로지스틱
• 양자화 유형: dequantize_op_quantize(logistic, operand, type(result)).

입력

라벨	이름	유형	제약 조건
(I1)	operand	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1)

출력

이름	유형	제약 조건
result	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(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]]

예시 더보기

지도

시맨틱

dimensions와 함께 inputs에 맵 함수 computation를 적용하고 result 텐서를 생성합니다.

더 공식적으로 result[result_index] = computation(inputs...[result_index])입니다. dimensions는 현재 사용되지 않으며 향후 삭제될 수 있습니다. (#487)

입력

라벨	이름	유형	제약 조건
(I1)	inputs	가변 텐서 또는 텐서당 양자화 텐서	(C1-C4)
(I2)	dimensions	si64 유형의 1차원 텐서 상수	(C3)
(I3)	computation	함수	(C4)

출력

이름	유형	제약 조건
result	텐서 또는 텐서당 양자화 텐서	(C1), (C4)

제약 조건

• (C1) shape(inputs...) = shape(result).
• (C2) 0 < size(inputs) = N.
• (C3) dimensions = range(rank(inputs[0])).
• (C4) computation에는 Ei = element_type(inputs[i]) 및 E' = element_type(result)인 (tensor<E0>, ..., tensor<EN-1>) -> tensor<E'> 유형이 있습니다.

예

// %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 = dense<[0, 1]> : tensor<2xi64>
} : (tensor<2x2xi64>, tensor<2x2xi64>) -> tensor<2x2xi64>
// %result: [[0, 5], [12, 21]]

예시 더보기

최대

시맨틱

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

• 불리언: 논리 OR입니다.
• 정수: 정수 최댓값
• 부동 소수점 수: IEEE-754의 maximum
• 복소수: (real, imaginary) 쌍의 사전식 최댓값입니다. 복소수에 정렬을 적용하는 데는 놀라운 의미 체계가 포함되므로 앞으로 이 연산에서 복소수에 대한 지원을 중단할 계획입니다. (#560)
• 양자화 유형:
  • dequantize_op_quantize(maximum, lhs, rhs, type(result)).

입력

라벨	이름	유형	제약 조건
(I1)	lhs	텐서 또는 텐서당 양자화 텐서	(C1)
(I2)	rhs	텐서 또는 텐서당 양자화 텐서	(C1)

출력

이름	유형	제약 조건
result	텐서 또는 텐서당 양자화 텐서	(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]]

예시 더보기

최소

시맨틱

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

• 불리언: 논리 AND
• 정수: 정수 최솟값
• 부동 소수점 수: IEEE-754의 minimum
• 복소수: (real, imaginary) 쌍의 사전식 최솟값입니다. 복소수에 정렬을 적용하는 데는 놀라운 의미 체계가 포함되므로 앞으로 이 연산에서 복소수에 대한 지원을 중단할 계획입니다. (#560)
• 양자화 유형:
  • dequantize_op_quantize(minimum, lhs, rhs, type(result)).

입력

라벨	이름	유형	제약 조건
(I1)	lhs	텐서 또는 텐서당 양자화 텐서	(C1)
(I2)	rhs	텐서 또는 텐서당 양자화 텐서	(C1)

출력

이름	유형	제약 조건
result	텐서 또는 텐서당 양자화 텐서	(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]]

예시 더보기

곱하기

시맨틱

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

• 불리언: 논리 AND
• 정수: 정수 곱셈.
• 부동 소수점 수: IEEE-754의 multiplication
• 복소수: 복소수 곱셈.
• 양자화 유형:
  • dequantize_op_quantize(multiply, lhs, rhs, type(result)).

입력

라벨	이름	유형	제약 조건
(I1)	lhs	텐서 또는 텐서당 양자화 텐서	(C1)
(I2)	rhs	텐서 또는 텐서당 양자화 텐서	(C1)

출력

이름	유형	제약 조건
result	텐서 또는 텐서당 양자화 텐서	(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

시맨틱

operand 텐서의 요소별 부정을 수행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 부호 있는 정수: 정수 부정
• 부호 없는 정수: 부호 있는 정수로 비트캐스트, 정수 부정, 부호 없는 정수로 비트캐스트
• 부동 소수점 수: IEEE-754의 negate
• 복소수: 복합 부정입니다.
• 양자화 유형: dequantize_op_quantize(negate, operand, type(result)).

입력

라벨	이름	유형	제약 조건
(I1)	operand	정수, 부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1)

출력

이름	유형	제약 조건
result	정수, 부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(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]

예시 더보기

않는

시맨틱

텐서 operand의 요소별 NOT을 실행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 불리언: 논리 NOT.
• 정수: 비트 NOT.

인수

이름	유형	제약 조건
operand	불리언 또는 정수 유형의 텐서	(C1)

출력

이름	유형	제약 조건
result	불리언 또는 정수 유형의 텐서	(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

시맨틱

operand를 생성하는 작업이 result에 종속되는 작업보다 먼저 실행되도록 하고 컴파일러 변환이 배리어를 통해 작업을 이동하지 않도록 합니다. 그 외에는 작업은 ID(예: result = operand)입니다.

인수

이름	유형	제약 조건
operand	다양한 개수의 텐서, 텐서당 양자화 텐서 또는 토큰	(C1)

출력

이름	유형	제약 조건
result	다양한 개수의 텐서, 텐서당 양자화 텐서 또는 토큰	(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

예시 더보기

또는

시맨틱

두 텐서 lhs 및 rhs에 대해 요소별 OR을 수행하여 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 불리언: 논리 OR입니다.
• 정수의 경우 비트 OR입니다.

입력

라벨	이름	유형	제약 조건
(I1)	lhs	정수 또는 불리언 유형의 텐서	(C1)
(I2)	rhs	정수 또는 불리언 유형의 텐서	(C1)

출력

이름	유형	제약 조건
result	정수 또는 불리언 유형의 텐서	(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]]

아웃피드

시맨틱

아웃피드에 inputs를 작성하고 result 토큰을 생성합니다.

outfeed_config의 시맨틱은 구현을 통해 정의됩니다.

입력

라벨	이름	유형
(I1)	inputs	가변 텐서 또는 양자화 텐서
(I2)	token	token
(I3)	outfeed_config	string 유형의 상수

출력

이름	유형
result	token

예

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

예시 더보기

패드

시맨틱

텐서 주변과 지정된 padding_value이 있는 텐서 요소 사이의 패딩으로 operand를 확장합니다.

edge_padding_low와 edge_padding_high는 각 측정기준의 하위 값 (색인 0 옆)과 높은 값 (가장 높은 색인 옆)에 추가된 패딩의 양을 각각 지정합니다. 패딩 값은 음수일 수 있습니다. 여기서 음수 패딩의 절댓값은 지정된 크기에서 삭제할 요소의 수를 나타냅니다.

interior_padding은 각 차원에서 음수가 될 수 없는 두 요소 사이에 추가되는 패딩의 양을 지정합니다. 내부 패딩은 가장자리 패딩 전에 발생합니다. 그러면 음의 가장자리 패딩이 내부 패딩 처리된 피연산자에서 요소를 삭제할 수 있습니다.

더 공식적으로 result[result_index]는 다음과 같이 정의됩니다.

• result_index = edge_padding_low + operand_index * (interior_padding + 1)인 경우 operand[operand_index]입니다.
• 그 밖의 경우에는 padding_value입니다.

입력

라벨	이름	유형	제약 조건
(I1)	operand	텐서 또는 텐서당 양자화 텐서	(C1), (C2), (C4)
(I2)	padding_value	0차원 텐서 또는 텐서당 양자화 텐서	(C1)
(I3)	edge_padding_low	si64 유형의 1차원 텐서 상수	(C1), (C4)
(I4)	edge_padding_high	si64 유형의 1차원 텐서 상수	(C1), (C4)
(I5)	interior_padding	si64 유형의 1차원 텐서 상수	(C2-C4)

출력

이름	유형	제약 조건
result	텐서 또는 텐서당 양자화 텐서	(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 = dense<[0, 1]> : tensor<2xi64>,
  edge_padding_high = dense<[2, 1]> : tensor<2xi64>,
  interior_padding = dense<[1, 2]> : tensor<2xi64>
} : (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

시맨틱

현재 프로세스의 partition_id를 생성합니다.

출력

이름	유형
result	ui32 유형의 0차원 텐서

예

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

예시 더보기

팝컨트

시맨틱

operand 텐서에 설정된 비트 수의 요소별 개수를 실행하고 result 텐서를 생성합니다.

입력

라벨	이름	유형	제약 조건
(I1)	operand	정수 유형의 텐서	(C1)

출력

이름	유형	제약 조건
result	정수 유형의 텐서	(C1)

제약 조건

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

예

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

예시 더보기

전력

시맨틱

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

• 정수: 정수 지수
• 부동 소수점 수: IEEE-754의 pow
• 복소수: 복소수 지수
• 양자화 유형: dequantize_op_quantize(power, lhs, rhs, type(result))

입력

라벨	이름	유형	제약 조건
(I1)	lhs	정수, 부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1)
(I2)	rhs	정수, 부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1)

출력

이름	유형	제약 조건
result	정수, 부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(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]

예시 더보기

real

시맨틱

operand에서 요소별로 실수 부분을 추출하고 result 텐서를 생성합니다. 더 공식적으로 각 요소 x에 대해 real(x) = is_complex(x) ? real_part(x) : x입니다.

입력

라벨	이름	유형	제약 조건
(I1)	operand	부동 소수점 또는 복합 유형의 텐서	(C1), (C2)

출력

이름	유형	제약 조건
result	부동 소수점 유형의 텐서	(C1), (C2)

제약 조건

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

예

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

예시 더보기

Recv

시맨틱

channel_id를 사용하여 채널에서 데이터를 수신하고 results를 생성합니다.

is_host_transfer이 true이면 작업은 호스트에서 데이터를 전송합니다. 그렇지 않으면 다른 기기에서 데이터를 전송합니다. 이것이 의미하는 바는 구현이 정의되어 있습니다. 이 플래그는 channel_type에 제공된 정보를 복제하므로 향후에는 이 중 하나만 유지할 계획입니다. (#666)

results는 먼저 나오는 페이로드 값과 마지막에 나오는 토큰으로 구성됩니다. 앞으로는 명확성을 높이기 위해 페이로드와 토큰을 별도의 두 출력으로 분할할 계획입니다. (#670)

입력

라벨	이름	유형	제약 조건
(I1)	token	token	(C4)
(I2)	channel_id	si64 유형의 상수
(I3)	channel_type	DEVICE_TO_DEVICE 및 HOST_TO_DEVICE의 열거형	(C1)
(I4)	is_host_transfer	i1 유형의 상수	(C1)

출력

이름	유형	제약 조건
results	텐서, 양자화 텐서 또는 토큰의 가변 수	(C2-C4)

제약 조건

• (C1) channel_type는 다음과 같이 정의됩니다.
  • is_host_transfer = true인 경우 HOST_TO_DEVICE
  • 그 밖의 경우에는 DEVICE_TO_DEVICE입니다.
• (C2) 0 < size(results).
• (C3) is_empty(result[:-1]) 또는 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)

예시 더보기

reduce

시맨틱

축소 함수 body를 dimensions와 함께 inputs 및 init_values에 적용하고 results 텐서를 생성합니다.

축소 순서는 구현에 따라 정의됩니다. 즉, body 및 init_values는 작업이 모든 구현의 모든 입력에 대해 동일한 결과를 생성하도록 모노이드를 형성해야 합니다. 그러나 널리 사용되는 여러 감소에는 이 조건이 적용되지 않습니다. 예를 들어 body의 경우 부동 소수점 덧셈, init_values의 경우 0은 실제로 모노이드를 형성하지 않습니다. 부동 소수점 덧셈이 결합 법칙이 아니기 때문입니다.

더 공식적으로 results...[j0, ..., jR-1] = reduce(input_slices_converted)이며, 각 항목의 의미는 다음과 같습니다.

• input_slices = inputs...[j0, ..., :, ..., jR-1], 여기서 :는 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:])...).
• 일부 바이너리 트리 schedule의 경우 reduce(input_slices_converted) = exec(schedule)이며 각 항목의 의미는 다음과 같습니다.
  • exec(node) = body(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value.
• schedule는 구현 정의 전체 바이너리 트리이며, 순서 순회는 다음과 같이 구성됩니다.
  • index_space(input_slices_converted)에 있는 모든 index의 경우 input_slices_converted...[index] 값(index의 사전 오름차순으로 표시됨)
  • 구현 정의 위치에 구현에서 정의한 양의 init_values_converted가 산재되어 있습니다.

입력

라벨	이름	유형	제약 조건
(I1)	inputs	가변 텐서 또는 텐서당 양자화 텐서	(C1-C4), (C6), (C7)
(I2)	init_values	0차원 텐서 또는 텐서당 양자화 텐서의 가변 수	(C2), (C3)
(I3)	dimensions	si64 유형의 1차원 텐서 상수	(C4), (C5), (C7)
(I4)	body	함수	(C6)

출력

이름	유형	제약 조건
results	가변 텐서 또는 텐서당 양자화 텐서	(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는 (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>) 유형이며 여기서 is_promotable(element_type(inputs[i]), Ei)입니다.
• (C7) shape(results...) = shape(inputs...), dimensions에 해당하는 inputs... 크기 크기가 포함되지 않음
• (C8) [0,N)의 모든 i에 대해 element_type(results[i]) = Ei.

예

// %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 = dense<1> : tensor<1xi64>
} : (tensor<1x6xi64>, tensor<i64>) -> tensor<1xi64>
// %result = [15]

예시 더보기

reduce_precision

시맨틱

exponent_bits 및 mantissa_bits를 사용하고 원래의 부동 소수점 유형으로 다시 operand를 다른 부동 소수점 유형으로 변환하여 output 텐서를 생성하는 요소별 변환을 수행합니다.

더 공식적으로

• 원래 값의 가수 비트는 roundToIntegralTiesToEven 시맨틱스를 사용하여 원래 값을 mantissa_bits로 나타낼 수 있는 가장 가까운 값으로 반올림하도록 업데이트됩니다.
• 그런 다음 mantissa_bits가 원래 값의 가수 비트 수보다 작으면 가수 비트가 mantissa_bits로 잘립니다.
• 그런 다음 중간 결과의 지수 비트가 exponent_bits에서 제공하는 범위에 맞지 않으면 중간 결과가 원래 부호를 사용하여 무한대로 오버플로되고, 원래 부호를 사용하면 언더플로가 0이 됩니다.
• 양자화 유형의 경우 dequantize_op_quantize( lambda operand: reduce_precision(operand, exponent_bits, mantissa_bits), operand, type(result))를 실행합니다.

입력

라벨	이름	유형	제약 조건
(I1)	operand	부동 소수점 유형 또는 텐서당 양자화 텐서의 텐서	(C1)
(I2)	exponent_bits	si32 유형의 상수	(C2)
(I3)	mantissa_bits	si32 유형의 상수	(C3)

출력

이름	유형	제약 조건
output	부동 소수점 유형 또는 텐서당 양자화 텐서의 텐서	(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

시맨틱

StableHLO 프로세스 그리드의 각 프로세스 그룹 내에서 computations를 사용하여 각 프로세스의 operand 텐서 값에 대해 축소를 실행하고 scatter_dimension에 따라 축소 결과를 여러 부분으로 분할하고 분할된 부분을 프로세스 간에 분산해 result를 생성합니다.

이 작업은 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 내에서 다음을 실행합니다.

• 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).
• process_group의 모든 sender에 대한 result@receiver = parts@sender[receiver_index]. 여기서 receiver_index = process_group.index(receiver).

입력

라벨	이름	유형	제약 조건
(I1)	operand	텐서 또는 텐서당 양자화 텐서	(C1), (C2), (C7), (C8)
(I2)	scatter_dimension	si64 유형의 상수	(C1), (C2), (C8)
(I3)	replica_groups	si64 유형의 2차원 텐서 상수	(C3-C5)
(I4)	channel_id	si64 유형의 상수	(C6)
(I5)	use_global_device_ids	i1 유형의 상수	(C6)
(I6)	computation	함수	(C7)

출력

이름	유형	제약 조건
result	텐서 또는 텐서당 양자화 텐서	(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)는 다음과 같이 정의됩니다.
  • cross_replica가 사용되면 num_replicas입니다.
  • cross_replica_and_partition가 사용되면 num_replicas입니다.
  • flattened_ids가 사용되면 num_processes입니다.
• (C5) 0 <= replica_groups < size(replica_groups).
• (C6) use_global_device_ids = true인 경우 channel_id > 0입니다.
• (C7) computation에는 is_promotable(element_type(operand), E)인 (tensor<E>, tensor<E>) -> (tensor<E>) 유형이 있습니다.
• (C8) shape(result) = shape(operand) 제외:
  • 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

시맨틱

축소 함수 body를 inputs 및 init_values의 창에 적용하고 results를 생성합니다.

다음 다이어그램은 구체적인 예를 사용하여 inputs...에서 results...의 요소가 계산되는 방식을 보여줍니다.

보다 공식적으로 results...[result_index] = reduce(windows, init_values, axes(inputs...), body)(감소 참고) 각 항목의 의미는 다음과 같습니다.

• 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	가변 텐서 또는 텐서당 양자화 텐서	(C1~C4), (C6), (C8), (C10), (C12), (C13), (C15)
(I2)	init_values	0차원 텐서 또는 텐서당 양자화 텐서의 가변 수	(C1), (C13)
(I3)	window_dimensions	si64 유형의 1차원 텐서 상수	(C4), (C5), (C15)
(I4)	window_strides	si64 유형의 1차원 텐서 상수	(C6), (C7), (C15)
(I5)	base_dilations	si64 유형의 1차원 텐서 상수	(C8), (C9), (C15)
(I6)	window_dilations	si64 유형의 1차원 텐서 상수	(C10), (C11), (C15)
(I7)	padding	si64 유형의 2차원 텐서 상수	(C12), (C15)
(I8)	body	함수	(C13)

출력

이름	유형	제약 조건
results	가변 텐서 또는 텐서당 양자화 텐서	(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에는 is_promotable(element_type(inputs[i]), Ei)가 있는 (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>) 유형이 있습니다.
• (C14) same(shape(results...)).
• (C15) shape(results[0]) = num_windows 각 항목의 의미는 다음과 같습니다.
  • 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) [0,N)의 모든 i에 대해 element_type(results[i]) = Ei.

예

// %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 = dense<[2, 1]> : tensor<2xi64>,
  window_strides = dense<[4, 1]> : tensor<2xi64>,
  base_dilations = dense<[2, 1]> : tensor<2xi64>,
  window_dilations = dense<[3, 1]> : tensor<2xi64>,
  padding = dense<[[2, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<3x2xi64>, tensor<i64>) -> tensor<2x2xi64>
// %result = [[0, 0], [3, 4]]

예시 더보기

나머지

시맨틱

피제수 lhs 및 제수 rhs 텐서의 요소별 나머지를 실행하고 result 텐서를 생성합니다.

더 공식적으로 결과의 부호를 피제수에서 가져오고, 결과의 절대값은 항상 제수의 절댓값보다 작습니다. 나머지는 lhs - d * rhs로 계산되며, 여기서 d는 다음과 같이 지정됩니다.

• 정수의 경우 stablehlo.divide(lhs, rhs).
• 부동 소수점 수: 반올림 속성 roundTowardZero가 있는 IEEE-754의 division(lhs, rhs)
• 복소수는 미정입니다. (#997)
• 양자화 유형:
  • dequantize_op_quantize(remainder, lhs, rhs, type(result)).

부동 소수점 요소 유형의 경우 이 연산은 IEEE-754 사양의 remainder 연산과 대조됩니다. 여기서 d는 정확한 lhs/rhs 값에 가장 가까운 정수이며 짝수에 연결됩니다.

입력

라벨	이름	유형	제약 조건
(I1)	lhs	정수의 텐서, 부동 소수점 또는 복합 유형 또는 텐서당 양자화 텐서	(C1)
(I2)	rhs	정수의 텐서, 부동 소수점 또는 복합 유형 또는 텐서당 양자화 텐서	(C1)

출력

이름	유형	제약 조건
result	정수의 텐서, 부동 소수점 또는 복합 유형 또는 텐서당 양자화 텐서	(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

시맨틱

현재 프로세스의 replica_id를 생성합니다.

출력

이름	유형
result	ui32 유형의 0차원 텐서

예

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

예시 더보기

형태 변경

시맨틱

operand 텐서의 형태를 result 텐서로 변경합니다. 개념적으로는 동일한 표준 표현을 유지하지만 잠재적으로 모양을 변경(예: tensor<2x3xf32>에서 tensor<3x2xf32> 또는 tensor<6xf32>로)하는 것입니다.

좀 더 공식적으로 result[result_index] = operand[operand_index], 여기서 result_index와 operand_index는 index_space(result) 및 index_space(operand)의 사전순으로 동일한 위치에 있습니다.

입력

라벨	이름	유형	제약 조건
(I1)	operand	텐서 또는 양자화 텐서	(C1-C3)

출력

이름	유형	제약 조건
result	텐서 또는 양자화 텐서	(C1-C3)

제약 조건

• (C1) element_type(result)는 다음과 같이 주어집니다.
  • !is_per_axis_quantized(operand)인 경우 element_type(operand)입니다.
  • quantization_dimension(operand)와 quantization_dimension(result)가 다를 수 있다는 점을 제외하고 element_type(operand)입니다.
• (C2) size(operand) = size(result).
• (C3) 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]]

예시 더보기

reverse

시맨틱

지정된 dimensions에 따라 operand에 있는 요소의 순서를 반대로 바꾸고 result 텐서를 생성합니다. 보다 공식적으로 result[result_index] = operand[operand_index] 각 항목의 의미는 다음과 같습니다.

• dimensions에서 d인 경우 operand_index[d] = dim(result, d) - result_index[d] - 1
• 그 밖의 경우에는 operand_index[d] = result_index[d]입니다.

입력

라벨	이름	유형	제약 조건
(I1)	operand	텐서 또는 텐서당 양자화 텐서	(C1), (C3)
(I2)	dimensions	si64 유형의 1차원 텐서 상수	(C2), (C3)

출력

이름	유형	제약 조건
result	텐서 또는 텐서당 양자화 텐서	(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 = dense<1> : tensor<1xi64>
} : (tensor<3x2xi32>) -> tensor<3x2xi32>
// %result: [[2, 1], [4, 3], [6, 5]]

예시 더보기

RN

시맨틱

rng_distribution 알고리즘을 사용하여 랜덤 숫자를 생성하고 지정된 형태 shape의 result 텐서를 생성합니다.

rng_distribution = UNIFORM인 경우 구간 [a, b)에 걸쳐 균일한 분포를 따라 난수가 생성됩니다. a >= b인 경우 동작이 정의되지 않습니다.

rng_distribution = NORMAL인 경우 평균 = a, 표준 편차 = b인 정규 분포를 따라 난수가 생성됩니다. b < 0인 경우 동작이 정의되지 않습니다.

랜덤 숫자가 생성되는 정확한 방법은 구현을 통해 정의됩니다. 예를 들어 확정적일 수도 있고 확정적이지 않을 수도 있으며 숨겨진 상태를 사용할 수도 있고 사용하지 않을 수도 있습니다.

많은 이해관계자와의 대화에서 이 작업은 사실상 지원 중단되었기 때문에 향후 삭제를 살펴볼 계획입니다. (#597)

입력

라벨	이름	유형	제약 조건
(I1)	a	정수, 불리언 또는 부동 소수점 유형의 0차원 텐서	(C1), (C2)
(I2)	b	정수, 불리언 또는 부동 소수점 유형의 0차원 텐서	(C1), (C2)
(I3)	shape	si64 유형의 1차원 텐서 상수	(C3)
(I4)	rng_distribution	UNIFORM 및 NORMAL의 열거형	(C2)

출력

이름	유형	제약 조건
result	정수, 불리언 또는 부동 소수점 유형의 텐서	(C1-C3)

제약 조건

• (C1) element_type(a) = element_type(b) = element_type(result).
• (C2) rng_distribution = NORMAL인 경우 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

시맨틱

초기 상태 initial_state이 지정된 의사 난수 생성기 알고리즘 rng_algorithm를 사용하여 균일한 랜덤 비트로 채워진 output와 업데이트된 출력 상태 output_state를 반환합니다. 출력은 initial_state의 확정적 함수가 보장되지만 구현 간에 확정적이 되지 않을 수도 있습니다.

rng_algorithm는 다음 중 하나입니다.

• DEFAULT: 구현 정의 알고리즘입니다.
• THREE_FRY: Threefry 알고리즘의 구현 정의 변형*
• PHILOX: Philox 알고리즘의 구현 정의 변형*

* 참고: Salmon 외 SC 2011년. 병렬 랜덤 숫자: 1, 2, 3만큼 쉬움

입력

라벨	이름	유형	제약 조건
(I1)	rng_algorithm	DEFAULT, THREE_FRY, PHILOX의 열거형	(C2)
(I2)	initial_state	ui64 유형의 1차원 텐서	(C1), (C2)

출력

이름	유형	제약 조건
output_state	ui64 유형의 1차원 텐서	(C1)
output	정수 또는 부동 소수점 유형의 텐서

제약 조건

• (C1) type(initial_state) = type(output_state).
• (C2) size(initial_state)는 다음과 같이 정의됩니다.
  • rng_algorithm = DEFAULT인 경우 구현이 정의됩니다.
  • rng_algorithm = THREE_FRY인 경우 2입니다.
  • 2 또는 3인 경우 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

시맨틱

operand 텐서에서 가장 가까운 정수를 향해 요소별 반올림을 수행하여 0으로부터의 관계를 끊고 result 텐서를 생성합니다. IEEE-754 사양에서 roundToIntegralTiesToAway 작업을 구현합니다. 양자화 유형의 경우 dequantize_op_quantize(round_nearest_afz, operand, type(result))를 실행합니다.

입력

라벨	이름	유형	제약 조건
(I1)	operand	부동 소수점 유형 또는 텐서당 양자화 텐서의 텐서	(C1)

출력

이름	유형	제약 조건
result	부동 소수점 유형 또는 텐서당 양자화 텐서의 텐서	(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

시맨틱

operand 텐서에서 가장 가까운 정수에 대한 요소별 반올림을 수행하여 짝수 정수와의 연결을 해제하고 result 텐서를 생성합니다. IEEE-754 사양의 roundToIntegralTiesToEven 작업을 구현합니다. 양자화 유형의 경우 dequantize_op_quantize(round_nearest_even, operand, type(result))를 실행합니다.

입력

라벨	이름	유형	제약 조건
(I1)	operand	부동 소수점 유형 또는 텐서당 양자화 텐서의 텐서	(C1)

출력

이름	유형	제약 조건
result	부동 소수점 유형 또는 텐서당 양자화 텐서의 텐서	(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

시맨틱

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

• 부동 소수점 수: IEEE-754의 rSqrt
• 복소수: 복소수 역수 제곱근입니다.
• 양자화 유형: dequantize_op_quantize(rsqrt, operand, type(result))

입력

라벨	이름	유형	제약 조건
(I1)	operand	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1)

출력

이름	유형	제약 조건
result	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(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]]

예시 더보기

scatter

시맨틱

scatter_indices로 지정된 여러 슬라이스가 update_computation를 사용하여 updates 값으로 업데이트되는 것을 제외하고 inputs 텐서와 동일한 results 텐서를 생성합니다.

다음 다이어그램은 구체적인 예를 사용하여 updates...의 요소가 results...의 요소에 매핑되는 방식을 보여줍니다. 이 다이어그램은 몇 가지 예시 updates... 색인을 선택하여 해당하는 results... 색인을 자세히 설명합니다.

좀 더 공식적으로 index_space(updates[0])의 모든 update_index에 대해 다음을 실행합니다.

• 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는 다음과 같이 정의됩니다.
  • scatter_indices[si0, ..., :, ..., siN], 여기서 si는 update_scatter_index의 개별 요소이고 :는 index_vector_dim < rank(scatter_indices)인 경우 index_vector_dim 색인에 삽입됩니다.
  • 그 밖의 경우에는 [scatter_indices[update_scatter_index]]입니다.
• axes(inputs[0])의 d_input:
  • d_input = scatter_dims_to_operand_dims[d_start]인 경우 full_start_index[d_input] = start_index[d_start]입니다.
  • 그 밖의 경우에는 full_start_index[d_input] = 0입니다.
• update_window_index = update_index[update_window_dims...].
• full_window_index = [wi0, ..., 0, ..., wiN], 여기서 wi는 update_window_index의 개별 요소이고 0는 inserted_window_dims의 색인에 삽입됩니다.
• result_index = full_start_index + full_window_index.

각 항목의 의미는 results = exec(schedule, inputs)입니다. 각 항목의 의미는 다음과 같습니다.

• schedule는 index_space(updates[0])의 구현 정의 순열입니다.
• exec([update_index, ...], results) = exec([...], updated_results) 각 항목의 의미는 다음과 같습니다.
  • result_index이 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는 results...[result_index]가 updated_values...로 설정된 results의 사본입니다.
  • 그렇지 않은 경우 다음 단계를 따릅니다.
  • updated_results = results.
• exec([], results) = results.

indices_are_sorted이 true이면 구현에서는 scatter_indices이 scatter_dims_to_operand_dims를 기준으로 정렬되었다고 가정할 수 있습니다. 그렇지 않으면 동작이 정의되지 않습니다. 좀 더 공식적으로 indices(result)의 모든 i1 < i2에 대해 full_start_index(i1) <= full_start_index(i2)입니다.

unique_indices가 true이면 구현은 분산되는 모든 result_index 색인이 고유한 것이라고 가정할 수 있습니다. unique_indices가 true이지만 분산되는 색인이 고유하지 않으면 동작이 정의되지 않습니다.

입력

라벨	이름	유형	제약 조건
(I1)	inputs	가변 텐서 또는 텐서당 양자화 텐서	(C1), (C2), (C4-C6), (C10), (C13), (C15-C16)
(I2)	scatter_indices	정수 유형의 텐서	(C4), (C11), (C14)
(I3)	updates	가변 텐서 또는 텐서당 양자화 텐서	(C3-C6), (C8)
(I4)	update_window_dims	si64 유형의 1차원 텐서 상수	(C2), (C4), (C7), (C8)
(I5)	inserted_window_dims	si64 유형의 1차원 텐서 상수	(C2), (C4), (C9), (C10)
(I6)	scatter_dims_to_operand_dims	si64 유형의 1차원 텐서 상수	(C11-C13)
(I7)	index_vector_dim	si64 유형의 상수	(C4), (C11), (C14)
(I8)	indices_are_sorted	i1 유형의 상수
(I9)	unique_indices	i1 유형의 상수
(I10)	update_computation	함수	(C15)

출력

이름	유형	제약 조건
results	가변 텐서 또는 텐서당 양자화 텐서	(C15-C17)

제약 조건

• (C1) same(shape(inputs...)).
• (C2) rank(inputs[0]) = size(update_window_dims) + size(inserted_window_dims).
• (C3) same(shape(updates...)).
• (C4) shape(updates[0]) = combine(update_scatter_dim_sizes, update_window_dim_sizes) 각 항목의 의미는 다음과 같습니다.
  • update_scatter_dim_sizes = shape(scatter_indices). 단, index_vector_dim에 해당하는 scatter_indices의 크기 크기는 포함되어 있지 않습니다.
  • update_window_dim_sizes <= shape(inputs[0])(inserted_window_dims에 해당하는 inputs[0]의 크기 크기는 포함되지 않음)
  • combine은 update_scatter_dims에 해당하는 축에 update_scatter_dim_sizes를 배치하고 update_window_dims에 해당하는 축에 update_window_dim_sizes를 배치합니다.
• (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(inserted_window_dims) and is_sorted(update_window_dims).
• (C10) 0 <= inserted_window_dims < rank(inputs[0]).
• (C11) size(scatter_dims_to_operand_dims) = index_vector_dim < rank(scatter_indices) ? dim(scatter_indices, index_vector_dim) : 1.
• (C12) is_unique(scatter_dims_to_operand_dims).
• (C13) 0 <= scatter_dims_to_operand_dims < rank(inputs[0]).
• (C14) 0 <= index_vector_dim <= rank(scatter_indices).
• (C15) update_computation에는 is_promotable(element_type(inputs[i]), Ei)인 (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>) 유형이 있습니다.
• (C16) shape(inputs...) = shape(results...).
• (C17) [0,N)의 모든 i에 대해 element_type(results[i]) = Ei.

예

// %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]]
//         ]
// %scatter_indices: [[[0, 2], [1, 0], [2, 1]], [[0, 1], [1, 0], [0, 9]]]
// %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 = [2, 3],
    inserted_window_dims = [0],
    scatter_dims_to_operand_dims = [1, 0],
    index_vector_dim = 2>,
  indices_are_sorted = false,
  unique_indices = false
} : (tensor<3x4x2xi64>, tensor<2x3x2xi64>, tensor<2x3x2x2xi64>) -> tensor<3x4x2xi64>
// %result: [
//           [[1, 2], [5, 6], [7, 8], [7, 8]],
//           [[10, 11], [12, 13], [14, 15], [16, 17]],
//           [[18, 19], [20, 21], [21, 22], [23, 24]]
//          ]

예시 더보기

select

시맨틱

result 텐서를 생성합니다. 여기서 각 요소는 pred의 상응하는 요소 값에 따라 on_true 또는 on_false 텐서에서 선택됩니다. 더 공식적으로 result[result_index] = pred_element ? on_true[result_index] : on_false[result_index], 여기서 pred_element = rank(pred) = 0 ? pred[] : pred[result_index]. 양자화 유형의 경우 dequantize_select_quantize(pred, on_true, on_false, type(result))를 실행합니다.

입력

라벨	이름	유형	제약 조건
(I1)	pred	i1 유형의 텐서	(C1)
(I2)	on_true	텐서 또는 텐서당 양자화 텐서	(C1-C2)
(I3)	on_false	텐서 또는 텐서당 양자화 텐서	(C2)

출력

이름	유형	제약 조건
result	텐서 또는 텐서당 양자화 텐서	(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

시맨틱

select를 사용하는 input 텐서의 reduce_window 결과에 따라 scatter를 사용하여 source 텐서의 값을 분산하고 result 텐서를 생성합니다.

다음 다이어그램은 구체적인 예를 사용하여 operand와 source에서 result의 요소가 계산되는 방식을 보여줍니다.

더 공식적으로

• selected_values = reduce_window_without_init(...)를 다음 입력으로 바꿉니다.

  • `inputs = [피연산자].
  • 있는 그대로 사용되는 window_dimensions, window_strides, padding
  • base_dilations = windows_dilations = 1.
  • body는 다음과 같이 정의됩니다.

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

  여기서 E = element_type(operand) 및 reduce_window_without_init는 기본 reduce의 schedule (reduce 참고)에 init 값이 포함되지 않는다는 점을 제외하고 reduce_window과 정확히 동일하게 작동합니다. 현재 해당 창에 값이 없으면 어떻게 되는지가 지정되어 있지 않습니다. (#731)

• result[result_index] = reduce([source_values], [init_value], [0], scatter) 각 항목의 의미는 다음과 같습니다.

  • source_values = [source[source_index] for source_index in source_indices].
  • selected_values[source_index]에 operand_index의 operand 요소가 있는 경우 selected_index(source_index) = operand_index입니다.
  • source_indices = [source_index for source_index in indices(source) if selected_index(source_index) = result_index].

입력

라벨	이름	유형	제약 조건
(I1)	operand	텐서 또는 텐서당 양자화 텐서	(C1-C4), (C6), (C8-C11)
(I2)	source	텐서 또는 텐서당 양자화 텐서	(C1), (C2)
(I3)	init_value	0차원 텐서 또는 텐서당 양자화 텐서	(C3)
(I4)	window_dimensions	si64 유형의 1차원 텐서 상수	(C2), (C4), (C5)
(I5)	window_strides	si64 유형의 1차원 텐서 상수	(C2), (C6), (C7)
(I6)	padding	si64 유형의 2차원 텐서 상수	(C2), (C8)
(I7)	select	함수	(C9)
(I8)	scatter	함수	(C10)

출력

이름	유형	제약 조건
result	텐서 또는 텐서당 양자화 텐서	(C11-C12)

제약 조건

• (C1) element_type(operand) = element_type(source).
• (C2) shape(source) = num_windows 각 항목의 의미는 다음과 같습니다.
  • 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에는 E = element_type(operand)인 (tensor<E>, tensor<E>) -> tensor<i1> 유형이 있습니다.
• (C10) scatter에는 is_promotable(element_type(operand), E)가 있는 (tensor<E>, tensor<E>) -> tensor<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 = dense<[3, 1]> : tensor<2xi64>,
  window_strides = dense<[2, 1]> : tensor<2xi64>,
  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]]

예시 더보기

보내기

시맨틱

inputs를 channel_id 채널에 전송하고 result 토큰을 생성합니다.

is_host_transfer이 true이면 작업은 데이터를 호스트로 전송합니다. 그렇지 않으면 데이터를 다른 기기로 전송합니다. 이것이 의미하는 바는 구현이 정의되어 있습니다. 이 플래그는 channel_type에 제공된 정보를 복제하므로 향후에는 이 중 하나만 유지할 계획입니다. (#666)

입력

라벨	이름	유형	제약 조건
(I1)	inputs	가변 텐서 또는 양자화 텐서
(I2)	token	token
(I3)	channel_id	si64 유형의 상수
(I4)	channel_type	DEVICE_TO_DEVICE 및 DEVICE_TO_HOST의 열거형	(C1)
(I5)	is_host_transfer	i1 유형의 상수	(C1)

출력

이름	유형
result	token

제약 조건

• (C1) channel_type는 다음과 같이 정의됩니다.
  • is_host_transfer = true인 경우 DEVICE_TO_HOST
  • 그 밖의 경우에는 DEVICE_TO_DEVICE입니다.

예

%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

시맨틱

lhs 텐서에서 rhs비트 수만큼 요소별 왼쪽 시프트 연산을 실행하고 result 텐서를 생성합니다.

입력

라벨	이름	유형	제약 조건
(I1)	lhs	정수 유형의 텐서	(C1)
(I2)	rhs	정수 유형의 텐서	(C1)

출력

이름	유형	제약 조건
result	정수 유형의 텐서	(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

시맨틱

lhs 텐서에서 비트 rhs만큼 요소별 오른쪽 시프트 연산을 실행하고 result 텐서를 생성합니다.

입력

라벨	이름	유형	제약 조건
(I1)	lhs	정수 유형의 텐서	(C1)
(I2)	rhs	정수 유형의 텐서	(C1)

출력

이름	유형	제약 조건
result	정수 유형의 텐서	(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

시맨틱

lhs 텐서에서 rhs비트 수만큼 요소별 논리 오른쪽 시프트 연산을 실행하고 result 텐서를 생성합니다.

입력

라벨	이름	유형	제약 조건
(I1)	lhs	정수 유형의 텐서	(C1)
(I2)	rhs	정수 유형의 텐서	(C1)

출력

이름	유형	제약 조건
result	정수 유형의 텐서	(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]

예시 더보기

표지판

시맨틱

operand의 부호를 요소별로 반환하고 result 텐서를 생성합니다. 좀 더 공식적으로 각 요소 x에 대해 다음과 같이 Python 구문을 사용하여 의미 체계를 표현할 수 있습니다.

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

양자화 유형의 경우 dequantize_op_quantize(sign, operand, type(result))를 실행합니다.

입력

라벨	이름	유형	제약 조건
(I1)	operand	부호 있는 정수, 부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1)

출력

이름	유형	제약 조건
result	부호 있는 정수, 부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(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]

예시 더보기

사인

시맨틱

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

• 부동 소수점 수: IEEE-754의 sin
• 복소수: 복소수 사인
• 양자화 유형: dequantize_op_quantize(sine, 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.sine"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[0.0, 1.0], [0.0, -1.0]]

예시 더보기

slice

시맨틱

정적으로 계산된 시작 색인을 사용하여 operand에서 슬라이스를 추출하고 result 텐서를 생성합니다. start_indices에는 각 차원에 대한 슬라이스의 시작 색인이 포함되며, limit_indices에는 각 차원에 대한 슬라이스의 종료 색인(제외)이 포함되며, strides에는 각 차원의 스트라이드가 포함됩니다.

더 공식적으로 result[result_index] = operand[operand_index]입니다(여기서 operand_index = start_indices + result_index * strides).

입력

라벨	이름	유형	제약 조건
(I1)	operand	텐서 또는 텐서당 양자화 텐서	(C1-C3), (C5)
(I2)	start_indices	si64 유형의 1차원 텐서 상수	(C2), (C3), (C5)
(I3)	limit_indices	si64 유형의 1차원 텐서 상수	(C2), (C3), (C5)
(I4)	strides	si64 유형의 1차원 텐서 상수	(C2), (C4)

출력

이름	유형	제약 조건
result	텐서 또는 텐서당 양자화 텐서	(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 = dense<[1, 2]> : tensor<2xi64>,
  limit_indices = dense<[3, 4]> : tensor<2xi64>,
  strides = dense<1> : tensor<2xi64>
} : (tensor<3x4xi64>) -> tensor<2x2xi64>
// % result: [
//            [1, 1],
//            [1, 1]
//           ]

예시 더보기

sort

시맨틱

comparator에 따라 차원 dimension를 따라 inputs의 1차원 슬라이스를 함께 정렬하고 results를 생성합니다.

다른 연산의 유사한 입력과 달리 dimension는 아래에 설명된 시맨틱을 사용하여 음수 값을 허용합니다. 앞으로는 일관성을 위해 이 방식이 허용되지 않을 수 있습니다. (#1377)

is_stable이 true이면 정렬이 안정적입니다. 즉, 비교 연산자에 의해 같은 것으로 간주되는 요소의 상대적 순서가 유지됩니다. 단일 입력이 있는 경우 comparator(e1, e2) = comparator(e2, e1) = false인 경우에만 두 요소 e1와 e2가 비교 연산자와 같은 것으로 간주됩니다. 이것이 여러 입력으로 일반화되는 방식은 아래 공식을 참고하세요.

좀 더 공식적으로 index_space(results[0])의 모든 result_index에 대해 다음을 실행합니다.

• adjusted_dimension = dimension >= 0 ? dimension : rank(inputs[0]) + dimension.
• result_slice = [ri0, ..., :, ..., riR-1], 여기서 riN는 result_index에 있는 개별 요소이고 :는 adjusted_dimension에 삽입됩니다.
• inputs_together = (inputs[0]..., ..., inputs[N-1]...).
• results_together[result_slice] = sort(inputs_together[result_slice], comparator_together).
• 여기서 sort는 왼쪽 인수가 오른쪽 두 번째 인수보다 작으면 comparator_together가 true를 반환할 것으로 예상하고 내림차순이 아닌 1차원 슬라이스를 정렬합니다.

• 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	가변 텐서 또는 텐서당 양자화 텐서	(C1-C5)
(I2)	dimension	si64 유형의 상수	(C4)
(I3)	is_stable	i1 유형의 상수
(I4)	comparator	함수	(C5)

출력

이름	유형	제약 조건
results	가변 텐서 또는 텐서당 양자화 텐서	(C2), (C3)

제약 조건

• (C1) 0 < size(inputs).
• (C2) type(inputs...) = type(results...).
• (C3) same(shape(inputs...) + shape(results...)).
• (C4) -R <= dimension < R, 여기서 R = rank(inputs[0]).
• (C5) comparator는 (tensor<E1>, tensor<E1>, ..., tensor<EN-1>, tensor<EN-1>) -> tensor<i1> 유형(여기서 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

시맨틱

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

• 부동 소수점 수: IEEE-754의 squareRoot
• 복소수: 복소수 제곱근
• 양자화 유형: dequantize_op_quantize(sqrt, operand, type(result))

입력

라벨	이름	유형	제약 조건
(I1)	operand	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1)

출력

이름	유형	제약 조건
result	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(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]]

예시 더보기

빼기

시맨틱

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

• 정수: 정수 뺄셈.
• 부동 소수점 수: IEEE-754의 subtraction
• 복소수: 복소수 뺄셈.
• 양자화 유형:
  • dequantize_op_quantize(subtract, lhs, rhs, type(result)).

입력

라벨	이름	유형	제약 조건
(I1)	lhs	정수, 부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1)
(I2)	rhs	정수, 부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1)

출력

이름	유형	제약 조건
result	정수, 부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(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

시맨틱

operand 텐서에서 요소별 쌍곡선 탄젠트 연산을 수행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 부동 소수점 수: IEEE-754의 tanh
• 복소수: 복소수 하이퍼볼릭 탄젠트.
• 양자화 유형:
  • dequantize_op_quantize(tanh, operand, type(result)).

입력

라벨	이름	유형	제약 조건
(I1)	operand	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1)

출력

이름	유형	제약 조건
result	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(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]

예시 더보기

전치

시맨틱

permutation를 사용하여 operand 텐서의 차원을 삭제하고 result 텐서를 생성합니다. 더 공식적으로 result[result_index] = operand[operand_index]입니다(여기서 result_index[d] = operand_index[permutation[d]]).

입력

라벨	이름	유형	제약 조건
(I1)	operand	텐서 또는 양자화 텐서	(C1-C4)
(I2)	permutation	si64 유형의 1차원 텐서 상수	(C2-C4)

출력

이름	유형	제약 조건
result	텐서 또는 양자화 텐서	(C1), (C3-C4)

제약 조건

• (C1) element_type(result)는 다음과 같이 주어집니다.
  • !is_per_axis_quantized(operand)인 경우 element_type(operand)입니다.
  • quantization_dimension(operand)와 quantization_dimension(result)가 다를 수 있다는 점을 제외하고 element_type(operand)입니다.
• (C2) permutation는 range(rank(operand))의 순열입니다.
• (C3) shape(result) = dim(operand, permutation...).
• (C4) is_per_axis_quantized(result)인 경우 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 = dense<[2, 1, 0]> : tensor<3xi64>
} : (tensor<2x3x2xi32>) -> tensor<2x3x2xi32>
// %result: [
//           [[1,7], [3,9], [5,11]],
//           [[2,8], [4,10], [6,12]]
//          ]

예시 더보기

triangular_solve

시맨틱

하위 또는 위쪽 삼각 계수 행렬이 있는 일차 방정식 시스템 배치를 풉니다.

좀 더 공식적으로 a와 b이 주어지면 result[i0, ..., iR-3, :, :]는 left_side가 true일 때 또는 x * op(a[i0, ..., iR-3, :, :]) = b[i0, ..., iR-3, :, :]일 때 op(a[i0, ..., iR-3, :, :]) * x = b[i0, ..., iR-3, :, :]에 대한 해법이며, op(a)가 transpose_a에 의해 결정되는 x 변수를 해결합니다. 여기서 op(a)는 다음 중 하나일 수 있습니다.left_sidefalse

• NO_TRANSPOSE: a를 있는 그대로 사용하여 작업을 실행합니다.
• TRANSPOSE: a 전치 작업을 실행합니다.
• ADJOINT: a의 켤레 전치 연산을 실행합니다.

입력 데이터는 lower이 true이거나 a의 위쪽 삼각형인 경우 a의 아래쪽 삼각형에서만 읽힙니다. 출력 데이터는 같은 삼각형에 반환되며, 다른 삼각형의 값은 구현을 통해 정의됩니다.

unit_diagonal이 true이면 구현은 a의 대각선 요소가 1이라고 가정할 수 있습니다. 그렇지 않으면 동작이 정의되지 않습니다.

양자화 유형의 경우 dequantize_op_quantize(lambda x, y: triangular_solve(x, y, left_side, lower, unit_diagonal, transpose_a), a, b, type(result))를 실행합니다.

입력

라벨	이름	유형	제약 조건
(I1)	a	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1-C3)
(I2)	b	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1-C4)
(I3)	left_side	i1 유형의 상수	(C3)
(I4)	lower	i1 유형의 상수
(I5)	unit_diagonal	i1 유형의 상수
(I6)	transpose_a	NO_TRANSPOSE, TRANSPOSE, ADJOINT의 열거형

출력

이름	유형	제약 조건
result	부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서	(C1)

제약 조건

• (C1) baseline_element_type(a) = baseline_element_type(b).
• (C2) 2 <= rank(a) = rank(b) = R.
• (C3) shape(a)와 shape(b) 간의 관계는 다음과 같이 정의됩니다.
  • 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]
//          ]

tuple

시맨틱

val 값에서 result 튜플을 생성합니다.

입력

라벨	이름	유형	제약 조건
(I1)	val	값의 다양한 개수	(C1)

출력

이름	유형	제약 조건
result	tuple	(C1)

제약 조건

• (C1) result에는 Ei = type(val[i])인 tuple<E0, ..., EN-1> 유형이 있습니다.

예

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

시맨틱

operand 유형으로 정의된 양자화 매개변수에 따라 양자화 텐서 operand를 부동 소수점 텐서 result로 요소별 변환을 실행합니다.

더 공식적으로 result = dequantize(operand)입니다.

입력

라벨	이름	유형	제약 조건
(I1)	operand	양자화 텐서	(C1), (C2)

출력

이름	유형	제약 조건
result	부동 소수점 유형의 텐서	(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

시맨틱

result 유형으로 정의된 양자화 매개변수에 따라 부동 소수점 텐서 또는 양자화 텐서 operand를 양자화 텐서 result로 요소별 변환을 실행합니다.

더 공식적으로

• is_float(operand)인 경우:
  • result = quantize(operand, type(result)).
• is_quantized(operand)인 경우:
  • float_result = dequantize(operand).
  • result = quantize(float_result, type(result)).

입력

라벨	이름	유형	제약 조건
(I1)	operand	부동 소수점 또는 양자화 유형의 텐서	(C1), (C2)

출력

이름	유형	제약 조건
result	양자화 텐서	(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]

동안

시맨틱

cond 함수가 true를 출력하는 동안 body 함수를 0회 이상 실행하여 출력을 생성합니다. 보다 공식적으로 시맨틱스는 다음과 같이 Python 구문을 사용하여 표현할 수 있습니다.

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

무한 루프의 동작은 미정입니다. (#383)

입력

라벨	이름	유형	제약 조건
(I1)	operand	텐서, 양자화 텐서 또는 토큰의 가변 수	(C1-C3)
(I2)	cond	함수	(C1)
(I3)	body	함수	(C2)

출력

이름	유형	제약 조건
results	텐서, 양자화 텐서 또는 토큰의 가변 수	(C3)

제약 조건

• (C1) cond는 (T0, ..., TN-1) -> tensor<i1> 유형(여기서 Ti = type(operand[i]))입니다.
• (C2) body는 (T0, ..., TN-1) -> (T0, ..., TN-1) 유형(여기서 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

시맨틱

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

• 불리언: 논리 XOR
• 정수: 비트 XOR

입력

라벨	이름	유형	제약 조건
(I1)	lhs	불리언 또는 정수 유형의 텐서	(C1)
(I2)	rhs	불리언 또는 정수 유형의 텐서	(C1)

출력

이름	유형	제약 조건
result	불리언 또는 정수 유형의 텐서	(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]]

실행

순차적 실행

StableHLO 프로그램은 main 함수에 입력 값을 제공하고 출력 값을 계산하여 실행됩니다. 함수의 출력 값은 상응하는 return 작업에 루팅된 연산의 그래프를 실행하여 계산됩니다.

실행 순서는 Dataflow와 일치한다면(즉, 작업이 사용 전에 실행되는 경우) 구현으로 정의됩니다. StableHLO에서 모든 부수 효과 작업은 하나의 토큰을 소비하고 하나의 토큰을 생성하므로 (after_all를 통해 여러 토큰을 하나의 토큰으로 멀티플렉스할 수 있음) 부작용의 실행 순서도 Dataflow와 일치합니다. 위 예시 프로그램의 가능한 실행 순서는 %0 → %1 → %2 → %3 → %4 → return 또는 %3 → %0 → %1 → %2 → %4 → return입니다.

좀 더 공식적으로 StableHLO 프로세스는 1) StableHLO 프로그램, 2) 작업 상태 (아직 실행되지 않았으며 이미 실행됨), 3) 프로세스에서 작업 중인 중간 값의 조합입니다. 프로세스는 main 함수의 입력 값으로 시작하여 작업 상태와 중간 값을 업데이트하는 작업의 그래프를 진행하며 출력 값으로 작업을 마칩니다. 추가 공식화는 미정입니다. (#484)

동시 실행

StableHLO 프로그램은 num_partitions에 의해 num_replicas의 2D 프로세스 그리드로 구성되며 둘 다 유형이 ui32인 상태로 동시에 실행될 수 있습니다.

StableHLO 프로세스 그리드에서는 StableHLO 프로세스 num_replicas * num_partitions개가 동시에 실행됩니다. 각 프로세스에는 고유한 process_id = (replica_id, partition_id)가 있으며 여기서 replica_ids = range(num_replicas)의 replica_id 및 partition_ids = range(num_partitions)의 partition_id 모두 ui32 유형을 가집니다.

프로세스 그리드의 크기는 모든 프로그램에서 정적으로 알려져 있으며 (앞으로는 StableHLO 프로그램 #650의 명시적인 부분으로 만들 계획임) 프로세스 그리드 내 위치는 모든 프로세스에 대해 정적으로 알려져 있습니다. 각 프로세스는 replica_id 및 partition_id 작업을 통해 프로세스 그리드 내의 위치에 액세스할 수 있습니다.

프로세스 그리드 내에서는 프로그램이 모두 같을 수 있고('단일 프로그램, 다중 데이터' 스타일에서) 모두 다를 수도 있고('다중 프로그램, 다중 데이터' 스타일에서) 모두 다를 수도 있습니다. 앞으로 GSPMD를 포함하여 병렬 StableHLO 프로그램을 정의하는 다른 관용구에 대한 지원을 도입할 계획입니다. (#619)

프로세스 그리드 내에서 프로세스는 대부분 서로 독립적입니다. 별도의 작업 상태와 별도의 입력/중간/출력 값을 가지며 아래에 설명된 소수의 집합 작업을 제외하고 대부분의 작업은 프로세스 간에 개별적으로 실행됩니다.

대부분의 작업 실행은 동일한 프로세스의 값만 사용하기 때문에 일반적으로 이러한 값을 이름으로 참조하는 것이 명확합니다. 그러나 집합 작업의 의미 체계를 설명하는 것은 충분하지 않으며 특정 프로세스 내의 name 값을 참조하도록 name@process_id 표기법을 사용하기도 합니다. (이러한 관점에서 정규화되지 않은 name은 name@(replica_id(), partition_id())의 약칭으로 볼 수 있습니다.)

프로세스 간 실행 순서는 구현을 통해 정의됩니다. 단, 아래에 설명된 대로 지점 간 통신 및 집합 운영에 의해 도입된 동기화는 예외입니다.

지점 간 통신

StableHLO 프로세스는 StableHLO 채널을 통해 서로 통신할 수 있습니다. 채널은 si64 유형의 양수 ID로 표시됩니다. 다양한 작업을 통해 채널에 값을 전송하고 채널에서 수신할 수 있습니다.

이러한 채널 ID의 출처, 프로세스에서 채널을 인식하는 방법, 프로그램에서 발생하는 동기화 종류와 같은 추가적인 공식화는 미정입니다. (#484)

스트리밍 통신

모든 StableHLO 프로세스는 두 가지 스트리밍 인터페이스에 액세스할 수 있습니다.

• 읽을 수 있는 인피드.
• 작성할 수 있는 아웃피드.

프로세스 간에 통신하여 양쪽 끝에 프로세스가 있는 채널과 달리 인피드 및 아웃피드는 다른 종단 구현이 정의되어 있습니다.

스트리밍 통신이 실행 순서에 미치는 영향 및 이로 인해 도입되는 동기화의 종류와 같은 추가 공식화는 미정입니다. (#484)

집단 운영

StableHLO에는 all_gather, all_reduce, all_to_all, collective_broadcast, collective_permute, reduce_scatter의 6가지 집단 작업이 있습니다. 이러한 모든 작업은 StableHLO 프로세스 그리드의 프로세스를 StableHLO 프로세스 그룹으로 분할하고 다른 프로세스 그룹과 독립적으로 각 프로세스 그룹 내에서 공동 계산을 실행합니다.

각 프로세스 그룹 내에서 집합 작업은 동기화 배리어를 유발할 수 있습니다. 추가 공식화, 예를 들어 이 동기화가 정확히 언제 발생하는지, 프로세스가 이 장벽에 정확히 어떻게 도달하는지, 그렇지 않으면 어떻게 되는지에 관한 자세한 내용은 미정입니다. (#484)

프로세스 그룹에 파티션 간 통신이 포함된 경우(예: 파티션 ID가 다른 프로세스 그룹에 프로세스가 있는 경우) 집합 작업의 실행에는 채널이 필요하며 집합 작업은 si64 유형의 양수 channel_id를 제공해야 합니다. 교차 복제본 통신에는 채널이 필요하지 않습니다.

집합 오퍼레이션에서 수행되는 계산은 개별 오퍼레이션마다 다르며, 위의 개별 오퍼레이션 섹션에 설명되어 있습니다. 그러나 프로세스 그리드가 프로세스 그룹으로 분할되는 전략은 이러한 작업 간에 공유되며 이 섹션에 설명되어 있습니다. 좀 더 공식적으로 StableHLO는 다음 4가지 전략을 지원합니다.

cross_replica

교차 복제본 통신만 각 프로세스 그룹 내에서 발생합니다. 이 전략은 복제본 ID 목록인 replica_groups를 사용하고 partition_ids에 의해 replica_groups의 카티전 프로덕트를 계산합니다. replica_groups는 고유한 요소를 포함하고 모든 replica_ids를 포함해야 합니다. 좀 더 공식적으로 Python 구문을 사용하면 다음과 같습니다.

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

예를 들어 replica_groups = [[0, 1], [2, 3]] 및 num_partitions = 2의 경우 cross_replica는 [[(0, 0), (1, 0)], [(0, 1), (1, 1)], [(2, 0), (3, 0)], [(2, 1), (3, 1)]]를 생성합니다.

cross_partition

파티션 간 통신만 각 프로세스 그룹 내에서 발생합니다. 이 전략은 파티션 ID 목록인 partition_groups를 사용하여 replica_ids에 의해 partition_groups의 카티전 프로덕트를 계산합니다. partition_groups는 고유한 요소를 포함하고 모든 partition_ids을 포함해야 합니다. 보다 공식적으로 Python 구문을 사용하면 다음과 같습니다.

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

예를 들어 partition_groups = [[0, 1]] 및 num_replicas = 4의 경우 cross_partition는 [[(0, 0), (0, 1)], [(1, 0), (1, 1)], [(2, 0), (2, 1)], [(3, 0), (3, 1)]]를 생성합니다.

cross_replica_and_partition

복제본 간 통신과 파티션 간 통신은 모두 각 프로세스 그룹 내에서 발생할 수 있습니다. 이 전략은 복제본 ID 목록인 replica_groups를 사용하고 partition_ids로 각 replica_group의 데카르트 곱을 계산합니다. replica_groups는 고유 요소가 있어야 하며 모든 replica_ids를 포함해야 합니다. 보다 공식적으로 Python 구문을 사용하면 다음과 같습니다.

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

예를 들어 replica_groups = [[0, 1], [2, 3]] 및 num_partitions = 2의 경우 cross_replica_and_partition는 [[(0, 0), (1, 0), (0, 1), (1, 1)], [(2, 0), (3, 0), (2, 1), (3, 1)]]를 생성합니다.

flattened_ids

이 전략은 replica_id * num_partitions + partition_id 형식의 '평면화된' 프로세스 ID 목록인 flattened_id_groups를 사용하여 프로세스 ID로 변환합니다. flattened_id_groups는 고유한 요소를 포함하고 모든 process_ids를 포함해야 합니다. 보다 공식적으로 Python 구문을 사용하면 다음과 같습니다.

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

예를 들어 flattened_id_groups = [[0, 1, 2, 3], [4, 5, 6, 7]], num_replicas = 4, num_partitions = 2의 경우 flattened_ids는 [[(0, 0), (0, 1), (1, 0), (1, 1)], [(2, 0), (2, 1), (3, 0), (3, 1)]]를 생성합니다.

정확성

현재 StableHLO는 숫자 정확성을 보장하지 않지만 향후 변경될 수 있습니다. (#1156)

오류

StableHLO 프로그램은 개별 작업에 관한 광범위한 제약조건을 통해 검증되며, 이러한 제약 조건은 실행 시간 전에 많은 클래스의 오류를 배제합니다. 그러나 오류 조건은 여전히 발생할 수 있습니다(예: 정수 오버플로, 범위를 벗어난 액세스 등을 통해). 명시적으로 언급된 경우를 제외하고 이러한 모든 오류가 구현 정의 동작을 초래하지만 이는 향후 변경될 수 있습니다. (#1157)

이 규칙의 예외로, StableHLO 프로그램의 부동 소수점 예외는 잘 정의된 동작을 보입니다. IEEE-754 표준에 의해 정의된 예외 (잘못된 연산, 0으로 나누기, 오버플로, 언더플로 또는 부정확한 예외)를 초래하는 연산은 기본 결과 (표준에 정의된 대로)를 생성하고 상응하는 상태 플래그를 발생시키지 않고 실행을 계속합니다. 표준의 raiseNoFlag 예외 처리와 유사합니다. 비표준 연산 (예: 복잡한 산술 함수, 특정 초월 함수)의 예외는 구현이 정의되어 있습니다.

Notation

구문을 설명하기 위해 이 문서에서는 수정된 ISO 버전의 EBNF 구문 (ISO/IEC 14977:1996, 위키백과)을 사용하며, 두 가지 수정사항이 있습니다. 1) 규칙이 =가 아닌 ::=를 사용하여 정의됩니다.

2) 연결은 ,보다는 병치를 사용하여 표현됩니다.

'유형', '상수', '작업' 섹션 내에서 의미 체계를 설명하기 위해, Google은 아래에 설명된 대로 배열 연산을 간결하게 표현할 수 있도록 지원되는 확장된 Python 문법을 기반으로 하는 수식을 사용합니다. 이 방법은 작은 코드 스니펫에 적합하지만, 드문 경우지만 더 큰 코드 스니펫이 필요한 경우에는 항상 명시적으로 도입되는 바닐라 Python 구문을 사용합니다.

수식

dot_general 사양의 예를 기반으로 수식의 작동 방식을 살펴보겠습니다. 이 작업의 제약 조건 중 하나는 다음과 같습니다. dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...)

이 수식에 사용되는 이름은 두 가지 소스에서 가져옵니다. 1) 전역 함수(예: dim, 2) 해당하는 프로그램 요소의 멤버 정의(즉, dot_general의 '입력' 섹션에 정의된 lhs, lhs_batching_dimensions, rhs, rhs_batching_dimensions 입력)

위에서 언급했듯이 이 수식의 구문은 Python을 기반으로 하며 일부 간결성 중심의 확장 프로그램이 있습니다. 공식을 이해하기 위해 바닐라 Python 구문으로 변환해 보겠습니다.

A) 이 수식에서는 =를 사용하여 동등성을 나타내므로 Python 구문을 얻기 위한 첫 번째 단계는 =를 다음과 같이 ==로 바꾸는 것입니다. dim(lhs, lhs_batching_dimensions...) == dim(rhs, rhs_batching_dimensions...)

B) 또한 이 수식은 스칼라 표현식을 텐서 표현식으로 변환하는 줄임표 (...)를 지원합니다. 간단히 말해 f(xs...)는 '텐서 xs의 각 스칼라 x에 대해 스칼라 f(x)를 계산한 다음 이러한 모든 스칼라 결과를 함께 텐서 결과로 반환'하는 것을 의미합니다. 일반적인 Python 구문에서 예시 수식은 [dim(lhs, dim1) for dim1 in lhs_batching_dimensions] == [dim(rhs, dim2) for dim2 in rhs_batching_dimensions]로 변환됩니다.

생략 부호 덕분에 개별 스칼라 수준에서 작업을 방지할 수 있는 경우가 많습니다. 그러나 일부 까다로운 경우에는 gather 사양의 start_indices[bi0, ..., :, ..., biN] 수식에서처럼 하위 수준의 반비공식 문법을 사용할 수도 있습니다. 간결성을 제공하기 위해 사례별로 직관적으로 이해할 수 있기 때문에 이러한 구문을 바닐라 Python으로 변환하는 정확한 형식은 제공하지 않습니다. 일부 특정 수식이 불투명하게 보이는 경우 Google에 알려주시면 이를 개선하도록 노력하겠습니다.

또한 공식에서는 생략 부호를 사용하여 텐서, 텐서 목록 (예: 다양한 수의 텐서에서 발생할 수 있음) 등을 포함한 모든 종류의 목록을 확장합니다. 이는 정확한 공식을 제공하지 않는 또 다른 영역입니다 (예: 목록은 StableHLO 유형 시스템의 일부도 아닌데).

C) 마지막으로 주목할 만한 표기법은 암시적 브로드캐스트입니다. StableHLO opset은 암시적 브로드캐스트를 지원하지 않지만 수식은 간결성을 제공하기 위해 지원합니다. 간단히 말해, 텐서가 예상되는 컨텍스트에서 스칼라가 사용되면 스칼라는 예상 형태로 브로드캐스트됩니다.

dot_general 예시를 계속 진행하려면 또 다른 제약조건은 0 <= lhs_batching_dimensions < rank(lhs)입니다. dot_general 사양에 정의된 대로 lhs_batching_dimensions은 텐서이지만 0 및 rank(lhs)는 모두 스칼라입니다. 암시적 브로드캐스트를 적용한 후에는 수식이 [0, ..., 0] <= lhs_batching_dimensions < [rank(lhs), ..., rank(lhs)]가 됩니다.

특정 dot_general 연산에 적용될 때 이 수식은 불리언 텐서로 평가됩니다. 수식이 제약 조건으로 사용되는 경우 수식이 true로 평가되거나 true 요소만 있는 텐서로 평가되면 제약 조건이 유지됩니다.

이름

수식에서 어휘 범위에는 1) 전역 함수, 2) 멤버 정의,

3) 지역적 정의 전역 함수 목록은 아래와 같습니다. 요소 정의 목록은 표기법이 적용되는 프로그램 요소에 따라 다릅니다.

• 작업의 경우 멤버 정의에는 '입력' 및 '출력' 섹션에 도입된 이름이 포함됩니다.
• 그 밖의 모든 경우의 경우, 멤버 정의에는 해당하는 EBNF 비단말기 이름을 따서 명명된 프로그램 요소의 구조적 부분이 포함됩니다. 대부분의 경우, 이러한 구조적 부분의 이름은 비터미널의 이름을 스네이크 표기법으로 변환하여 얻지만 (예: IntegerLiteral => integer_literal), 프로세스에서 이름이 축약되는 경우도 있습니다 (예: QuantizationStorageType => storage_type). 이때 이름은 연산별 출력의 '입력' / '출력' 섹션과 명시적으로 비슷합니다.
• 또한 멤버 정의에는 상응하는 프로그램 요소를 참조하는 self가 항상 포함됩니다.

값

수식은 평가될 때 다음 유형의 값을 사용합니다. 1) Value (실제 값, 예: dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>, 항상 유형을 알 수 있음), 2) Placeholder (미래 값 예: lhs, rhs 또는 result, 실제 값은 아직 알려지지 않고 유형만 알려짐), 3) Type('유형' 섹션에 정의된 유형), 4) Function('유형' 섹션에 정의된 함수)

컨텍스트에 따라 이름은 다른 값을 참조할 수 있습니다. 더 구체적으로 말하면, 연산의 '시맨틱' 섹션 (및 다른 프로그램 요소와 동등한 항목)은 런타임 로직을 정의하므로 모든 입력을 Value로 사용할 수 있습니다. 반면에 작업 (및 이와 동등한 항목)의 'Constraints' 섹션은 '컴파일 시간' 로직, 즉 일반적으로 런타임 전에 실행되는 로직을 정의하므로 상수 입력만 Value로 사용할 수 있고 다른 입력은 Placeholder로만 사용할 수 있습니다.

이름	'시맨틱'에서	'Constraints'에서
전역 함수	Function	Function
상수 입력	Value	Value
상수가 아닌 입력	Value	Placeholder
출력	Value	Placeholder
로컬 정의	정의에 따라 다름	정의에 따라 다름

transpose 작업의 예시를 살펴보겠습니다.

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

이 연산의 경우 permutation는 상수이므로 시맨틱스와 제약 조건 모두에서 Value로 사용할 수 있습니다. 반면에 operand 및 result는 시맨틱에서는 Value로 사용할 수 있지만 제약 조건에서는 Placeholder로만 사용할 수 있습니다.

함수

유형 생성

유형을 생성하는 데 사용할 수 있는 함수가 없습니다. 대신 유형 문법이 일반적으로 더 간결하기 때문에 직접 사용합니다. 예를 들어 function_type( [tensor_type([], E), tensor_type([], E)], [tensor_type([], E)])가 아닌 (tensor<E>, tensor<E>) -> (tensor<E>)를 사용합니다.

유형에 대한 함수

• element_type는 텐서 유형과 양자화 텐서 유형에 정의되며 각각 해당하는 TensorType 또는 QuantizedTensorType의 TensorElementType 또는 QuantizedTensorElementType 부분을 반환합니다.

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_quantized(x) and quantization_dimension(x) is not None의 바로가기입니다.

• is_per_tensor_quantized(x: Value | Placeholder | Type) -> Value는 is_quantized(x) and quantization_dimension(x) is None의 바로가기입니다.

• is_promotable(x: Type, y: Type) -> bool는 x 유형을 y 유형으로 승격할 수 있는지 확인합니다. x 및 y가 QuantizedTensorElementType이면 프로모션이 storage_type에만 적용됩니다. 이 특정 버전의 프로모션은 현재 축소 계산의 컨텍스트에서 사용되고 있습니다 (자세한 내용은 RFC 참조).

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_quantized_tensor_element_type(x)의 바로가기입니다.

• is_type_name(x: Value | Placeholder | Type) -> Value. 모든 유형에 사용할 수 있습니다. 예를 들어 x이 FloatType이면 is_float(x)는 true를 반환합니다. x가 값 또는 자리표시자이면 이 함수는 is_type_name(type(x))의 단축키입니다.

• max_value(x: Type) -> Value는 TensorElementType의 최댓값을 반환합니다. x이 TensorElementType가 아니면 None을 반환합니다.

• min_value(x: Type) -> Value는 TensorElementType의 가능한 최솟값을 반환합니다. x이 TensorElementType가 아니면 None을 반환합니다.

• member_name(x: Value | Placeholder | Type) -> Any. 모든 유형의 모든 구성원 정의 member_name에 사용할 수 있습니다. 예를 들어 tensor_element_type(x)는 상응하는 TensorType의 TensorElementType 부분을 반환합니다. x가 값 또는 자리표시자이면 이 함수는 member_name(type(x))의 단축키입니다. x가 적절한 구성원이 있는 유형이나 이러한 유형의 값 또는 자리표시자가 아닌 경우 None가 반환됩니다.

가치 구성

• operation_name(*xs: Value | Type) -> Value. 모든 작업에 사용할 수 있습니다. 예를 들어 add(lhs, rhs)는 텐서 값 2개 lhs 및 rhs를 취하고 이러한 입력으로 add 작업을 평가한 출력을 반환합니다. broadcast_in_dim와 같은 일부 연산의 경우 출력 유형이 '부하 방위'입니다. 즉, 연산을 평가하는 데 필요합니다. 이 경우 함수는 이러한 유형을 인수로 사용합니다.

값 함수

• 모든 Python의 연산자와 함수를 사용할 수 있습니다. 예를 들어 Python의 구독 및 슬라이싱 표기법을 모두 텐서, 양자화 텐서, 튜플로 색인을 생성할 수 있습니다.

• to_destination_type(x: Value, destination_type: Type) -> Value는 텐서에 정의되며, 다음과 같이 type(x) 및 destination_type에 따라 x의 변환된 값을 반환합니다.

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)

convert, uniform_quantize, uniform_dequantize 작업의 병합에 관한 내용은 초기부터 논의됩니다. (#1576) 병합 후에는 위의 함수가 필요하지 않으며 대신 convert에 작업 이름을 사용할 수 있습니다.

• is_nan(x: Value) -> Value는 텐서에 정의되었으며 x의 모든 요소가 NaN이면 true을 반환하고 그렇지 않으면 false를 반환합니다. x가 텐서가 아닌 경우 None을 반환합니다.

• is_sorted(x: Value) -> Value는 텐서에 정의되며, x의 요소가 색인의 사전 오름차순을 기준으로 오름차순으로 정렬되면 true를 반환하고 그렇지 않으면 false를 반환합니다. x가 텐서가 아닌 경우 None을 반환합니다.

• is_unique(x: Value) -> Value는 텐서에 정의되며 x에 중복 요소가 없으면 true, 중복 요소가 없으면 false를 반환합니다. x가 텐서가 아닌 경우 None을 반환합니다.

• member_name(x: Value) -> Any는 모든 값의 모든 멤버 정의 member_name에 대해 정의됩니다. 예를 들어 real_part(x)는 상응하는 ComplexConstant의 RealPart 부분을 반환합니다. x가 적절한 구성원을 가진 값이 아닌 경우 None이 반환됩니다.

• same(x: Value) -> Value는 텐서에 정의되었으며 x의 요소가 모두 서로 같으면 true를 반환하고 그렇지 않으면 false를 반환합니다. 텐서에 요소가 없으면 '모두 서로 같음'으로 계산됩니다. 즉, 함수는 true를 반환합니다. x가 텐서가 아닌 경우 None을 반환합니다.

• split(x: Value, num_results: Value, axis: Value) -> Value는 텐서에 정의되며 axis 축을 따라 x의 num_results 슬라이스를 반환합니다. x가 텐서 또는 dim(x, axis) % num_results != 0가 아닌 경우 None을 반환합니다.

셰이프 계산

• axes(x: Value | Placeholder | Type) -> Value는 range(rank(x))의 바로가기입니다.

• dim(x: Value | Placeholder | Type, axis: Value) -> Value는 shape(x)[axis]의 바로가기입니다.

• dims(x: Value | Placeholder | Type, axes: List) -> List는 list(map(lambda axis: dim(x, axis), axes))의 바로가기입니다.

• index_space(x: Value | Placeholder | Type) -> Value는 텐서에서 정의되며 사전순으로 정렬된 해당 TensorType에 대해 size(x) 색인을 반환합니다(예: [0, ..., 0], [0, ..., 1], ..., shape(x) - 1). x가 텐서 유형, 양자화 텐서 유형, 이러한 유형 중 하나의 값이나 자리표시자가 아닌 경우 None이 반환됩니다.

• rank(x: Value | Placeholder | Type) -> Value는 size(shape(x))의 바로가기입니다.

• shape(x: Value | Placeholder | Type) -> Value는 member_name를 통해 '유형에 대한 함수' 섹션에서 정의됩니다.

• size(x: Value | Placeholder | Type) -> Value는 reduce(lambda x, y: x * y, shape(x))의 바로가기입니다.

양자화 계산

• def baseline_element_type(x: Value | Placeholder | Type) -> Type는 element_type(baseline_type(x))의 바로가기입니다.

• baseline_type는 텐서 유형과 양자화 텐서 유형에 정의되며, '기준', 즉 형태는 동일하지만 요소 유형의 양자화 매개변수가 기본값으로 재설정되는 유형인 '기준'으로 변환합니다. 이는 텐서 유형과 양자화 텐서 유형을 균일하게 비교하는 데 사용되는 편리한 방법으로, 매우 자주 필요합니다. 양자화 유형의 경우 이를 통해 양자화 매개변수를 무시하는 유형을 비교할 수 있습니다. 즉, shape, storage_type, expressed_type, storage_min, storage_max, quantization_dimension (축별 양자화 유형)는 모두 일치해야 하지만 scales와 zero points는 다를 수 있습니다.

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는 양자화 텐서 유형에 정의되며 부동 소수점 텐서 유형으로 변환합니다. 이는 양자화된 요소 유형과 연결된 0점과 배율을 사용하여 저장소 유형의 정수 값을 나타내는 양자화된 요소를 표현된 유형의 상응하는 부동 소수점 값으로 변환함으로써 발생합니다.

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는 부동 소수점 텐서 유형에 정의되며 양자화 텐서 유형으로 변환합니다. 이는 양자화된 요소 유형과 연결된 0점과 배율을 사용하여 표현된 유형의 부동 소수점 값을 저장소 유형의 상응하는 정수 값으로 변환함으로써 발생합니다.

def quantize(x: Value, type: Type) -> Value:
  assert is_float(x) and is_quantized(type)
  x_expressed_rounded = round_nearest_even(x / compute_scales(type, type(x)))
  x_storage_rounded = convert(x_expressed_rounded, storage_type(type))
  x_storage_add = x_storage_rounded + compute_zero_points(type, type(x_storage_rounded))
  x_storage = clamp(storage_min(type), x_storage_add, storage_max(type))
  return bitcast_convert(x_storage, type)

• dequantize_op_quantize는 양자화 텐서에서 요소별 계산을 지정하는 데 사용됩니다. 비양자화, 즉 양자화 요소를 표현된 유형으로 변환한 다음 연산을 실행한 다음 양자화합니다. 즉, 결과를 다시 스토리지 유형으로 변환합니다. 현재 이 함수는 텐서당 양자화에만 작동합니다. 축당 양자화 작업이 진행 중입니다. (#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)

그리드 계산

• cross_partition(replica_groups: Value) -> Value. 위의 'cross_Replica' 섹션을 참조하세요.

• cross_replica(replica_groups: Value) -> Value. 위의 'cross_Replica' 섹션을 참조하세요.

• cross_replica_and_partition(replica_groups: Value) -> Value. 위의 'cross_Replica_and_partition' 섹션을 참조하세요.

• flattened_ids(replica_groups: Value) -> Value. 위의 'flattened_ids' 섹션을 참고하세요.