StableHLO Specification

StableHLO is an operation set for high-level operations (HLO) in machine learning (ML) models. StableHLO works as a portability layer between different ML frameworks and ML compilers: ML frameworks that produce StableHLO programs are compatible with ML compilers that consume StableHLO programs.

Our goal is to simplify and accelerate ML development by creating more interoperability between various ML frameworks (such as TensorFlow, JAX and PyTorch) and ML compilers (such as XLA and IREE). Towards that end, this document provides a specification for the StableHLO programming language.

This specification contains three major sections. First, the Programs section describes the structure of StableHLO programs which consist of StableHLO functions which themselves consist of StableHLO ops. Within that structure, the Ops section specifies the semantics of individual ops. The Execution section provides semantics for all these ops executing together within a program. Finally, the Notation section discusses the notation used throughout the specification.

To view the spec from a previous release of StableHLO, open the repo at the tagged release of interest. For example, the StableHLO v0.19.0 Spec. To view changes that occurred at each minor version bump of StableHLO, refer to the version log in VhloDialect.td.

Programs

Program ::= {Func}

StableHLO programs consist of an arbitrary number of StableHLO functions. Below is an example program with a function @main which has 3 inputs (%image, %weights and %bias) and 1 output. The body of the function has 6 ops.

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

Functions

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

StableHLO functions (which are also called named functions) have an identifier, inputs/outputs and a body. In the future, we are planning to introduce additional metadata for functions to achieve better compatibility with HLO (#425, #626, #740, #744).

Identifiers

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

StableHLO identifiers are similar to identifiers in many programming languages, with two peculiarities: 1) all identifiers have sigils which distinguish different kinds of identifiers, 2) value identifiers can be completely numeric to simplify generation of StableHLO programs.

Types

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

StableHLO types are categorized into value types (which are also called first-class types) which represent StableHLO values and non-value types which describe other program elements. StableHLO types are similar to types in many programming languages, with the main peculiarity being StableHLO's domain-specific nature which results in some unusual outcomes (e.g. scalar types are not value types).

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

Tensor types represent tensors, i.e. multidimensional arrays. They have a shape and an element type, where a shape represents non-negative or unknown dimension sizes in the ascending order of the corresponding dimensions (which are also called axes) numbered from 0 to R-1. The number of dimensions R is called rank. For example, tensor<2x3xf32> is a tensor type with shape 2x3 and element type f32. It has two dimensions (or, in other words, two axes) - 0th dimension and 1st dimension - whose sizes are 2 and 3. Its rank is 2.

Shapes can be partially or completely unknown (dynamic), e.g. tensor<?x2xf64> is partially unknown and tensor<?x?xf64> is completely unknown. Dynamic dimension sizes are represented using a ?. Shapes cannot be unranked.

In the future, we are planning to explore extending tensor types beyond dimension sizes and element types, for example, to include layouts (#629) and sparsity (#1078).

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

Quantized element types represent integer values of a storage type in the range from storage_min to storage_max (inclusive) that correspond to floating-point values of an expressed type. For a given integer value i, the corresponding floating-point value f can be computed as f = (i - zero_point) * scale, where scale and zero_point are called quantization parameters. The storage_min and storage_max are optional in the grammar, but have default values of min_value(storage_type) and max_value(storage_type) respectively. Quantized element types have the following constraints:

• (C1) type(storage_min) = storage_type.
• (C2) type(storage_max) = storage_type.
• (C3) min_value(storage_type) <= storage_min < storage_max <= max_value(storage_type).
• (C4) type(scales...) = expressed_type.
• (C5) 0 < scales.
• (C6) is_finite(scales...).
• (C7) storage_min <= zero_points <= storage_max.
• (C8) type(zero_points...) = storage_type.
• (C9) size(scales) = size(zero_points).
• (C10) If is_empty(quantization_dimension), then size(scales) = 1.
• (C11) 0 <= quantization_dimension.

At the moment, QuantizationScale is a floating-point constant, but there is strong interest in integer-based scales, represented with multipliers and shifts. We are planning to explore this in the near future (#1404).

There is an ongoing discussion on the semantics of QuantizationZeroPoint, including the type, the values and whether there can be just one or potentially multiple zero points in a quantized tensor type. Based on the results of this discussion, the specification around zero points may change in the future (#1405).

Another ongoing discussion involves the semantics of QuantizationStorageMin and QuantizationStorageMax to determine whether any constraints should be imposed on these values and on the values of quantized tensors (#1406).

Finally, we are planning to explore representing unknown scales and zero points, similarly to how we are planning to explore representing unknown dimension sizes (#1407).

Quantized tensor types represent tensors with quantized elements. These tensors are exactly the same as regular tensors, except that their elements have quantized element types, instead of regular element types.

In quantized tensors, quantization can be per-tensor, meaning, having one scale and zero_point for the entire tensor or can be per-axis, meaning, having multiple scales and zero_points, one pair per slice of a particular dimension quantization_dimension. More formally, in a tensor t with per-axis quantization, there are dim(t, quantization_dimension) slices of the quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :], etc. All elements in the ith slice use scales[i] and zero_points[i] as their quantization parameters. Quantized tensor types have the following constraints:

• For per-tensor quantization:
  • No additional constraints.
• For per-axis quantization:
  • (C12) quantization_dimension < rank(self).
  • (C13) dim(self, quantization_dimension) = size(scales).

TokenType ::= 'token'

Token types represent tokens, i.e. opaque values produced and consumed by some operations. Tokens are used for imposing execution order on operations as described in the Execution section.

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

Tuple types represent tuples, i.e. heterogeneous lists. Tuples are a legacy feature which only exists for compatibility with HLO. In HLO, tuples are used to represent variadic inputs and outputs. In StableHLO, variadic inputs and outputs are supported natively, and the only use of tuples in StableHLO is to comprehensively represent HLO ABI where e.g. T, tuple<T> and tuple<tuple<T>> may be materially different depending on a particular implementation. In the future, we are planning to make changes to HLO ABI which may allow us to remove tuple types from StableHLO (#598).

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

Element types represent elements of tensor types. Unlike in many programming languages, these types are not first class in StableHLO. This means that StableHLO programs cannot directly represent values of these types (as a result, it is idiomatic to represent scalar values of type T with 0-dimensional tensor values of type tensor<T>).

• Boolean type represents boolean values true and false.
• Integer types can be either signed (si) or unsigned (ui) and have one of the supported bit widths (2, 4, 8, 16, 32 or 64). Signed siN types represent integer values from -2^(N-1) to 2^(N-1)-1 inclusive, and unsigned uiN types represent integer values from 0 to 2^N-1 inclusive.
• Floating-point types can be one of the following:
  • f8E3M4, f8E4M3 and f8E5M2 8-bit floating point numbers following IEEE-754 conventions.
  • f8E4M3FN and f8E5M2 types corresponding to respectively the E4M3 and E5M2 encodings of the FP8 format described in FP8 Formats for Deep Learning.
  • f8E4M3FNUZ and f8E5M2FNUZ types corresponding to the E4M3 and E5M2 encodings of the FP8 formats described in 8-bit Numerical Formats for Deep Neural Networks.
  • f8E4M3B11FNUZ type corresponding to the E4M3 encoding of the FP8 formats described in Hybrid 8-bit Floating Point (HFP8) Training and Inference for Deep Neural Networks.
  • bf16 type corresponding to the bfloat16 format described in BFloat16: The secret to high performance on Cloud TPUs.
  • f16, f32 and f64 types corresponding to respectively binary16 ("half precision"), binary32 ("single precision") and binary64 ("double precision") formats described in the IEEE 754 standard.
  • tf32 type corresponds to the TensorFloat32 format and has limited support in StableHLO.
  • f4E2M1FN, f6E2M3FN, f6E3M2FN and f8E8M0FNU MX (microscaling) types described in OCP Microscaling Formats Specification.
• Complex types represent complex values that have a real part and an imaginary part of the same element type. Supported complex types are complex<f32> (both parts are of type f32) and complex<f64> (both parts are of type f64).

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

Function types represent both named and anonymous functions. They have input types (the list of types on the left-hand side of ->) and output types (the list of types on the right-hand side of ->). In many programming languages, function types are first class, but not in StableHLO.

StringType ::= 'string'

String type represents sequences of bytes. Unlike in many programming languages, string type is not first class in StableHLO and is only used to specify static metadata for program elements.

Operations

StableHLO operations (which are also called ops) represent a closed set of high-level operations in machine learning models. As discussed above, StableHLO syntax is heavily inspired by MLIR, which is not necessarily the most ergonomic alternative, but is arguably the best fit for StableHLO's goal of creating more interoperability between ML frameworks and ML compilers.

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

StableHLO operations (which are also called ops) have a name, inputs/outputs and a signature. The name consists of the stablehlo. prefix and a mnemonic which uniquely identifies one of the supported ops. See below for a comprehensive list of all supported ops.

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

Ops consume inputs and produce outputs. Inputs are categorized into input values (computed during execution), input functions (provided statically, because in StableHLO functions are not first-class values) and input attributes (also provided statically). The kind of inputs and outputs consumed and produced by an op depends on its mnemonic. For example, the add op consumes 2 input values and produces 1 output value. In comparison, the select_and_scatter op consumes 3 input values, 2 input functions and 3 input attributes.

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

Input functions (which are also called anonymous functions) are very similar to named functions except that: 1) they don't have an identifier (hence the name "anonymous"), 2) they don't declare output types (output types are inferred from the return op within the function).

The syntax for input functions includes a currently unused part (see the Unused production above) which is there for compatibility with MLIR. In MLIR, there is a more general concept of "regions" which can have multiple "blocks" of ops connected together via jump ops. These blocks have ids which correspond to the Unused production, so that they can be distinguished from each other. StableHLO doesn't have jump ops, so the corresponding part of MLIR syntax is unused (but is still there).

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

Input attributes have a name and a value which is one of the supported constants. They are the primary way to specify static metadata for program elements. For example, the concatenate op uses the attribute dimension to specify the dimension along which its input values are concatenated. Similarly, the slice op uses multiple attributes like start_indices and limit_indices to specify the bounds that are used to slice the input value.

At the moment, StableHLO programs in the wild sometimes contain attributes which are not described in this document. In the future, we are planning to either absorb these attributes into the StableHLO opset or prohibit them from appearing in StableHLO programs. In the meanwhile, here is the list of these attributes:

• layout (#629).
• mhlo.frontend_attributes (#628).
• mhlo.sharding (#619).
• output_operand_aliases (#740).
• Location metadata (#594).

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

Op signature consists of the types of all input values (the list of types on the left-hand side of ->) and the types of all output values (the list of types on the right-hand side of ->). Strictly speaking, input types are redundant, and output types are almost always redundant as well (because for most StableHLO ops, output types can be inferred from inputs). Nonetheless, op signature is deliberately part of StableHLO syntax for compatibility with MLIR.

Below is an example op whose mnemonic is select_and_scatter. It consumes 3 input values (%operand, %source and %init_value), 2 input functions and 3 input attributes (window_dimensions, window_strides and padding). Note how the signature of the op only includes the types of its input values (but not the types of input functions and attributes which are provided inline).

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

Constants

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

StableHLO constants have a literal and a type which together represent a StableHLO value. Generally, the type is part of the constant syntax, except when it's unambiguous (e.g. a boolean constant unambiguously has type i1, whereas an integer constant can have multiple possible types).

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

Boolean constants represent boolean values true and false. Boolean constants have type i1.

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

Integer constants represent integer values via strings that use decimal or hexadecimal notation. Other bases, e.g. binary or octal, are not supported. Integer constants have the following constraints:

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

Floating-point constants represent floating-point values via strings that use decimal or scientific notation. Additionally, hexadecimal notation can be used to directly specify the underlying bits in the floating-point format of the corresponding type. Floating-point constants have the following constraints:

• (C1) If non-hexadecimal notation is used, is_wellformed(float_literal, float_type).
• (C2) If hexadecimal notation is used, size(hexadecimal_digits) = num_bits(float_type) / 4.

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

Complex constants represent complex values using lists of a real part (comes first) and an imaginary part (comes second). For example, (1.0, 0.0) : complex<f32> represents 1.0 + 0.0i, and (0.0, 1.0) : complex<f32> represents 0.0 + 1.0i. The order in which these parts are then stored in memory is implementation-defined. Complex constants have the following constraints:

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

Tensor constants represent tensor values using nested lists specified via NumPy notation. For example, dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> represents a tensor value with the following mapping from indices to elements: {0, 0} => 1, {0, 1} => 2, {0, 2} => 3, {1, 0} => 4, {1, 1} => 5, {1, 2} => 6. The order in which these elements are then stored in memory is implementation-defined. Tensor constants have the following constraints:

• (C1) has_syntax(tensor_literal, element_type(tensor_type)), where:
  • 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)), where:
  • 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:]).
  • otherwise, false.

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

Quantized tensor constants represent quantized tensor values using the same notation as tensor constants, with elements specified as constants of their storage type. Quantized tensor constants have the following constraints:

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

String literals consist of bytes specified using ASCII characters and escape sequences. They are encoding-agnostic, so the interpretation of these bytes is implementation-defined. String literals have type string.

Ops

abs

Semantics

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

• For signed integers: integer modulus.
• For floats: abs from IEEE-754.
• For complex numbers: complex modulus.
• For quantized types: dequantize_op_quantize(abs, operand, type(result)).

Inputs

Outputs

Constraints

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

Examples

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

 More Examples

add

Semantics

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

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

Inputs

Outputs

Constraints

• If the operation uses non-quantized tensors:
  • (C1) type(lhs) = type(rhs) = type(result).
• If the operation uses quantized tensors:
  • (C2) is_quantized(lhs) and is_quantized(rhs) and is_quantized(result).
  • (C3) storage_type(lhs) = storage_type(rhs) = storage_type(result).
  • (C4) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C5) (is_per_axis_quantized(lhs) or is_per_axis_quantized(rhs)) = is_per_axis_quantized(result).
  • (C6) If is_per_axis_quantized(lhs), then quantization_dimension(lhs) = quantization_dimension(result).
  • (C7) If is_per_axis_quantized(rhs), then quantization_dimension(rhs) = quantization_dimension(result).

Examples

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

 More Examples

after_all

Semantics

Ensures that the operations producing the inputs are executed before any operations that depend on result. Execution of this operation does nothing, it only exists to establish data dependencies from result to inputs.

Inputs

Outputs

Examples

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

 More Examples

all_gather

Semantics

Within each process group in the StableHLO process grid, concatenates the values of the operands tensors from each process along all_gather_dim and produces results tensors.

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

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

Afterwards, within each process_group:

• operands...@receiver = [operand@sender for sender in process_group] for all receiver in process_group.
• results...@process = concatenate(operands...@process, all_gather_dim) for all process in process_group.

Inputs

Outputs

Constraints

• (C1) 0 <= all_gather_dim < rank(operands...).
• (C2) is_unique(replica_groups).
• (C3) size(replica_groups) is defined as:
  • num_replicas if cross_replica is used.
  • num_replicas if cross_replica_and_partition is used.
  • num_processes if flattened_ids is used.
• (C4) 0 <= replica_groups < size(replica_groups).
• (C5) If use_global_device_ids = true, then channel_id > 0.
• (C6) type(results...) = type(operands...) except:
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1).

Examples

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

 More Examples

all_reduce

Semantics

Within each process group in the StableHLO process grid, applies a reduction function computation to the values of the operands tensors from each process and produces results tensors.

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

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

Afterwards, within each process_group:

• results...@process[result_index] = exec(schedule) for some binary tree schedule where:
  • exec(node) = computation(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value.
• schedule is an implementation-defined binary tree whose in-order traversal is to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0])).

Inputs

Outputs

Constraints

• (C1) is_unique(replica_groups).
• (C2) size(replica_groups) is defined as:
  • num_replicas if cross_replica is used.
  • num_replicas if cross_replica_and_partition is used.
  • num_processes if flattened_ids is used.
• (C3) 0 <= replica_groups < size(replica_groups).
• (C4) If use_global_device_ids = true, then channel_id > 0.
• (C5) computation has type (tensor<E>, tensor<E>) -> (tensor<E>) where is_promotable(element_type(operand), E).
• (C6) shape(results...) = shape(operands...).
• (C7) element_type(results...) = E.

Examples

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

 More Examples

all_to_all

Semantics

Within each process group in the StableHLO process grid, splits the values of the operands tensors along split_dimension into parts, scatters the split parts between the processes, concatenates the scattered parts along concat_dimension and produces results tensors. The operation splits the StableHLO process grid into process_groups which is defined as follows:

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

Afterwards, within each process_group:

• split_parts...@sender = split(operands...@sender, split_count, split_dimension) for all sender in process_group.
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group] where receiver_index = process_group.index(receiver).
• results...@process = concatenate(scattered_parts...@process, concat_dimension).

Inputs

Outputs

Constraints

• (C1) 0 <= split_dimension < rank(operands...).
• (C2) dim(operands..., split_dimension) % split_count = 0.
• (C3) 0 <= concat_dimension < rank(operands...).
• (C4) 0 < split_count.
• (C5) is_unique(replica_groups).
• (C6) size(replica_groups) is defined as:
  • num_replicas if cross_replica is used.
  • num_partitions if cross_partition is used.
• (C7) 0 <= replica_groups < size(replica_groups).
• (C8) dim(replica_groups, 1) = split_count.
• (C9) type(results...) = type(operands...) except, if split_dimension != concat_dimension:
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count.
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count.

Examples

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

 More Examples

and

Semantics

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

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

Inputs

Outputs

Constraints

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

Examples

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

 More Examples

atan2

Semantics

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

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

Inputs

Outputs

Constraints

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

Examples

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

 More Examples

batch_norm_grad

Semantics

Computes gradients of several inputs of batch_norm_training backpropagating from grad_output, and produces grad_operand, grad_scale and grad_offset tensors. More formally, this operation can be expressed as a decomposition to existing StableHLO operations using Python syntax as follows:

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

For quantized types, performs 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)).

Inputs

Outputs

Constraints

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, mean, variance, grad_output, grad_operand, grad_scale and grad_offset have the same baseline_element_type.
• (C3) operand, grad_output and grad_operand have the same shape.
• (C4) scale, mean, variance, grad_scale and grad_offset have the same shape.
• (C5) size(scale) = dim(operand, feature_index).

Examples

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

Semantics

Normalizes the operand tensor across all dimensions except for the feature_index dimension and produces a result tensor. More formally, this operation can be expressed as a decomposition to existing StableHLO operations using Python syntax as follows:

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)

For quantized types, performs 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)).

Inputs

Outputs

Constraints

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, mean, variance and result have the same 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).

Examples

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

Semantics

Computes mean and variance across all dimensions except for the feature_index dimension and normalizes the operand tensor producing output, batch_mean and batch_var tensors. More formally, this operation can be expressed as a decomposition to existing StableHLO operations using Python syntax as follows:

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

For quantized types, performs 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)).

Inputs

Outputs

Constraints

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, batch_mean, batch_var and output have the same 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).

Examples

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

Semantics

Performs a bitcast operation on operand tensor and produces a result tensor where the bits of the entire operand tensor are reinterpreted using the type of the result tensor.

More formally, given E = element_type(operand), E' = element_type(result), and R = rank(operand):

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

bits returns in-memory representation of a given value, and its behavior is implementation-defined because the exact representation of tensors is implementation-defined, and the exact representation of element types is implementation-defined as well.

Inputs

Outputs

Constraints

• (C1) Given E = is_quantized(operand) ? storage_type(operand) : element_type(operand), E' = is_quantized(result) ? storage_type(result) : element_type(result), and R = rank(operand):
  • If num_bits(E') = num_bits(E), shape(result) = shape(operand).
  • If num_bits(E') < num_bits(E):
  • rank(result) = R + 1.
  • dim(result, i) = dim(operand, i) for all 0 <= i < R.
  • dim(result, R) * num_bits(E') = num_bits(E).
  • If num_bits(E') > num_bits(E):
  • rank(result) = R - 1.
  • dim(result, i) = dim(operand, i) for all 0 <= i < R.
  • dim(operand, R - 1) * num_bits(E) = num_bits(E').
• (C2) If is_complex(operand) or is_complex(result), then is_complex(operand) and is_complex(result).

Examples

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

 More Examples

broadcast_in_dim

Semantics

Expands the dimensions and/or rank of an input tensor by duplicating the data in the operand tensor and produces a result tensor. More formally, result[result_index] = operand[operand_index] where for all d in axes(operand):

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

Inputs

Outputs

Constraints

• (C1) element_type(result) is given by:
  • element_type(operand), if !is_per_axis_quantized(operand).
  • element_type(operand) except that quantization_dimension(operand), scales(operand), and zero_points(operand) may differ from quantization_dimension(result), scales(result), and zero_points(result) resp., otherwise.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) For all d in axes(operand):
  • dim(operand, d) = 1 or
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) If is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • If dim(operand, quantization_dimension(operand)) = 1, then scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))).

Examples

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

 More Examples

case

Semantics

Produces the output from executing exactly one function from branches depending on the value of index. More formally, result = selected_branch() where:

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

Inputs

Outputs

Constraints

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

Examples

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

 More Examples

cbrt

Semantics

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

• For floats: rootn(x, 3) from IEEE-754.
• For complex numbers: complex cubic root.
• For quantized types: dequantize_op_quantize(cbrt, operand, type(result))

Inputs

Outputs

Constraints

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

Examples

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

 More Examples

ceil

Semantics

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

Inputs

Outputs

Constraints

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

Examples

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

 More Examples

cholesky

Semantics

Computes the Cholesky decomposition of a batch of matrices.

More formally, for all i in index_space(result), result[i0, ..., iR-3, :, :] is a Cholesky decomposition of a[i0, ..., iR-3, :, :], in the form of either of a lower-triangular (if lower is true) or upper-triangular (if lower is false) matrix. The output values in the opposite triangle, i.e. the strict upper triangle or strict lower triangle correspondingly, are implementation-defined.

If there exists i where the input matrix is not an Hermitian positive-definite matrix, then the behavior is undefined.

For quantized types, performs dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)).

Inputs

Outputs

Constraints

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

Examples

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

clamp

Semantics

Clamps every element of the operand tensor between a minimum and maximum value and produces a result tensor. More formally, result[result_index] = minimum(maximum(operand[result_index], min_element), max_element), where min_element = rank(min) = 0 ? min[] : min[result_index], max_element = rank(max) = 0 ? max[] : max[result_index]. For quantized types, performs dequantize_op_quantize(clamp, min, operand, max, type(result)).

Imposing an ordering on complex numbers involves surprising semantics, so in the future we are planning to remove support for complex numbers for this operation (#560).

Inputs

Outputs

Constraints

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

Examples

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

 More Examples

collective_broadcast

Semantics

Within each process group in the StableHLO process grid, send the value of the operand tensor from the source process to the target processes and produce a result tensor.

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

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

Afterwards, result@process is given by:

• operand@process_groups[i, 0] if there exists an i such that the process is in process_groups[i].
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) otherwise.

Inputs

Outputs

Constraints

• (C1) is_unique(replica_groups).
• (C2) 0 <= replica_groups < N where N is defined as:
  • num_replicas if cross_replica is used.
  • num_partitions if cross_partition is used.
• (C3) type(result) = type(operand).

Examples

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

Semantics

Within each process group in the StableHLO process grid, sends the value of the operand tensor from the source process to the target process and produces a result tensor.

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

• cross_replica(source_target_pairs) if channel_id <= 0.
• cross_partition(source_target_pairs) if channel_id > 0.

Afterwards, result@process is given by:

• operand@process_groups[i, 0], if there exists an i such that process_groups[i, 1] = process.
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) otherwise.

Inputs

Outputs

Constraints

• (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, where N is defined as:
  • num_replicas if cross_replica is used.
  • num_partitions if cross_partition is used.
• (C5) type(result) = type(operand).

Examples

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

 More Examples

compare

Semantics

Performs element-wise comparison of lhs and rhs tensors according to comparison_direction and compare_type, and produces a result tensor.

The values of comparison_direction and compare_type have the following semantics:

For boolean and integer element types:

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

For floating-point element types with compare_type = FLOAT, the op implements the following IEEE-754 operations:

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

For floating-point element types with compare_type = TOTALORDER, the op uses the combination of totalOrder and compareQuietEqual operations from IEEE-754.

For complex element types, lexicographic comparison of (real, imag) pairs is performed using the provided comparison_direction and compare_type. Imposing an ordering on complex numbers involves surprising semantics, so in the future we are planning to remove support for complex numbers when comparison_direction is GE, GT, LE or LT (#560).

For quantized types. performs dequantize_compare(lhs, rhs, comparison_direction).

Inputs

Outputs

Constraints

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs).
• (C2) shape(lhs) = shape(rhs) = shape(result).
• (C3) compare_type is defined as:
  • SIGNED if is_signed_integer(element_type(lhs)).
  • UNSIGNED if is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)).
  • FLOAT or TOTALORDER if is_float(element_type(lhs)).
  • FLOAT if is_complex(element_type(lhs)).

Examples

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

 More Examples

complex

Semantics

Performs element-wise conversion to a complex value from a pair of real and imaginary values, lhs and rhs, and produces a result tensor.

Inputs

Outputs

Constraints

• (C1) type(lhs) = type(rhs).
• (C2) shape(result) = shape(lhs).
• (C3) element_type(result) has type complex<E> where E = element_type(lhs).

Examples

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

 More Examples

composite

Semantics

Encapsulates an operation made up (composed) of other StableHLO operations, taking inputs and composite_attributes and producing results. The semantics of the op are implemented by the decomposition attribute. The composite op can be replaced with its decomposition without changing program semantics. In cases where inlining the decomposition does not provide the same op semantics, prefer using custom_call.

The version field (defaults to 0) is used to denote when a composite's semantics change.

Inputs

Outputs

Constraints

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

Examples

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

 More Examples

concatenate

Semantics

Concatenates inputs along dimension dimension in the same order as the given arguments and produces a result tensor. More formally, result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1], where:

1. id = d0 + ... + dk-1 + kd.
2. d is equal to dimension, and d0, ... are dth dimension sizes of inputs.

Inputs

Outputs

Constraints

• (C1) same(element_type(inputs...)).
• (C2) same(shape(inputs...)) except for 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]) except for:
  • dim(result, dimension) = dim(inputs[0], dimension) + ....

Examples

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

 More Examples

constant

Semantics

Produces an output tensor from a constant value.

Inputs

Outputs

Constraints

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

Examples

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

 More Examples

convert

Semantics

Performs an element-wise conversion from one element type to another on operand tensor and produces a result tensor.

For boolean-to-any-supported-type conversions, the value false is converted to zero, and the value true is converted to one. For any-supported-type-to-boolean conversions, a zero value is converted to false, and non-zero values are converted to true. See below for how this work for complex types.

For conversions involving integer-to-integer, integer-to-floating-point or floating-point-to-floating-point, if the source value can be exactly represented in the destination type, the result value is that exact representation. Otherwise, the behavior is TBD (#180).

For conversions involving floating-point-to-integer, the fractional part is truncated. If the truncated value cannot be represented in the destination type, the behavior is TBD (#180).

Conversion involving complex-to-complex follow the same behavior of floating-point-to-floating-point conversions for converting real and imaginary parts.

For complex-to-any-other-type and any-other-type-to-complex conversions, the source imaginary value is ignored or the destination imaginary value is zeroed, respectively. The conversion of the real part follows the floating-point conversions.

In principle, this operation could express dequantization (conversion from quantized tensors to regular tensors), quantization (conversion from regular tensors to quantized tensors) and requantization (conversion between quantized tensors), but at the moment we have dedicated operations for that - uniform_dequantize for the first use case and uniform_quantize for the second and the third use cases. In the future, these two ops may be merged into convert (#1576).

Inputs

Outputs

Constraints

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

Examples

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

 More Examples

convolution

Semantics

Computes dot products between windows of lhs and slices of rhs and produces result. The following diagram shows how elements in result are computed from lhs and rhs using a concrete example.

More formally, consider the following reframing of the inputs in terms of lhs in order to be able to express windows of 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).

This reframing uses the following helper functions:

• 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] where j[d] = i[permutation[d]].

If feature_group_count = 1 and batch_group_count = 1, then for all output_spatial_index in index_space(dim(result, output_spatial_dimensions...)), result[result_shape(:, output_spatial_index, :)] = dot_product where:

• 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]). This feature appears to be unused, so in the future we are planning to remove it (#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]).

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

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

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

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

Inputs

Outputs

Constraints

• (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) Given input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]:
  • is_unique(input_dimensions).
  • 0 <= input_dimensions < N.
• (C14) dim(rhs, kernel_input_feature_dimension) = dim(lhs, input_feature_dimension) / feature_group_count.
• (C15) dim(rhs, kernel_output_feature_dimension) % batch_group_count = 0.
• (C16) dim(rhs, kernel_output_feature_dimension) % feature_group_count = 0.
• (C17) size(kernel_spatial_dimensions) = N - 2.
• (C18) Given kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]:
  • is_unique(kernel_dimensions).
  • 0 <= kernel_dimensions < N.
• (C19) size(output_spatial_dimensions) = N - 2.
• (C20) Given output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]:
  • is_unique(output_dimensions).
  • 0 <= output_dimensions < N.
• (C21) 0 < feature_group_count.
• (C22) 0 < batch_group_count.
• (C23) feature_group_count = 1 or batch_group_count = 1.
• (C24) size(precision_config) = 2.
• (C25) dim(result, result_dim) is defined as:
  • dim(lhs, input_batch_dimension) / batch_group_count if result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) if result_dim = output_feature_dimension.
  • num_windows otherwise, where:
  • output_spatial_dimensions[spatial_dim] = result_dim.
  • lhs_dim = input_spatial_dimensions[spatial_dim].
  • rhs_dim = kernel_spatial_dimensions[spatial_dim].
  • dilated_input_shape[lhs_dim] = dim(lhs, lhs_dim) = 0 ? 0 : (dim(lhs, lhs_dim) - 1) * lhs_dilation[spatial_dim] + 1.
  • padded_input_shape[lhs_dim] = padding[spatial_dim, 0] + dilated_input_shape[lhs_dim] + padding[spatial_dim, 1].
  • dilated_window_shape[lhs_dim] = dim(rhs, rhs_dim) = 0 ? 0 : (dim(rhs, rhs_dim) - 1) * rhs_dilation[spatial_dim] + 1.
  • is_empty_window[lhs_dim] = padded_input_shape[lhs_dim] = 0 || dilated_window_shape[lhs_dim] > padded_input_shape[lhs_dim].
  • num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1.
• (C26) rank(result) = N.
• If the operation uses non-quantized tensors:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• If the operation uses quantized tensors:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) If is_per_axis_quantized(rhs), then quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) If is_per_axis_quantized(result), then quantization_dimension(result) = output_feature_dimension.
  • If is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) If is_per_tensor_quantized(rhs), then is_per_tensor_quantized(result).
  • If !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Examples

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

 More Examples

cosine

Semantics

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

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

Inputs

Outputs

Constraints

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

Examples

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

 More Examples

count_leading_zeros

Semantics

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

Inputs

Outputs

Constraints

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

Examples

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

 More Examples

custom_call

Semantics

Encapsulates an implementation-defined operation call_target_name that takes inputs and called_computations and produces results. has_side_effect, backend_config and api_version may be used to provide additional implementation-defined metadata.

At the moment, this operation contains a fairly disorganized collection of metadata which reflects organic evolution of its counterpart operation in the XLA compiler. In the future, we are planning to unify this metadata (#741).

Inputs

Outputs

Examples

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

divide

Semantics

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

• For integers: integer division which produces the algebraic quotient with any fractional part discarded.
• For floats: division from IEEE-754.
• For complex numbers: complex division.
• For quantized types:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)).

Inputs

Outputs

Constraints

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

Examples

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

 More Examples

dot_general

Semantics

Computes dot products between slices of lhs and slices of rhs and produces a result tensor.

More formally, result[result_index] = dot_product, where:

• 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 where size(result_batching_index) = size(lhs_batching_dimensions), size(result_lhs_index) = size(lhs_result_dimensions) and 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)).

For quantized types, performs 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)).

For hybrid quantized types, performs hybrid_dequantize_then_op( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs).

precision_config controls the tradeoff between speed and accuracy for computations on accelerator backends. This can be one of the following (at the moment, the semantics of these enum values is underspecified, but we are planning to address this in #755):

• DEFAULT: Fastest calculation, but least accurate approximation to the original number.
• HIGH: Slower calculation, but more accurate approximation to the original number.
• HIGHEST: Slowest calculation, but most accurate approximation to the original number.

A DotAlgorithm defines the main properties of the algorithm used to implement the dot operation, which also defines the precision. If the algorithm attribute fields are set, then the precision_config must be DEFAULT. DotAlgorithms do not have a default value, as the default parameters are implementation defined. As such, all dot algorithm fields may be set to None to specify an empty dot algorithm, which will instead use the precision_config value.

DotAlgorithm fields include:

• lhs_precision_type and rhs_precision_type, the precisions that the LHS and RHS of the operation are rounded to. Precision types are independent from the storage types of the inputs and the output.
• accumulation_type the precision used for accumulation.
• lhs_component_count, rhs_component_count, and num_primitive_operations apply when we are doing an algorithm which decomposes the LHS and/or RHS into multiple components and does multiple "primitive" dot operations on those values - usually to emulate a higher precision (e.g. Leveraging the bfloat16 Artificial Intelligence Datatype For Higher-Precision Computations: bf16_6x tf32_3x, etc). For algorithms with no decomposition, these values should be set to 1.
• allow_imprecise_accumulation to specify if accumulation in lower precision is permitted for some steps (e.g. CUBLASLT_MATMUL_DESC_FAST_ACCUM).

Example DotAlgorithm attributes:

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


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


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

It is up to the implementations to decide which combinations are supported. In general, it is not guaranteed that each algorithm is supported on each accelerator type by the consumer of the StableHLO. If a given algorithm is not supported, an error should be raised as opposed to falling back to an alternative. StableHLO verification will provide best effort verification, preventing algorithms that are not known to be supported on any hardware.

See xla_data.proto > Algorithm for some supported algorithm values. Ticket #2483 captures the plan to create a centralized doc on supported algorithms by backend.

Inputs

Outputs

Constraints

• (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).
• If the operation uses non-quantized tensors:
  • (C13) element_type(lhs) = element_type(rhs).
• If the operation uses quantized tensors:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C15) zero_points(rhs) = 0.
  • (C16) If is_per_axis_quantized(rhs), then quantization_dimension(rhs) not in rhs_contracting_dimensions.
  • If is_quantized(lhs):
  • (C17) storage_type(lhs) = storage_type(rhs).
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C19) If is_per_tensor_quantized(rhs), then is_per_tensor_quantized(result).
  • If !is_quantized(lhs):
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result).
• If !is_empty_algorithm(lhs_precision_type, rhs_precision_type, accumulation_type, lhs_component_count, rhs_component_count, num_primitive_operations allow_imprecise_accumulation):
  • (C21) precision_config... = DEFAULT.
  • (C22) 0 < lhs_component_count.
  • (C23) 0 < rhs_component_count.
  • (C24) 0 < num_primitive_operations.

Examples

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

 More Examples

dynamic_broadcast_in_dim

Semantics

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

The operation also accepts optional attributes known_expanding_dimensions, known_nonexpanding_dimensions to express static knowledge about the expanding behavior of dimensions. If not specified, all dimensions are assumed to be possibly expanding.

Inputs

Outputs

Constraints

• (C1) element_type(result) is given by:
  • element_type(operand), if !is_per_axis_quantized(operand).
  • element_type(operand) except that quantization_dimension(operand), scales(operand), and zero_points(operand) may differ from quantization_dimension(result), scales(result), and zero_points(result) resp., otherwise.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) For all d in axes(operand):
  • dim(operand, d) = 1 or
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) If is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • If dim(operand, quantization_dimension(operand)) = 1, then scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))).
• (C7) size(output_dimensions) = rank(result).
• (C8) is_unique(known_expanding_dimensions + known_nonexpanding_dimensions).
• (C9) 0 <= known_expanding_dimensions < rank(operand).
• (C10) 0 <= known_nonexpanding_dimensions < rank(operand).

Examples

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

 More Examples

dynamic_conv

Semantics

This operation is functionally identical to convolution op, but the padding is specified dynamically via padding.

Inputs

Outputs

Constraints

• (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) Given input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]:
  • is_unique(input_dimensions).
  • 0 <= input_dimensions < N.
• (C14) dim(rhs, kernel_input_feature_dimension) = dim(lhs, input_feature_dimension) / feature_group_count.
• (C15) dim(rhs, kernel_output_feature_dimension) % batch_group_count = 0.
• (C16) dim(rhs, kernel_output_feature_dimension) % feature_group_count = 0.
• (C17) size(kernel_spatial_dimensions) = N - 2.
• (C18) Given kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]:
  • is_unique(kernel_dimensions).
  • 0 <= kernel_dimensions < N.
• (C19) size(output_spatial_dimensions) = N - 2.
• (C20) Given output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]:
  • is_unique(output_dimensions).
  • 0 <= output_dimensions < N.
• (C21) 0 < feature_group_count.
• (C22) 0 < batch_group_count.
• (C23) feature_group_count = 1 or batch_group_count = 1.
• (C24) size(precision_config) = 2.
• (C25) dim(result, result_dim) is defined as:
  • dim(lhs, input_batch_dimension) / batch_group_count if result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) if result_dim = output_feature_dimension.
  • num_windows otherwise, where:
  • output_spatial_dimensions[spatial_dim] = result_dim.
  • lhs_dim = input_spatial_dimensions[spatial_dim].
  • rhs_dim = kernel_spatial_dimensions[spatial_dim].
  • dilated_input_shape[lhs_dim] = dim(lhs, lhs_dim) = 0 ? 0 : (dim(lhs, lhs_dim) - 1) * lhs_dilation[spatial_dim] + 1.
  • padded_input_shape[lhs_dim] = padding[spatial_dim, 0] + dilated_input_shape[lhs_dim] + padding[spatial_dim, 1].
  • dilated_window_shape[lhs_dim] = dim(rhs, rhs_dim) = 0 ? 0 : (dim(rhs, rhs_dim) - 1) * rhs_dilation[spatial_dim] + 1.
  • is_empty_window[lhs_dim] = padded_input_shape[lhs_dim] = 0 || dilated_window_shape[lhs_dim] > padded_input_shape[lhs_dim].
  • num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1.
• (C26) rank(result) = N.
• If the operation uses non-quantized tensors:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• If the operation uses quantized tensors:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) If is_per_axis_quantized(rhs), then quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) If is_per_axis_quantized(result), then quantization_dimension(result) = output_feature_dimension.
  • If is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) If is_per_tensor_quantized(rhs), then is_per_tensor_quantized(result).
  • If !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Examples

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