blob: f318563c142308bcb8387132beabc3f8d562b7df [file] [view] [edit]
# 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](#programs) section describes the structure of StableHLO programs
which consist of StableHLO functions which themselves consist of StableHLO ops.
Within that structure, the [Ops](#ops) section specifies the semantics of
individual ops. The [Execution](#execution) section provides semantics for all
these ops executing together within a program. Finally, the
[Notation](#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](https://github.com/openxla/stablehlo/tags) of interest.
For example, the [StableHLO v0.19.0 Spec](https://github.com/openxla/stablehlo/blob/v0.19.0/docs/spec.md).
To view changes that occurred at each minor version bump of StableHLO, refer to
the version log in [VhloDialect.td](https://github.com/openxla/stablehlo/blob/main/stablehlo/dialect/VhloDialect.td).
## Programs
```ebnf
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.
```mlir
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
```ebnf
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](https://github.com/openxla/stablehlo/issues/425),
[#626](https://github.com/openxla/stablehlo/issues/626),
[#740](https://github.com/openxla/stablehlo/issues/740),
[#744](https://github.com/openxla/stablehlo/issues/744)).
### Identifiers
```ebnf
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
```ebnf
Type ::= ValueType | NonValueType
ValueType ::= TensorType | QuantizedTensorType | TokenType | TupleType | BufferType | FutureType
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).
```ebnf
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](https://github.com/openxla/stablehlo/issues/629)) and sparsity
([#1078](https://github.com/openxla/stablehlo/issues/1078)).
```ebnf
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
```
| Name | Type | Constraints |
|--------------------------|---------------------------------------------|-----------------------------|
| `storage_type` | integer type | (C1-C3), (C8) |
| `storage_min` | integer constant | (C1), (C3), (C7) |
| `storage_max` | integer constant | (C2), (C3), (C7) |
| `expressed_type` | floating-point type | (C4) |
| `quantization_dimension` | optional integer constant | (C10-C12) |
| `scales` | variadic number of floating-point constants | (C4-C6), (C9), (C10), (C13) |
| `zero_points` | variadic number of integer constants | (C7-C9) |
**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](https://github.com/openxla/stablehlo/issues/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](https://github.com/openxla/stablehlo/issues/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](https://github.com/openxla/stablehlo/issues/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](https://github.com/openxla/stablehlo/issues/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 `i`th 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)`.
```ebnf
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](#execution) section.
```ebnf
TupleType ::= 'tuple' '<' TupleElementTypes '>'
TupleElementTypes ::= [ValueType {',' ValueType}]
```
**Buffer types** represent buffers. For example, in XLA, buffers are
multidimensional arrays with consistent storage. Similar to **tensor types**,
buffer types 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,
`memref<2x3xf32>` is a buffer 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.
Buffers can be allocated using a `custom_call` to `CreateBuffer` or `Pin` and
deallocated via a `custom_call` to `Unpin`. Only `custom_call` ops can read and
write the content inside buffers. See [custom_call](#custom_call) for more
detail.
```ebnf
FutureType ::= 'future' '<' FutureElementTypes '>'
FutureElementTypes ::= TensorType {',' TensorType}
```
**Future types** represent futures, i.e. values produced by an asynchronous
operation that become available at some point in the future.
**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](https://github.com/openxla/stablehlo/issues/598)).
```ebnf
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](https://arxiv.org/abs/2209.05433).
* `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](https://arxiv.org/abs/2206.02915).
* `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](https://proceedings.neurips.cc/paper_files/paper/2019/file/65fc9fb4897a89789352e211ca2d398f-Paper.pdf).
* `bf16` type corresponding to the `bfloat16` format described in
[BFloat16: The secret to high performance on Cloud TPUs](https://cloud.google.com/blog/products/ai-machine-learning/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](https://ieeexplore.ieee.org/document/8766229).
* `tf32` type corresponds to the [TensorFloat32 format](https://blogs.nvidia.com/blog/tensorfloat-32-precision-format/)
and has limited support in StableHLO.
* `f4E2M1FN`, `f6E2M3FN`, `f6E3M2FN` and `f8E8M0FNU` MX (microscaling) types
described in
[OCP Microscaling Formats Specification](https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf).
* **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`).
```ebnf
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.
```ebnf
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.
```ebnf
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.
```ebnf
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.
```ebnf
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).
```ebnf
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](https://github.com/openxla/stablehlo/issues/629)).
* `mhlo.frontend_attributes`
([#628](https://github.com/openxla/stablehlo/issues/628)).
* `mhlo.sharding` ([#619](https://github.com/openxla/stablehlo/issues/619)).
* `output_operand_aliases`
([#740](https://github.com/openxla/stablehlo/issues/740)).
* Location metadata ([#594](https://github.com/openxla/stablehlo/issues/594)).
```ebnf
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).
```mlir
%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
```ebnf
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).
```ebnf
BooleanConstant ::= BooleanLiteral
BooleanLiteral ::= 'true' | 'false'
```
**Boolean constants** represent boolean values `true` and `false`. Boolean
constants have type `i1`.
```ebnf
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)`.
```ebnf
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`.
```ebnf
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))`.
```ebnf
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`.
```ebnf
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))`.
```ebnf
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
| Label | Name | Type | Constraints |
|-------|-----------|------------------------------------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of signed integer, floating-point, or complex type or per-tensor quantized tensor | (C1-C2) |
#### Outputs
| Name | Type | Constraints |
|----------|--------------------------------------------------------------------------------|-------------|
| `result` | tensor of signed integer or floating-point type or per-tensor quantized tensor | (C1-C2) |
#### 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
```mlir
// %operand: [-2, 0, 2]
%result = "stablehlo.abs"(%operand) : (tensor<3xi32>) -> tensor<3xi32>
// %result: [2, 0, 2]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/abs.mlir)
### 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
| Label | Name | Type | Constraints |
|-------|-------|----------------------------|---------------|
| (I1) | `lhs` | tensor or quantized tensor | (C1-C6) |
| (I2) | `rhs` | tensor or quantized tensor | (C1-C5), (C7) |
#### Outputs
| Name | Type | Constraints |
|----------|----------------------------|-------------|
| `result` | tensor or quantized tensor | (C1-C7) |
#### 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
```mlir
// %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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/add.mlir)
### 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
| Label | Name | Type |
|-------|----------|----------------------------|
| (I1) | `inputs` | variadic number of `token` |
#### Outputs
| Name | Type |
|----------|---------|
| `result` | `token` |
#### Examples
```mlir
// %input0: !stablehlo.token
// %input1: !stablehlo.token
%result = "stablehlo.after_all"(%input0, %input1) : (!stablehlo.token, !stablehlo.token) -> !stablehlo.token
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/after_all.mlir)
### 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
| Label | Name | Type | Constraints |
|-------|-------------------------|------------------------------------------------------------------------|-------------|
| (I1) | `operands` | variadic number of tensors or per-tensor quantized tensors | (C1), (C6) |
| (I2) | `all_gather_dim` | constant of type `si64` | (C1), (C6) |
| (I3) | `replica_groups` | 2-dimensional tensor constant of type `si64` or `ReplicaGroupMeshAxes` | (C2-C4) |
| (I4) | `channel_id` | constant of type `si64` | (C5) |
| (I5) | `use_global_device_ids` | constant of type `i1` | (C5) |
#### Outputs
| Name | Type | Constraints |
|-----------|------------------------------------------------------------|-------------|
| `results` | variadic number of tensors or per-tensor quantized tensors | (C6) |
#### 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
```mlir
// 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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/all_gather.mlir)
### 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
| Label | Name | Type | Constraints |
|-------|-------------------------|--------------------------------------------------------------------------------------------|-------------|
| (I1) | `operands` | variadic number of tensors or per-tensor quantized tensors | (C5), (C6) |
| (I2) | `replica_groups` | variadic number of 1-dimensional tensor constants of type `si64` or `ReplicaGroupMeshAxes` | (C1-C3) |
| (I3) | `channel_id` | constant of type `si64` | (C4) |
| (I4) | `use_global_device_ids` | constant of type `i1` | (C4) |
| (I5) | `computation` | function | (C5) |
#### Outputs
| Name | Type | Constraints |
|-----------|------------------------------------------------------------|-------------|
| `results` | variadic number of tensors or per-tensor quantized tensors | (C6-C7) |
#### 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
```mlir
// 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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/all_reduce.mlir)
### all_to_all
#### Semantics
![all_to_all](images/spec/all_to_all.svg)
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
| Label | Name | Type | Constraints |
|-------|--------------------|------------------------------------------------------------------------|------------------------|
| (I1) | `operands` | variadic number of tensors or per-tensor quantized tensors | (C1-C3), (C9) |
| (I2) | `split_dimension` | constant of type `si64` | (C1), (C2), (C9) |
| (I3) | `concat_dimension` | constant of type `si64` | (C3), (C9) |
| (I4) | `split_count` | constant of type `si64` | (C2), (C4), (C8), (C9) |
| (I5) | `replica_groups` | 2-dimensional tensor constant of type `si64` or `ReplicaGroupMeshAxes` | (C5-C8) |
| (I6) | `channel_id` | constant of type `si64` | |
#### Outputs
| Name | Type | Constraints |
|-----------|------------------------------------------------------------|-------------|
| `results` | variadic number of tensors or per-tensor quantized tensors | (C9) |
#### 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
```mlir
// 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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/all_to_all.mlir)
### 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
| Label | Name | Type | Constraints |
|-------|-------|-----------------------------------|-------------|
| (I1) | `lhs` | tensor of boolean or integer type | (C1) |
| (I2) | `rhs` | tensor of boolean or integer type | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-----------------------------------|-------------|
| `result` | tensor of boolean or integer type | (C1) |
#### Constraints
* (C1) `type(lhs) = type(rhs) = type(result)`.
#### Examples
```mlir
// %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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/and.mlir)
### async_done
#### Semantics
Waits for an asynchronous operation to complete and returns its value.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|-----------------|-------------|
| (I1) | `operand` | a tensor future | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------------------------------------------------|-------------|
| `result` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) If `async_done` receives a `future<T>`, then `async_done` returns `T`.
#### Examples
```mlir
%0 = "stablehlo.async_start"(%arg0) ({
%1 = "stablehlo.all_gather"(%arg0) {
all_gather_dim = 1 : i64,
replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7]]> : tensor<2x4xi64>
} : (tensor<8x2xf32>) -> tensor<8x8xf32>
"stablehlo.return"(%1) : (tensor<8x8xf32>) -> ()
}) : (tensor<8x2xf32>) -> !stablehlo.future<tensor<8x8xf32>>
%2 = "stablehlo.async_done"(%0) : (!stablehlo.future<tensor<8x8xf32>>) -> tensor<8x8xf32>
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/atan2.mlir)
### async_start
#### Semantics
Performs a computation asynchronously.
By default, StableHLO operations are performed synchronously. `async_start`
indicates that an operation should instead be performed asynchronously. The
asynchrony is not guaranteed. A backend may still choose to perform the
operation synchronously. You can use operations like `optimization_barrier` to
coerce the scheduling of asynchronous operations to run concurrently with other
operations, but again, this scheduling is not guaranteed.
#### Inputs
| Label | Name | Type | Constraints |
|-------|------------|------------------------------------------------------------|------------------------|
| (I1) | `operands` | variadic number of tensors or per-tensor quantized tensors | (C3) |
| (I2) | `body` | a region with a single operation | (C1), (C2), (C3), (C4) |
#### Outputs
| Name | Type | Constraints |
|----------|-----------------|-------------|
| `result` | a tensor future | (C4) |
#### Constraints
* (C1) `body` must contain a single operation and a return.
* (C2) The operation in `body` can only be a collective (`all_reduce`,
`all_gather`, `all_to_all`, `collective_permute`, `reduce_scatter`,
`collective_broadcast`) or a slice (`slice`, `dynamic_slice`,
`dynamic_update_slice`).
* (C3) The arguments to `async_start` must match the ones passed to the operation
in `body`.
* (C4) If the op in `body` returns type `T`, then `async_start` returns
`future<T>`.
#### Examples
```mlir
%0 = "stablehlo.async_start"(%arg0) ({
%1 = "stablehlo.all_gather"(%arg0) {
all_gather_dim = 1 : i64,
replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7]]> : tensor<2x4xi64>
} : (tensor<8x2xf32>) -> tensor<8x8xf32>
"stablehlo.return"(%1) : (tensor<8x8xf32>) -> ()
}) : (tensor<8x2xf32>) -> !stablehlo.future<tensor<8x8xf32>>
```
### 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
| Label | Name | Type | Constraints |
|-------|-------|-------------------------------------------------------------------------|-------------|
| (I1) | `lhs` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
| (I2) | `rhs` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------------------------------------------------|-------------|
| `result` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(lhs) = baseline_type(rhs) = baseline_type(result)`.
#### Examples
```mlir
// %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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/atan2.mlir)
### 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:
```python
def compute_sum(operand, feature_index):
(sum,) = reduce(
inputs=[operand],
init_values=[constant(0, element_type(operand))],
dimensions=[i for i in range(rank(operand)) if i != feature_index],
body=lambda x, y: add(x, y))
return sum
def compute_mean(operand, feature_index):
sum = compute_sum(operand, feature_index)
divisor = constant(size(operand) / dim(operand, feature_index),
element_type(operand))
divisor_bcast = broadcast_in_dim(divisor, [], type(sum))
return divide(sum, divisor_bcast)
def batch_norm_grad(operand, scale, mean, variance, grad_output, epsilon, feature_index):
# Broadcast inputs to type(operand)
scale_bcast = broadcast_in_dim(scale, [feature_index], type(operand))
mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
variance_bcast = broadcast_in_dim(variance, [feature_index], type(operand))
epsilon_bcast = broadcast_in_dim(constant(epsilon, element_type(operand)), [],
type(operand))
# Perform normalization using the provided `mean` and `variance`
# Intermediate values will be useful for computing gradients
centered_operand = subtract(operand, mean_bcast)
stddev = sqrt(add(variance_bcast, epsilon_bcast))
normalized_operand = divide(centered_operand, stddev)
# Use the implementation from batchnorm_expander.cc in XLA
# Temporary variables have exactly the same names as in the C++ code
elements_per_feature = broadcast_in_dim(
constant(divide(size(operand), dim(operand, feature_index)),
element_type(grad_output)),
[], type(operand))
i1 = multiply(grad_output, elements_per_feature)
i2 = broadcast_in_dim(
compute_sum(grad_output, feature_index), [feature_index], type(operand))
i3 = broadcast_in_dim(
compute_sum(multiply(grad_output, centered_operand), feature_index),
[feature_index], type(operand))
i4 = multiply(i3, centered_operand)
i5 = divide(i4, add(variance_bcast, epsilon_bcast))
i6 = subtract(subtract(i1, i2), i5)
grad_operand =
multiply(divide(divide(scale_bcast, stddev), elements_per_feature), i6)
grad_scale =
compute_sum(multiply(grad_output, normalized_operand), feature_index)
grad_offset = compute_sum(grad_output, feature_index)
return grad_operand, grad_scale, grad_offset
```
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
| Label | Name | Type | Constraints |
|-------|-----------------|---------------------------------------------------------------------|------------------|
| (I1) | `operand` | tensor of floating-point type or per-tensor quantized tensor | (C1-C3), (C5) |
| (I2) | `scale` | 1-dimensional tensor of floating-point or per-tensor quantized type | (C2), (C4), (C5) |
| (I3) | `mean` | 1-dimensional tensor of floating-point or per-tensor quantized type | (C2), (C4) |
| (I4) | `variance` | 1-dimensional tensor of floating-point or per-tensor quantized type | (C2), (C4) |
| (I5) | `grad_output` | tensor of floating-point type or per-tensor quantized tensor | (C2), (C3) |
| (I6) | `epsilon` | constant of type `f32` | |
| (I7) | `feature_index` | constant of type `si64` | (C1), (C5) |
#### Outputs
| Name | Type | Constraints |
|----------------|---------------------------------------------------------------------|-------------|
| `grad_operand` | tensor of floating-point type or per-tensor quantized tensor | (C2), (C3) |
| `grad_scale` | 1-dimensional tensor of floating-point or per-tensor quantized type | (C2), (C4) |
| `grad_offset` | 1-dimensional tensor of floating-point or per-tensor quantized type | (C2), (C4) |
#### 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
```mlir
// %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:
```python
def batch_norm_inference(operand, scale, offset, mean, variance, epsilon, feature_index):
# Broadcast inputs to shape(operand)
scale_bcast = broadcast_in_dim(scale, [feature_index], type(operand))
offset_bcast = broadcast_in_dim(offset, [feature_index], type(operand))
mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
variance_bcast = broadcast_in_dim(variance, [feature_index], type(operand))
epsilon_bcast = broadcast_in_dim(constant(epsilon, element_type(operand)), [],
type(operand))
# Perform normalization using the provided `mean` and `variance` instead of
# computing them like `batch_norm_training` does.
centered_operand = subtract(operand, mean_bcast)
stddev = sqrt(add(variance_bcast, epsilon_bcast))
normalized_operand = divide(centered_operand, stddev)
return add(multiply(scale_bcast, normalized_operand), offset_bcast)
```
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
| Label | Name | Type | Constraints |
|-------|-----------------|---------------------------------------------------------------------|---------------|
| (I1) | `operand` | tensor of floating-point type or per-tensor quantized tensor | (C1-C7) |
| (I2) | `scale` | 1-dimensional tensor of floating-point or per-tensor quantized type | (C2), (C3) |
| (I3) | `offset` | 1-dimensional tensor of floating-point or per-tensor quantized type | (C2), (C4) |
| (I4) | `mean` | 1-dimensional tensor of floating-point or per-tensor quantized type | (C5) |
| (I5) | `variance` | 1-dimensional tensor of floating-point or per-tensor quantized type | (C2), (C6) |
| (I6) | `epsilon` | constant of type `f32` | |
| (I7) | `feature_index` | constant of type `si64` | (C1), (C3-C6) |
#### Outputs
| Name | Type | Constraints |
|----------|--------------------------------------------------------------|-------------|
| `result` | tensor of floating-point type or per-tensor quantized tensor | (C2), (C7) |
#### 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
```mlir
// %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:
```python
def compute_mean(operand, feature_index):
(sum,) = reduce(
inputs=[operand],
init_values=[constant(0, element_type(operand))],
dimensions=[i for i in range(rank(operand)) if i != feature_index],
body=lambda x, y: add(x, y))
divisor = constant(size(operand) / dim(operand, feature_index),
element_type(operand))
divisor_bcast = broadcast_in_dim(divisor, [], type(sum))
return divide(sum, divisor_bcast)
def compute_variance(operand, feature_index):
mean = compute_mean(operand, feature_index)
mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
centered_operand = subtract(operand, mean_bcast)
return compute_mean(mul(centered_operand, centered_operand), feature_index)
def batch_norm_training(operand, scale, offset, epsilon, feature_index):
mean = compute_mean(operand, feature_index)
variance = compute_variance(operand, feature_index)
return batch_norm_inference(operand, scale, offset, mean, variance, epsilon,
feature_index),
mean, variance
```
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
| Label | Name | Type | Constraints |
|-------|-----------------|----------------------------------------------------------------|---------------|
| (I1) | `operand` | tensor of floating-point type or per-tensor quantized tensor | (C1) |
| (I2) | `scale` | 1-dimensional tensor of floating-point or per-tensor quantized | (C2), (C3) |
| (I3) | `offset` | 1-dimensional tensor of floating-point or per-tensor quantized | (C2), (C4) |
| (I4) | `epsilon` | constant of type `f32` | (C1), (C3-C6) |
| (I5) | `feature_index` | constant of type `si64` | (C1), (C3-C6) |
#### Outputs
| Name | Type | Constraints |
|--------------|----------------------------------------------------------------|-------------|
| `output` | tensor of floating-point type or per-tensor quantized tensor | (C7) |
| `batch_mean` | 1-dimensional tensor of floating-point or per-tensor quantized | (C2), (C5) |
| `batch_var` | 1-dimensional tensor of floating-point or per-tensor quantized | (C2), (C6) |
#### 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
```mlir
// %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
| Label | Name | Type | Constraints |
|-------|-----------|----------------------------|-------------|
| (I1) | `operand` | tensor or quantized tensor | (C1-C2) |
#### Outputs
| Name | Type | Constraints |
|----------|----------------------------|-------------|
| `result` | tensor or quantized tensor | (C1-C2) |
#### 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
```mlir
// %operand: 0x0123456789ABCDEF
%result = "stablehlo.bitcast_convert"(%operand) : (tensor<f64>) -> tensor<4xf16>
// %result: [0xCDEF, 0x89AB, 0x4567, 0x0123] // little-endian representation
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/bitcast_convert.mlir)
### broadcast_in_dim
#### Semantics
![broadcast_in_dim](images/spec/broadcast_in_dim.svg)
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
| Label | Name | Type | Constraints |
|-------|------------------------|----------------------------------------------|------------------|
| (I1) | `operand` | tensor or quantized tensor | (C1-C2), (C5-C6) |
| (I2) | `broadcast_dimensions` | 1-dimensional tensor constant of type `si64` | (C2-C6) |
#### Outputs
| Name | Type | Constraints |
|----------|----------------------------|---------------------|
| `result` | tensor or quantized tensor | (C1), (C3), (C5-C6) |
#### 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
```mlir
// %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]
// ]
// ]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/broadcast_in_dim.mlir)
### 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
| Label | Name | Type | Constraints |
|-------|------------|-------------------------------------|-------------|
| (I1) | `index` | 0-dimensional tensor of type `si32` | |
| (I2) | `branches` | variadic number of functions | (C1-C4) |
#### Outputs
| Name | Type | Constraints |
|-----------|---------------------------------------------------------|-------------|
| `results` | variadic number of tensors, quantized tensors or tokens | (C4) |
#### Constraints
* (C1) `0 < size(branches)`.
* (C2) `input_types(branches...) = []`.
* (C3) `same(output_types(branches...))`.
* (C4) `type(results...) = output_types(branches[0])`.
#### Examples
```mlir
// %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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/case.mlir)
### 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
| Label | Name | Type | Constraints |
|-------|-----------|-------------------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------------------------------------------------|-------------|
| `result` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// %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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/cbrt.mlir)
### 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
| Label | Name | Type | Constraints |
|-------|-----------|--------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|--------------------------------------------------------------|-------------|
| `result` | tensor of floating-point type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// %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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/ceil.mlir)
### 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
| Label | Name | Type | Constraints |
|-------|---------|-------------------------------------------------------------------------|-------------|
| (I1) | `a` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1-C3) |
| (I2) | `lower` | constant of type `i1` | |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------------------------------------------------|-------------|
| `result` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(a) = baseline_type(result)`.
* (C2) `2 <= rank(a)`.
* (C3) `dim(a, -2) = dim(a, -1)`.
#### Examples
```mlir
// %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](https://github.com/openxla/stablehlo/issues/560)).
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|---------------------------------------|-------------|
| (I1) | `min` | tensor or per-tensor quantized tensor | (C1), (C3) |
| (I2) | `operand` | tensor or per-tensor quantized tensor | (C1-C4) |
| (I3) | `max` | tensor or per-tensor quantized tensor | (C2), (C3) |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-------------|
| `result` | tensor or per-tensor quantized tensor | (C4) |
#### 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
```mlir
// %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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/clamp.mlir)
### collective_broadcast
#### Semantics
Within each process group in the StableHLO process grid, send the value of the
`operands` tensors from the source process to the target processes and produce
`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`.
Let `N = size(results)`. The source process for `process_groups[i]` computing
`results[j]` (for `0 <= j < N`) is determined by `has_dynamic_root`:
* If `has_dynamic_root = false`: the source is `process_groups[i, 0]`.
* If `has_dynamic_root = true`: the source is
`process_groups[i, operands[N][j]]`.
Afterwards, `results[j]@process` is given by:
* `operands[j]@source_process` if there exists an `i` such that the process is
in `process_groups[i]`.
* `broadcast_in_dim(constant(is_quantized(results[j]) ? quantize(0,
element_type(results[j])) : 0, element_type(results[j])), [], type(results[j]))`
otherwise.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------------------|--------------------------------------------------------------------------------------------|-----------------------|
| (I1) | `operands` | variadic number of tensors or per-tensor quantized tensors | (C3), (C4), (C5), (6) |
| (I2) | `replica_groups` | variadic number of 1-dimensional tensor constants of type `si64` or `ReplicaGroupMeshAxes` | (C1), (C2) |
| (I3) | `channel_id` | constant of type `si64` | |
| (I4) | `has_dynamic_root` | constant of type `i1` | (C3), (C4), (C5) |
#### Outputs
| Name | Type | Constraints |
|-----------|------------------------------------------------------------|-------------|
| `results` | variadic number of tensors or per-tensor quantized tensors | (C3) |
#### 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) If `has_dynamic_root = false`:
* `size(operands) >= 1`.
* `type(results[j]) = type(operands[j])` for all `j` in `[0, size(operands))`.
* (C4) If `has_dynamic_root = true`:
* `size(operands) >= 2`.
* `type(operands[size(operands)-1]) = tensor<[size(operands)-1]xi32>`.
* `type(results[j]) = type(operands[j])` for all `j` in `[0, size(operands)-1)`.
* (C5) `size(results) = size(operands) - (has_dynamic_root ? 1 : 0)`.
* (C6) `0 <= operands[size(operands)-1] < shape(replica_groups)[1]`.
#### Examples
```mlir
// 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
| Label | Name | Type | Constraints |
|-------|-----------------------|----------------------------------------------|-------------|
| (I1) | `operand` | tensor or per-tensor quantized tensor | (C5) |
| (I2) | `source_target_pairs` | 2-dimensional tensor constant of type `si64` | (C1-C4) |
| (I3) | `channel_id` | constant of type `si64` | |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-------------|
| `result` | tensor or per-tensor quantized tensor | (C1) |
#### 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
```mlir
// 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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/collective_permute.mlir)
### collective_reduce
#### Semantics
Within each process group in the StableHLO process grid, send the value of the
`operands` tensor to the source process, applies a reduction
function `computation` to the values of the `operands` tensors from each process
and produces `results` tensors and saves the
`results` tensors only on the source process.
If `has_dynamic_root` is set to true, the last operand in `operands` must be
a flat array of `si32` containing the root values with 1-to-1 mapping to each
data operand.
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(0, element_type(result)), [], type(result))`
otherwise.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------------------------|------------------------------------------------------------------|-------------|
| (I1) | `operands` | variadic number of tensors or per-tensor quantized tensors | (C3) |
| (I2) | `replica_groups` | variadic number of 1-dimensional tensor constants of type `si64` | (C1), (C2) |
| (I3) | `channel_id` | constant of type `si64` | |
| (I4) | `has_dynamic_root` | constant of type `i1` | |
| (I5) | `computation` | function | (C4) |
#### Outputs
| Name | Type | Constraints |
|----------|------------------------------------------------------------|-------------|
| `results`| variadic number of tensors or per-tensor quantized tensors | (C3) |
#### 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) if `has_dynamic_root` == false, `shape(results) = shape(operands)`;
* else, `shape(results) = shape(operands[0, size(operands)-2])`,
`dim(operands[size(operands)-1]) == 1` with
`num_elements(operands[size(operands)-1]) == size(operands)-1`.
* (C4) `computation` has type `(tensor<E>, tensor<E>) -> (tensor<E>)` where
`is_promotable(element_type(operand), E)`.
#### Examples
```mlir
// 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_reduce"(%operand) {
has_dynamic_root = false
replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor1x2xi64>) -> tensor<1x2xi64>
// %result@(0, 0): [[16, 20]]
// %result@(1, 0): [[0, 0]]
// %result@(2, 0): [[0, 0]]
// %result@(3, 0): [[0, 0]]
```
```mlir
// 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]]
// %root: [3]
%result = "stablehlo.collective_reduce"(%operand, %root) {
has_dynamic_root = true
replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor1x2xi64>) -> tensor<1x2xi64>
// %result@(0, 0): [[0, 0]]
// %result@(1, 0): [[0, 0]]
// %result@(2, 0): [[0, 0]]
// %result@(3, 0): [[16, 20]]
```
### 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 floating-point element types with `compare_type = WEAKORDER`, the op
implements a total preorder (strict weak ordering) where `-0.0` and `+0.0` are
treated as equal, and all `NaN` representations are treated as equal to each
other and greater than `+infinity`:
* `EQ`: `compareQuietEqual(lhs, rhs) || (is_nan(lhs) && is_nan(rhs))`.
* `NE`: `!EQ(lhs, rhs)`.
* `GE`: `GT(lhs, rhs) || EQ(lhs, rhs)`.
* `GT`: `LT(rhs, lhs)`.
* `LE`: `LT(lhs, rhs) || EQ(lhs, rhs)`.
* `LT`: `(!is_nan(lhs) && is_nan(rhs)) ||`
`(!is_nan(lhs) && !is_nan(rhs) && compareQuietLess(lhs, rhs))`.
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](https://github.com/openxla/stablehlo/issues/560)).
For quantized types. performs `dequantize_compare(lhs, rhs,
comparison_direction)`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|------------------------|-----------------------------------------------------------------------------------------|-------------|
| (I1) | `lhs` | tensor or per-tensor quantized tensor | (C1-C3) |
| (I2) | `rhs` | tensor or per-tensor quantized tensor | (C1-C2) |
| (I3) | `comparison_direction` | enum of `EQ`, `NE`, `GE`, `GT`, `LE`, and `LT` | |
| (I4) | `compare_type` | optional enum of `NOTYPE`, `FLOAT`, `TOTALORDER`, `WEAKORDER`, `SIGNED`, and `UNSIGNED` | (C3) |
#### Outputs
| Name | Type | Constraints |
|----------|------------------------|-------------|
| `result` | tensor of boolean type | (C2) |
#### Constraints
* (C1) `baseline_element_type(lhs) = baseline_element_type(rhs)`.
* (C2) `shape(lhs) = shape(rhs) = shape(result)`.
* (C3) `compare_type` is defined as:
* `NOTYPE`, `SIGNED` (deprecated), or `UNSIGNED` (deprecated) if
`is_signed_integer(element_type(lhs))`,
`is_unsigned_integer(element_type(lhs))`, or
`is_boolean(element_type(lhs))`.
* `FLOAT`, `TOTALORDER`, or `WEAKORDER` if `is_float(element_type(lhs))`.
* `FLOAT` if `is_complex(element_type(lhs))`.
#### Examples
```mlir
// %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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/compare.mlir)
### 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
| Label | Name | Type | Constraints |
|-------|-------|-------------------------------|-------------|
| (I1) | `lhs` | tensor of type `f32` or `f64` | (C1-C3) |
| (I2) | `rhs` | tensor of type `f32` or `f64` | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|------------------------|-------------|
| `result` | tensor of complex type | (C2), (C3) |
#### 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
```mlir
// %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)]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/complex.mlir)
### composite
#### Semantics
Encapsulates an operation made up (composed) of other StableHLO operations,
taking `inputs`, `composite_attributes` and `regions` 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.
The intent of `regions` is to only be used to model ops with bodies (e.g.
`my_op` may have a body like `while` or `reduce`). If the decomposition is
inlined, the `regions` are ignored.
#### Inputs
| Label | Name | Type |
|-------|------------------------|------------------------------|
| (I1) | `inputs` | variadic number of values |
| (I2) | `name` | constant of type `string` |
| (I3) | `composite_attributes` | attribute dictionary |
| (I4) | `decomposition` | constant of type `string` |
| (I5) | `version` | constant of type `si32` |
| (I6) | `regions` | variadic number of functions |
#### Outputs
| Name | Type |
|-----------|---------------------------|
| `results` | variadic number of values |
#### 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
```mlir
%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>
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/composite.mlir)
### 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`.
1. `d` is equal to `dimension`, and `d0`, ... are `d`th dimension sizes
of `inputs`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------------|------------------------------------------------------------|------------------|
| (I1) | `inputs` | variadic number of tensors or per-tensor quantized tensors | (C1-C6) |
| (I2) | `dimension` | constant of type `si64` | (C2), (C4), (C6) |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-------------|
| `result` | tensor or per-tensor quantized tensor | (C5-C6) |
#### 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
```mlir
// %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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/concatenate.mlir)
### constant
#### Semantics
Produces an `output` tensor from a constant `value`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|---------|----------|-------------|
| (I1) | `value` | constant | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|----------------------------|-------------|
| `output` | tensor or quantized tensor | (C1) |
#### Constraints
* (C1) `type(value) = type(output)`.
#### Examples
```mlir
%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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/constant.mlir)
### 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](https://github.com/openxla/stablehlo/issues/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](https://github.com/openxla/stablehlo/issues/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](https://github.com/openxla/stablehlo/issues/1576)).
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|--------|-------------|
| (I1) | `operand` | tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|--------|-------------|
| `result` | tensor | (C1) |
#### Constraints
* (C1) `shape(operand) = shape(result)`.
#### Examples
```mlir
// %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)]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/convert.mlir)
### 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.
![convolution](images/spec/convolution.svg)
More formally, consider the following reframing of the inputs in terms of `lhs`
in order to be able to express windows of `lhs`:
<!-- markdownlint-disable line-length -->
* `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](https://github.com/openxla/stablehlo/issues/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)`.
<!-- markdownlint-enable line-length -->
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
| Label | Name | Type | Constraints |
|-------|-----------------------------------|--------------------------------------------------------------|-----------------------------------------------------------|
| (I1) | `lhs` | tensor or per-tensor quantized tensor | (C1), (C10-C11), (C14) (C25), (C27-C28), (C31-C32), (C34) |
| (I2) | `rhs` | tensor or quantized tensor | (C1), (C14-C16), (C25), (C27-C29), (C31-C34) |
| (I3) | `window_strides` | 1-dimensional tensor constant of type `si64` | (C2-C3), (C25) |
| (I4) | `padding` | 2-dimensional tensor constant of type `si64` | (C4), (C25) |
| (I5) | `lhs_dilation` | 1-dimensional tensor constant of type `si64` | (C5-C6), (C25) |
| (I6) | `rhs_dilation` | 1-dimensional tensor constant of type `si64` | (C7-C8), (C25) |
| (I7) | `window_reversal` | 1-dimensional tensor constant of type `i1` | (C9) |
| (I8) | `input_batch_dimension` | constant of type `si64` | (C10), (C13), (C25) |
| (I9) | `input_feature_dimension` | constant of type `si64` | (C11), (C13-C14) |
| (I10) | `input_spatial_dimensions` | 1-dimensional tensor constant of type `si64` | (C12), (C13), (C25) |
| (I11) | `kernel_input_feature_dimension` | constant of type `si64` | (C14), (C18) |
| (I12) | `kernel_output_feature_dimension` | constant of type `si64` | (C15-C16), (C18), (C25), (C29) |
| (I13) | `kernel_spatial_dimensions` | 1-dimensional tensor constant of type `si64` | (C17-C18), (C25) |
| (I14) | `output_batch_dimension` | constant of type `si64` | (C20), (C25) |
| (I15) | `output_feature_dimension` | constant of type `si64` | (C20), (C25), (C30) |
| (I16) | `output_spatial_dimensions` | 1-dimensional tensor constant of type `si64` | (C19-C20), (C25) |
| (I17) | `feature_group_count` | constant of type `si64` | (C11), (C14), (C16), (C21), (C23) |
| (I18) | `batch_group_count` | constant of type `si64` | (C10), (C15), (C22), (C23), (C25) |
| (I19) | `precision_config` | variadic number of enums of `DEFAULT`, `HIGH`, and `HIGHEST` | (C24) |
#### Outputs
| Name | Type | Constraints |
|----------|----------------------------|----------------------------|
| `result` | tensor or quantized tensor | (C25-C28), (C30), (C32-34) |
#### Constraints
<!-- markdownlint-disable line-length -->
* (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)` or
(`is_fp8(element_type(lhs))` and `is_fp8(element_type(rhs))` and
`element_type(result) = element_type(lhs)`).
* 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)`.
<!-- markdownlint-enable line-length -->
#### Examples
```mlir
// %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]]
// ]]
```
&nbsp;[More Examples](../stablehlo/tests/interpret/convolution.mlir)
### 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
| Label | Name | Type | Constraints |
|-------|-----------|-------------------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------------------------------------------------|-------------|
| `result` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// %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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/cosine.mlir)
### 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
| Label | Name | Type | Constraints |
|-------|-----------|------------------------|-------------|
| (I1) | `operand` | tensor of integer type | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|------------------------|-------------|
| `result` | tensor of integer type | (C1) |
#### Constraints
* (C1) `type(operand) = type(result)`.
#### Examples
```mlir
// %operand: [[0, 1], [128, -1]]
%result = "stablehlo.count_leading_zeros"(%operand) : (tensor<2x2xi64>) -> tensor<2x2xi64>
// %result: [[64, 63], [56, 0]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/count_leading_zeros.mlir)
### 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](https://github.com/openxla/stablehlo/issues/741)).
#### Inputs
| Label | Name | Type |
|-------|--------------------------|--------------------------------------------------------|
| (I1) | `inputs` | variadic number of values |
| (I2) | `call_target_name` | constant of type `string` |
| (I3) | `has_side_effect` | constant of type `i1` |
| (I4) | `backend_config` | constant of type `string` or attribute dictionary |
| (I5) | `api_version` | constant of type `si32` |
| (I6) | `called_computations` | variadic number of constants of type `string` |
| (I7) | `output_operand_aliases` | specify the aliasing parts in the outputs and operands |
#### Outputs
| Name | Type |
|-----------|---------------------------|
| `results` | variadic number of values |
### (XLA GPU Support) Special custom_call targets
There are three special `call_target_name` related to `buffer` types:
`CreateBuffer` creates an uninitialized `buffer`, `Pin` creates an initialized
`buffer` and `Unpin` deallocates a `buffer` and returns the content of the
`buffer`.
```mlir
%uninitialized_buffer = "stablehlo.custom_call"() {
call_target_name = "CreateBuffer",
api_version = 4 : i32,
} : () -> memref<4xf64>
%initialized_buffer = "stablehlo.custom_call"(%init_value) {
call_target_name = "Pin",
api_version = 4 : i32,
} : (tensor<4xf64>) -> memref<4xf64>
%dealloc_buffer = "stablehlo.custom_call"(%initialized_buffer) {
call_target_name = "Unpin",
api_version = 4 : i32,
} : (memref<4xf64>) -> tensor<4xf64>
```
### Alias
Some custom_call ops may require a part in the outputs and a part in the
operands to share the same memory. This can be expressed via
`output_operand_aliases`. An alias pair representation consists a list of output
tuple indices representing the output part, and an operand_index along with a
list of operand tuple indices representing the operand part. The list of output
or operand tuple indices is empty if the corresponding type is not a `tuple`
type, and can be arbitrarily long for an arbitrarily nested tuple type. This
is similar to [the XLA alias representation](https://www.tensorflow.org/xla/aliasing).
The output part and the input part in an alias pair must have the same type. For
custom_call ops that aren't call to `CreateBuffer`, `Pin` and `Unpin`, a
`buffer` operand can appear in at most one pair of alias, and a `buffer` output
must appear in one pair of alias.
#### Examples
```mlir
%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>
%updated_buffer = "stablehlo.custom_call"(%buffer) {
call_target_name = "Update",
api_version = 4 : i32,
output_operand_aliases = [
#stablehlo.output_operand_alias<output_tuple_indices = [],
operand_index = 0,
operand_tuple_indices = []>]
} : (memref<4xf64>) -> memref<4xf64>
```
### 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
| Label | Name | Type | Constraints |
|-------|-------|----------------------------------------------------------------------------------|-------------|
| (I1) | `lhs` | tensor of integer, floating-point or complex type or per-tensor quantized tensor | (C1) |
| (I2) | `rhs` | tensor of integer, floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-----------------------------------------------------------------------------------|-------------|
| `result` | tensor of integer, floating-point, or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(lhs) = baseline_type(rhs) = baseline_type(result)`.
#### Examples
```mlir
// %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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/divide.mlir)
### 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:
<!-- markdownlint-disable line-length -->
* `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))`.
<!-- markdownlint-enable line-length -->
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](https://github.com/openxla/stablehlo/issues/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](https://arxiv.org/pdf/1904.06376.pdf):
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:
```txt
// 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`](https://github.com/openxla/xla/blob/e8a707554de6b3d6bfd891583a81ff7020a97b54/xla/xla_data.proto#L1022)
for some supported algorithm values. Ticket #2483 captures the plan to create a
centralized doc on supported algorithms by backend.
#### Inputs
| Label | Name | Type | Constraints |
|-------|--------------------------------|--------------------------------------------------------------|------------------------------------------------|
| (I1) | `lhs` | tensor or per-tensor quantized tensor | (C5-C6), (C9-C10), (C12-C14), (C17-C18), (C20) |
| (I2) | `rhs` | tensor or quantized tensor | (C7-C10), (C12-C20) |
| (I3) | `lhs_batching_dimensions` | 1-dimensional tensor constant of type `si64` | (C1), (C3), (C5), (C9), (C12) |
| (I4) | `rhs_batching_dimensions` | 1-dimensional tensor constant of type `si64` | (C1), (C4), (C7), (C9) |
| (I5) | `lhs_contracting_dimensions` | 1-dimensional tensor constant of type `si64` | (C2), (C3), (C6), (C10) |
| (I6) | `rhs_contracting_dimensions` | 1-dimensional tensor constant of type `si64` | (C2), (C4), (C8), (C10), (C16) |
| (I7) | `precision_config` | variadic number of enums of `DEFAULT`, `HIGH`, and `HIGHEST` | (C11), (C21) |
| (I8) | `lhs_precision_type` | FloatType or TensorFloat32 | (C21) |
| (I9) | `rhs_precision_type` | FloatType or TensorFloat32 | (C21) |
| (I10) | `accumulation_type` | FloatType or TensorFloat32 | (C21) |
| (I11) | `lhs_component_count` | constant of type `si32` | (C21), (C22) |
| (I12) | `rhs_component_count` | constant of type `si32` | (C21), (C23) |
| (I13) | `num_primitive_operations` | constant of type `si32` | (C21), (C24) |
| (I14) | `allow_imprecise_accumulation` | constant of type `bool` | (C21) |
#### Outputs
| Name | Type | Constraints |
|----------|----------------------------|-------------------------|
| `result` | tensor or quantized tensor | (C12), (C14), (C18-C20) |
#### 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
```mlir
// %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]]
// ]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/dot_general.mlir)
### dynamic_broadcast_in_dim
#### Semantics
This operation is functionally identical to
[broadcast_in_dim](https://github.com/openxla/stablehlo/blob/main/docs/spec.md#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
| Label | Name | Type | Constraints |
|-------|---------------------------------|-----------------------------------------------|------------------------|
| (I1) | `operand` | tensor or quantized tensor | (C1-C2), (C5-C6), (C9) |
| (I2) | `output_dimensions` | 1-dimensional tensor of integer type | (C7) |
| (I3) | `broadcast_dimensions` | 1-dimensional constant tensor of integer type | (C2-C6) |
| (I4) | `known_expanding_dimensions` | 1-dimensional constant tensor of integer type | (C8-C9) |
| (I5) | `known_nonexpanding_dimensions` | 1-dimensional constant tensor of integer type | (C8-C9) |
#### Outputs
| Name | Type | Constraints |
|----------|----------------------------|---------------------|
| `result` | tensor or quantized tensor | (C1), (C3), (C5-C7) |
#### 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
```mlir
// %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]
// ]
// ]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/dynamic_broadcast_in_dim.mlir)
### dynamic_conv
#### Semantics
This operation is functionally identical to
[convolution](https://github.com/openxla/stablehlo/blob/main/docs/spec.md#convolution)
op, but the padding is specified dynamically via `padding`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------------------------------|--------------------------------------------------------------|-----------------------------------------------------------|
| (I1) | `lhs` | tensor or per-tensor quantized tensor | (C1), (C10-C11), (C14) (C25), (C26-C27), (C30-C31), (C33) |
| (I2) | `rhs` | tensor or quantized tensor | (C1), (C14-C16), (C26-C28), (C30-C33) |
| (I3) | `padding` | 2-dimensional tensor of integer type | (C4) |
| (I4) | `window_strides` | 1-dimensional tensor constant of type `si64` | (C2-C3) |
| (I5) | `lhs_dilation` | 1-dimensional tensor constant of type `si64` | (C5-C6) |
| (I6) | `rhs_dilation` | 1-dimensional tensor constant of type `si64` | (C7-C8) |
| (I7) | `window_reversal` | 1-dimensional tensor constant of type `i1` | (C9) |
| (I8) | `input_batch_dimension` | constant of type `si64` | (C10), (C13) |
| (I9) | `input_feature_dimension` | constant of type `si64` | (C11), (C13-C14) |
| (I10) | `input_spatial_dimensions` | 1-dimensional tensor constant of type `si64` | (C12), (C13) |
| (I11) | `kernel_input_feature_dimension` | constant of type `si64` | (C14), (C18) |
| (I12) | `kernel_output_feature_dimension` | constant of type `si64` | (C15-C16), (C18), (C28) |
| (I13) | `kernel_spatial_dimensions` | 1-dimensional tensor constant of type `si64` | (C17-C18) |
| (I14) | `output_batch_dimension` | constant of type `si64` | (C20) |
| (I15) | `output_feature_dimension` | constant of type `si64` | (C20), (C29) |
| (I16) | `output_spatial_dimensions` | 1-dimensional tensor constant of type `si64` | (C19-C20) |
| (I17) | `feature_group_count` | constant of type `si64` | (C11), (C14), (C16), (C21), (C23) |
| (I18) | `batch_group_count` | constant of type `si64` | (C10), (C15), (C22), (C23) |
| (I19) | `precision_config` | variadic number of enums of `DEFAULT`, `HIGH`, and `HIGHEST` | (C24) |
#### Outputs
| Name | Type | Constraints |
|----------|----------------------------|-----------------------------|
| `result` | tensor or quantized tensor | (C25-C27), (C29), (C31-C33) |
#### Constraints
<!-- markdownlint-disable line-length -->
* (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)` or
(`is_fp8(element_type(lhs))` and `is_fp8(element_type(rhs))` and
`element_type(result) = element_type(lhs)`).
* 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)`.
<!-- markdownlint-enable line-length -->
#### Examples
```mlir
// %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]]
// ]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/dynamic_conv.mlir)
### dynamic_gather
#### Semantics
This operation is functionally identical to
[gather](https://github.com/openxla/stablehlo/blob/main/docs/spec.md#gather)
op, with the `slice_sizes` specified dynamically as a value.
#### Inputs
| Label | Name | Type | Constraints |
|-------|------------------------|----------------------------------------------|------------------------------|
| (I1) | `operand` | tensor or per-tensor quantized tensor | (C1), (C7), (C10-C12), (C14) |
| (I2) | `start_indices` | tensor of integer type | (C2), (C3), (C13) |
| (I3) | `slice_sizes` | 1-dimensional tensor of integer type | (C8), (C11-C13) |
| (I4) | `offset_dims` | 1-dimensional tensor constant of type `si64` | (C1), (C4-C5), (C13) |
| (I5) | `collapsed_slice_dims` | 1-dimensional tensor constant of type `si64` | (C1), (C6-C8), (C13) |
| (I6) | `start_index_map` | 1-dimensional tensor constant of type `si64` | (C3), (C9), (C10) |
| (I7) | `index_vector_dim` | constant of type `si64` | (C2), (C3), (C13) |
| (I8) | `indices_are_sorted` | constant of type `i1` | |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-----------------|
| `result` | tensor or per-tensor quantized tensor | (C5), (C13-C14) |
#### Constraints
* (C1) `rank(operand) = size(offset_dims) + size(collapsed_slice_dims)`.
* (C2) `0 <= index_vector_dim <= rank(start_indices)`.
* (C3) `size(start_index_map) =
index_vector_dim < rank(start_indices) ?
dim(start_indices, index_vector_dim) : 1`.
* (C4) `is_unique(offset_dims) and is_sorted(offset_dims)`.
* (C5) `0 <= offset_dims < rank(result)`.
* (C6) `is_unique(collapsed_slice_dims) and is_sorted(collapsed_slice_dims)`.
* (C7) `0 <= collapsed_slice_dims < rank(operand)`.
* (C8) `slice_sizes[collapsed_slice_dims...] <= 1`.
* (C9) `is_unique(start_index_map)`.
* (C10) `0 <= start_index_map < rank(operand)`.
* (C11) `size(slice_sizes) = rank(operand)`.
* (C12) `0 <= slice_sizes <= shape(operand)`.
* (C13) `shape(result) = combine(batch_dim_sizes, offset_dim_sizes)` where:
* `batch_dim_sizes = shape(start_indices)` except that the dimension size
of `start_indices` corresponding to `index_vector_dim` is not included.
* `offset_dim_sizes = shape(slice_sizes)` except that the dimension sizes
in `slice_sizes` corresponding to `collapsed_slice_dims` are not included.
* `combine` puts `batch_dim_sizes` at axes corresponding to `batch_dims` and
`offset_dim_sizes` at axes corresponding to `offset_dims`.
* (C14) `element_type(operand) = element_type(result)`.
#### Examples
```mlir
// %operand: [
// [[1, 2], [3, 4], [5, 6], [7, 8]],
// [[9, 10],[11, 12], [13, 14], [15, 16]],
// [[17, 18], [19, 20], [21, 22], [23, 24]]
// ]
// %start_indices: [
// [[0, 0], [1, 0], [2, 1]],
// [[0, 1], [1, 1], [0, 2]]
// ]
// %slize_sizes: [1, 2, 2]
%result = "stablehlo.dynamic_gather"(%operand, %start_indices, %slize_sizes) {
dimension_numbers = #stablehlo.gather<
offset_dims = [2, 3],
collapsed_slice_dims = [0],
start_index_map = [1, 0],
index_vector_dim = 2>,
indices_are_sorted = false
} : (tensor<3x4x2xi64>, tensor<2x3x2xi64>, tensor<3xi64>) -> tensor<2x3x2x2xi64>
// %result: [
// [
// [[1, 2], [3, 4]],
// [[3, 4], [5, 6]],
// [[13, 14], [15, 16]]
// ],
// [
// [[9, 10], [11, 12]],
// [[11, 12], [13, 14]],
// [[17, 18], [19, 20]]
// ]
// ]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/dynamic_gather.mlir)
### dynamic_iota
#### Semantics
This operation is functionally identical to
[iota](https://github.com/openxla/stablehlo/blob/main/docs/spec.md#iota)
op, but the result shape is specified dynamically via `output_shape`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|------------------|--------------------------------------|-------------|
| (I1) | `output_shape` | 1-dimensional tensor of integer type | (C1), (C2) |
| (I2) | `iota_dimension` | `si64` | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-----------------------------------------------------------------------------------|-------------|
| `result` | tensor of integer, floating-point, or complex type or per-tensor quantized tensor | (C2) |
#### Constraints
* (C1) `0 <= iota_dimension < size(output_shape)`.
* (C2) `rank(result) = size(output_shape)`.
#### Examples
```mlir
%output_shape = stablehlo.constant dense<[4, 5]> : tensor<2xi64>
%result = "stablehlo.dynamic_iota"(%output_shape) {
iota_dimension = 0 : i64
} : (tensor<2xi64>) -> tensor<4x5xi64>
// %result: [
// [0, 0, 0, 0, 0],
// [1, 1, 1, 1, 1],
// [2, 2, 2, 2, 2],
// [3, 3, 3, 3, 3]
// ]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/dynamic_iota.mlir)
### dynamic_pad
#### Semantics
This operation is functionally identical to
[pad](https://github.com/openxla/stablehlo/blob/main/docs/spec.md#pad)
op, but with `edge_padding_low`, `edge_padding_high`, and `interior_padding`
specified dynamically as values.
#### Inputs
| Label | Name | Type | Constraints |
|-------|---------------------|-----------------------------------------------------|------------------|
| (I1) | `operand` | tensor or per-tensor quantized tensor | (C1), (C2), (C4) |
| (I2) | `padding_value` | 0-dimensional tensor or per-tensor quantized tensor | (C1) |
| (I3) | `edge_padding_low` | 1-dimensional tensor of integer type | (C1), (C4) |
| (I4) | `edge_padding_high` | 1-dimensional tensor of integer type | (C1), (C4) |
| (I5) | `interior_padding` | 1-dimensional tensor of integer type | (C2-C4) |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-------------|
| `result` | tensor or per-tensor quantized tensor | (C3-C6) |
#### Constraints
* (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`.
#### Examples
```mlir
// %operand: [
// [1, 2, 3],
// [4, 5, 6]
// ]
// %padding_value: 0
// %edge_padding_low: [0, 1]
// %edge_padding_high: [2, 1]
// %interior_padding: [1, 2]
%result = "stablehlo.dynamic_pad"(%operand, %padding_value,
%edge_padding_low, %edge_padding_high, %interior_padding
) : (tensor<2x3xi64>, tensor<i64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> tensor<5x9xi64>
// %result: [
// [0, 1, 0, 0, 2, 0, 0, 3, 0],
// [0, 0, 0, 0, 0, 0, 0, 0, 0],
// [0, 4, 0, 0, 5, 0, 0, 6, 0],
// [0, 0, 0, 0, 0, 0, 0, 0, 0],
// [0, 0, 0, 0, 0, 0, 0, 0, 0]
// ]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/dynamic_pad.mlir)
### dynamic_reshape
#### Semantics
This operation is functionally identical to
[reshape](https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reshape)
op, but the result shape is specified dynamically via `output_shape`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|----------------|--------------------------------------|-------------|
| (I1) | `operand` | tensor or quantized tensor | (C1-C3) |
| (I2) | `output_shape` | 1-dimensional tensor of integer type | (C4) |
#### Outputs
| Name | Type | Constraints |
|----------|----------------------------|-------------|
| `result` | tensor or quantized tensor | (C1-C4) |
#### 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)` and
`quantization_dimension(result)` may differ, otherwise.
* (C2) `size(operand) = size(result)`.
* (C3) If `is_per_axis_quantized(operand)`:
* `reduce(dims(operand, [0, 1, ..., quantization_dimension(operand) - 1]),
init_values=1, dimensions=[0], body=lambda x, y: x * y) =
reduce(dims(result, [0, 1, ..., quantization_dimension(result) - 1]),
init_values=1, dimensions=[0], body=lambda x, y: x * y)`.
* `dim(operand, quantization_dimension(operand)) =
dim(result, quantization_dimension(result))`.
* `reduce(dims(operand,
[quantization_dimension(operand) + 1, ..., rank(operand) - 1]),
init_values=1, dimensions=[0], body=lambda x, y: x * y) =
reduce(dims(result,
[quantization_dimension(result) + 1, ..., rank(result) - 1]),
init_values=1, dimensions=[0], body=lambda x, y: x * y)`.
* (C4) `size(output_shape) = rank(result)`.
#### Examples
```mlir
// %operand: [[1, 2, 3], [4, 5, 6]]
// %output_shape: [3, 2]
%result = "stablehlo.dynamic_reshape"(%operand, %output_shape) : (tensor<2x3xi64>, tensor<2xi64>) -> tensor<3x2xi64>
// %result: [[1, 2], [3, 4], [5, 6]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/dynamic_reshape.mlir)
### dynamic_slice
#### Semantics
Extracts a slice from the `operand` using dynamically-computed starting indices
and produces a `result` tensor. `start_indices` contain the starting indices of
the slice for each dimension subject to potential adjustment, and `slice_sizes`
contain the sizes of the slice for each dimension. More formally,
`result[result_index] = operand[operand_index]` where:
* `adjusted_start_indices = clamp(0, start_indices, shape(operand) -
slice_sizes)`.
* `operand_index = adjusted_start_indices + result_index`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------------|----------------------------------------------------------|------------------|
| (I1) | `operand` | tensor or per-tensor quantized tensor | (C1), (C2), (C4) |
| (I2) | `start_indices` | variadic number of 0-dimensional tensors of integer type | (C2), (C3) |
| (I3) | `slice_sizes` | 1-dimensional tensor constant of type `si64` | (C2), (C4), (C5) |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-------------|
| `result` | tensor or per-tensor quantized tensor | (C1), (C5) |
#### Constraints
* (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`.
#### Examples
```mlir
// %operand: [
// [0, 0, 1, 1],
// [0, 0, 1, 1],
// [0, 0, 0, 0],
// [0, 0, 0, 0]
// ]
// %start_indices0: -1
// %start_indices1: 3
%result = "stablehlo.dynamic_slice"(%operand, %start_indices0, %start_indices1) {
slice_sizes = array<i64: 2, 2>
} : (tensor<4x4xi32>, tensor<i64>, tensor<i64>) -> tensor<2x2xi32>
// %result: [
// [1, 1],
// [1, 1]
// ]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/dynamic_slice.mlir)
### dynamic_update_slice
#### Semantics
Produces a `result` tensor which is equal to the `operand` tensor except that
the slice starting at `start_indices` is updated with the values in `update`.
More formally, `result[result_index]` is defined as:
* `update[update_index]` if `0 <= update_index < shape(update)` where:
* `adjusted_start_indices = clamp(0, start_indices, shape(operand) -
shape(update))`.
* `update_index = result_index - adjusted_start_indices`.
* `operand[result_index]` otherwise.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------------|----------------------------------------------------------|------------------|
| (I1) | `operand` | tensor or per-tensor quantized tensor | (C1-C4), (C6) |
| (I2) | `update` | tensor or per-tensor quantized tensor | (C2), (C3), (C6) |
| (I3) | `start_indices` | variadic number of 0-dimensional tensors of integer type | (C4), (C5) |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-------------|
| `result` | tensor or per-tensor quantized tensor | (C1) |
#### Constraints
* (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)`.
#### Examples
```mlir
// %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]
// ]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/dynamic_update_slice.mlir)
### exponential
#### Semantics
Performs element-wise exponential operation on `operand` tensor and produces a
`result` tensor. Depending on the element type, does the following:
* For floats: `exp` from IEEE-754.
* For complex numbers: complex exponential.
* For quantized types:
`dequantize_op_quantize(exponential, operand, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|-------------------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------------------------------------------------|-------------|
| `result` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// %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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/exponential.mlir)
### exponential_minus_one
#### Semantics
Performs element-wise exponential minus one operation on `operand` tensor and
produces a `result` tensor. Depending on the element type, does the following:
* For floats: `expm1` from IEEE-754.
* For complex numbers: complex exponential minus one.
* For quantized types:
`dequantize_op_quantize(exponential_minus_one, operand, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|-------------------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------------------------------------------------|-------------|
| `result` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// %operand: [0.0, 1.0]
%result = "stablehlo.exponential_minus_one"(%operand) : (tensor<2xf64>) -> tensor<2xf64>
// %result: [0.0, 1.71828187]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/exponential_minus_one.mlir)
### fft
#### Semantics
Performs the forward and inverse Fourier transforms for real and complex
inputs/outputs.
`fft_type` is one of the following:
* `FFT`: Forward complex-to-complex FFT.
* `IFFT`: Inverse complex-to-complex FFT.
* `RFFT`: Forward real-to-complex FFT.
* `IRFFT`: Inverse real-to-complex FFT (i.e. takes complex, returns real).
More formally, given the function `fft` which takes 1-dimensional tensors of
complex types as input, produces 1-dimensional tensors of same types as
output and computes the discrete Fourier transform:
For `fft_type = FFT`, `result` is defined as the final result of a series of L
computations where `L = size(fft_length)`. For example, for `L = 3`:
* `result1[i0, ..., :] = fft(operand[i0, ..., :])`.
* `result2[i0, ..., :, iR-1] = fft(result1[i0, ..., :, iR-1])`.
* `result[i0, ..., :, iR-2, iR-1] = fft(result2[i0, ..., :, iR-2, iR-1])`.
Furthermore, given the function `ifft` which has the same type signature and
computes the inverse of `fft`:
For `fft_type = IFFT`, `result` is defined as the inverse of the computations
for `fft_type = FFT`. For example, for `L = 3`:
* `result1[i0, ..., :, iR-2, iR-1] = ifft(operand[i0, ..., :, iR-2, iR-1])`.
* `result2[i0, ..., :, iR-1] = ifft(result1[i0, ..., :, iR-1])`.
* `result[i0, ..., :] = ifft(result2[i0, ..., :])`.
Furthermore, given the function `rfft` which takes 1-dimensional tensors of
floating-point types, produces 1-dimensional tensors of complex types of the
same floating-point semantics and works as follows:
* `rfft(real_operand) = truncated_result` where
* `complex_operand... = (real_operand..., 0.0)`.
* `complex_result = fft(complex_operand)`.
* `truncated_result = complex_result[:(rank(complex_result) / 2 + 1)]`.
(When the discrete Fourier transform is computed for real operands, the first
`N/2 + 1` elements of the result unambiguously define the rest of the result,
so the result of `rfft` is truncated to avoid computing redundant elements).
For `fft_type = RFFT`, `result` is defined as the final result of a series of L
computations where `L = size(fft_length)`. For example, for `L = 3`:
* `result1[i0, ..., :] = rfft(operand[i0, ..., :])`.
* `result2[i0, ..., :, iR-1] = fft(result1[i0, ..., :, iR-1])`.
* `result[i0, ..., :, iR-2, iR-1] = fft(result2[i0, ..., :, iR-2, iR-1])`.
Finally, given the function `irfft` which has the same type signature and
computes the inverse of `rfft`:
For `fft_type = IRFFT`, `result` is defined as the inverse of the computations
for `fft_type = RFFT`. For example, for `L = 3`:
* `result1[i0, ..., :, iR-2, iR-1] = ifft(operand[i0, ..., :, iR-2, iR-1])`.
* `result2[i0, ..., :, iR-1] = ifft(result1[i0, ..., :, iR-1])`.
* `result[i0, ..., :] = irfft(result2[i0, ..., :])`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|--------------|----------------------------------------------|------------------------|
| (I1) | `operand` | tensor of floating-point or complex type | (C1), (C2), (C4), (C5) |
| (I2) | `fft_type` | enum of `FFT`, `IFFT`, `RFFT`, and `IRFFT` | (C2), (C5) |
| (I3) | `fft_length` | 1-dimensional tensor constant of type `si64` | (C1), (C3), (C4) |
#### Outputs
| Name | Type | Constraints |
|----------|------------------------------------------|------------------|
| `result` | tensor of floating-point or complex type | (C2), (C4), (C5) |
#### Constraints
* (C1) `size(fft_length) <= rank(operand)`.
* (C2) The relationship between `operand` and `result` element types varies:
* If `fft_type = FFT`, `element_type(operand)` and `element_type(result)`
have the same complex type.
* If `fft_type = IFFT`, `element_type(operand)` and `element_type(result)`
have the same complex type.
* If `fft_type = RFFT`, `element_type(operand)` is a floating-point type and
`element_type(result)` is a complex type of the same floating-point
semantics.
* If `fft_type = IRFFT`, `element_type(operand)` is a complex type and
`element_type(result)` is a floating-point type of the same floating-point
semantics.
* (C3) `1 <= size(fft_length) <= 3`.
* (C4) If among `operand` and `result`, there is a tensor `real` of a
floating-point type, then `shape(real)[-size(fft_length):] = fft_length`.
* (C5) `shape(result) = shape(operand)` except for:
* If `fft_type = RFFT`,
`dim(result, -1) = dim(operand, -1) = 0 ? 0 : dim(operand, -1) / 2 + 1`.
* If `fft_type = IRFFT`,
`dim(operand, -1) = dim(result, -1) = 0 ? 0 : dim(result, -1) / 2 + 1`.
#### Examples
```mlir
// %operand: [(1.0, 0.0), (0.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
%result = "stablehlo.fft"(%operand) {
fft_type = #stablehlo<fft_type FFT>,
fft_length = array<i64: 4>
} : (tensor<4xcomplex<f32>>) -> tensor<4xcomplex<f32>>
// %result: [(1.0, 0.0), (1.0, 0.0), (1.0, 0.0), (1.0, 0.0)]
```
### floor
#### Semantics
Performs element-wise floor of `operand` tensor and produces a `result` tensor.
Implements the `roundToIntegralTowardNegative` operation from the IEEE-754
specification. For quantized types, performs
`dequantize_op_quantize(floor, operand, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|--------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|--------------------------------------------------------------|-------------|
| `result` | tensor of floating-point type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// %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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/floor.mlir)
### gather
#### Semantics
Gathers slices from `operand` tensor from offsets specified in `start_indices`
and produces a `result` tensor.
The following diagram shows how elements in `result` map on elements in
`operand` using a concrete example. The diagram picks a few example `result`
indices and explains in detail which `operand` indices they correspond to.
![gather](images/spec/gather.svg)
More formally, `result[result_index] = operand[operand_index]` where:
<!-- markdownlint-disable line-length -->
* `batch_dims = [d for d in axes(result) and d not in offset_dims]`.
* `batch_index = result_index[batch_dims...]`.
* `start_index` is defined as:
* `start_indices[bi0, ..., :, ..., biN]` where `bi` are individual elements in
`batch_index` and `:` is inserted at the `index_vector_dim` index, if
`index_vector_dim` < `rank(start_indices)`.
* `[start_indices[batch_index]]` otherwise.
* For `d_operand` in `axes(operand)`,
* `full_start_index[d_operand] = clamp(start_index[d_start], 0,
dim(operand, d_operand) - slice_sizes[d_operand])`
if `d_operand = start_index_map[d_start]`.
* `full_start_index[d_operand] = 0` otherwise.
* For `d_operand` in `axes(operand)`,
* `full_batching_index[d_operand] =
batch_index[d_start - (d_start < index_vector_dim ? 0 : 1)]`
if `d_operand = operand_batching_dims[i_batching]` and
`d_start = start_indices_batching_dims[i_batching]`.
* `full_batching_index[d_operand] = 0` otherwise.
* `offset_index = result_index[offset_dims...]`.
* `full_offset_index = [oi0, ..., 0, ..., oiN]` where `oi` are individual
elements in `offset_index`, and `0` is inserted at indices from
`collapsed_slice_dims` and `operand_batching_dims`.
* `operand_index = full_start_index + full_batching_index + full_offset_index`.
<!-- markdownlint-enable line-length -->
If `indices_are_sorted` is `true` then the implementation can assume that
`start_indices` are sorted with respect to `start_index_map`, otherwise the
behavior is undefined. More formally, for all `i1 < i2` from `indices(result)`,
`full_start_index(i1) <= full_start_index(i2)`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------------------------------|----------------------------------------------|--------------------------------------------|
| (I1) | `operand` | tensor or per-tensor quantized tensor | (C1), (C8), (C11), (C17), (C19-C21), (C23) |
| (I2) | `start_indices` | tensor of integer type | (C2-C3), (C14), (C17), (C22) |
| (I3) | `offset_dims` | 1-dimensional tensor constant of type `si64` | (C1), (C4-C5), (C22) |
| (I4) | `collapsed_slice_dims` | 1-dimensional tensor constant of type `si64` | (C1), (C6-C9), (C22) |
| (I5) | `operand_batching_dims` | 1-dimensional tensor constant of type `si64` | (C1), (C6), (C10-C12), (C16-C18), (C22) |
| (I6) | `start_indices_batching_dims` | 1-dimensional tensor constant of type `si64` | (C13-C17) |
| (I7) | `start_index_map` | 1-dimensional tensor constant of type `si64` | (C3), (C18-C19) |
| (I8) | `index_vector_dim` | constant of type `si64` | (C2-C3), (C15), (C22) |
| (I9) | `slice_sizes` | 1-dimensional tensor constant of type `si64` | (C9), (C12), (C20-C22) |
| (I10) | `indices_are_sorted` | constant of type `i1` | |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-----------------|
| `result` | tensor or per-tensor quantized tensor | (C5), (C22-C23) |
#### Constraints
* (C1) `rank(operand) = size(offset_dims) + size(collapsed_slice_dims) +
size(operand_batching_dims)`.
* (C2) `0 <= index_vector_dim <= rank(start_indices)`.
* (C3) `size(start_index_map) =
index_vector_dim < rank(start_indices) ?
dim(start_indices, index_vector_dim) : 1`.
* (C4) `is_unique(offset_dims) and is_sorted(offset_dims)`.
* (C5) `0 <= offset_dims < rank(result)`.
* (C6) `is_unique(concatenate(collapsed_slice_dims, operand_batching_dims))`
* (C7) `is_sorted(collapsed_slice_dims)`.
* (C8) `0 <= collapsed_slice_dims < rank(operand)`.
* (C9) `slice_sizes[collapsed_slice_dims...] <= 1`.
* (C10) `is_sorted(operand_batching_dims)`.
* (C11) `0 <= operand_batching_dims < rank(operand)`.
* (C12) `slice_sizes[operand_batching_dims...] <= 1`.
* (C13) `is_unique(start_indices_batching_dims)`.
* (C14) `0 <= start_indices_batching_dims < rank(start_indices)`.
* (C15) `index_vector_dim not in start_indices_batching_dims`.
* (C16) `size(operand_batching_dims) == size(start_indices_batching_dims)`.
* (C17) `dim(operand, operand_batching_dims...) =
dim(start_indices, start_indices_batching_dims...)`.
* (C18) `is_unique(concatenate(start_index_map, operand_batching_dims))`.
* (C19) `0 <= start_index_map < rank(operand)`.
* (C20) `size(slice_sizes) = rank(operand)`.
* (C21) `0 <= slice_sizes <= shape(operand)`.
* (C22) `shape(result) = combine(batch_dim_sizes, offset_dim_sizes)` where:
* `batch_dim_sizes = shape(start_indices)` except that the dimension size
of `start_indices` corresponding to `index_vector_dim` is not included.
* `offset_dim_sizes = slice_sizes` except that the dimension sizes in
`slice_sizes` corresponding to `collapsed_slice_dims` and
`operand_batching_dims` are not included.
* `combine` puts `batch_dim_sizes` at axes corresponding to `batch_dims` and
`offset_dim_sizes` at axes corresponding to `offset_dims`.
* (C23) `element_type(operand) = element_type(result)`.
#### Examples
```mlir
// %operand: [
// [
// [[1, 2], [3, 4], [5, 6], [7, 8]],
// [[9, 10],[11, 12], [13, 14], [15, 16]],
// [[17, 18], [19, 20], [21, 22], [23, 24]]
// ],
// [
// [[25, 26], [27, 28], [29, 30], [31, 32]],
// [[33, 34], [35, 36], [37, 38], [39, 40]],
// [[41, 42], [43, 44], [45, 46], [47, 48]]
// ]
// ]
// %start_indices: [
// [
// [[0, 0], [1, 0], [2, 1]],
// [[0, 1], [1, 1], [0, 9]]
// ],
// [
// [[0, 0], [2, 1], [2, 2]],
// [[1, 2], [0, 1], [1, 0]]
// ]
// ]
%result = "stablehlo.gather"(%operand, %start_indices) {
dimension_numbers = #stablehlo.gather<
offset_dims = [3, 4],
collapsed_slice_dims = [1],
operand_batching_dims = [0],
start_indices_batching_dims = [1],
start_index_map = [2, 1],
index_vector_dim = 3>,
slice_sizes = array<i64: 1, 1, 2, 2>,
indices_are_sorted = false
} : (tensor<2x3x4x2xi32>, tensor<2x2x3x2xi64>) -> tensor<2x2x3x2x2xi32>
// %result: [
// [
// [
// [[1, 2], [3, 4]],
// [[3, 4], [5, 6]],
// [[13, 14], [15, 16]]
// ],
// [
// [[33, 34], [35, 36]],
// [[35, 36], [37, 38]],
// [[41, 42], [43, 44]]
// ]
// ],
// [
// [
// [[1, 2], [3, 4]],
// [[13, 14], [15, 16]],
// [[21, 22], [23, 24]]
// ],
// [
// [[43, 44], [45, 46]],
// [[33, 34], [35, 36]],
// [[27, 28], [29, 30]]
// ]
// ]
// ]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/gather.mlir)
### get_dimension_size
#### Semantics
Produces the size of the given `dimension` of the `operand`. More formally,
`result = dim(operand, dimension)`. The Semantics concerns only with the shape
component of the type. The element-type could be anything.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------------|----------------------------|-------------|
| (I1) | `operand` | tensor or quantized tensor | (C1) |
| (I2) | `dimension` | constant of type `si64` | (C1) |
#### Outputs
| Name | Type |
|----------|-------------------------------------|
| `result` | 0-dimensional tensor of type `si32` |
#### Constraints
* (C1) `0 <= dimension < rank(operand)`.
#### Examples
```mlir
// %operand: [[1, 2, 3], [4, 5, 6]]
%result = "stablehlo.get_dimension_size"(%operand) {
dimension = 1 : i64
} : (tensor<2x3xi64>) -> tensor<i32>
// %result: 3
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/get_dimension_size.mlir)
### get_tuple_element
> Note: Per [StableHLO v1.0 Cleanup #2283](https://github.com/openxla/stablehlo/pull/2283),
> this op is being explored for deprecation as it appears to be unused by both
> frameworks and compilers. As such, it has limited compatibility guarantees
> (6 months).
#### Semantics
Extracts element at `index` position of the `operand` tuple and produces a
`result`. More formally, `result = operand[index]`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|-------------------------|-------------|
| (I1) | `operand` | tuple | (C1), (C2) |
| (I2) | `index` | constant of type `si32` | (C1), (C2) |
#### Outputs
| Name | Type | Constraints |
|----------|-----------|-------------|
| `result` | any value | (C2) |
#### Constraints
* (C1) `0 <= index < size(operand)`.
* (C2) `type(result) = tuple_element_types(operand)[index]`.
#### Examples
```mlir
// %operand: ([1.0, 2.0], (3))
%result = "stablehlo.get_tuple_element"(%operand) <{index = 0 : i32}> : (tuple<tensor<2xf64>, tuple<tensor<i64>>>) -> tensor<2xf64>
// %result: [1.0, 2.0]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/tuple_and_get_tuple_element.mlir)
### if
#### Semantics
Produces the output from executing exactly one function from `true_branch` or
`false_branch` depending on the value of `pred`. More formally, `result =
pred ? true_branch() : false_branch()`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|----------------|-----------------------------------|-------------|
| (I1) | `pred` | 0-dimensional tensor of type `i1` | |
| (I2) | `true_branch` | function | (C1-C3) |
| (I3) | `false_branch` | function | (C1), (C2) |
#### Outputs
| Name | Type | Constraints |
|-----------|---------------------------------------------------------|-------------|
| `results` | variadic number of tensors, quantized tensors or tokens | (C3) |
#### Constraints
* (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)`.
#### Examples
```mlir
// %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
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/if.mlir)
### imag
#### Semantics
Extracts the imaginary part, element-wise, from the `operand` and produces a
`result` tensor. More formally, for each element `x`:
`imag(x) = is_complex(x) ? imaginary_part(x) :
constant(0, element_type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point or complex type | (C1), (C2) |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------|-------------|
| `result` | tensor of floating-point type | (C1), (C2) |
#### Constraints
* (C1) `shape(result) = shape(operand)`.
* (C2) `element_type(result)` is defined as:
* `complex_element_type(element_type(operand))` if `is_complex(operand)`.
* `element_type(operand)` otherwise.
#### Examples
```mlir
// %operand: [(1.0, 2.0), (3.0, 4.0)]
%result = "stablehlo.imag"(%operand) : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
// %result: [2.0, 4.0]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/imag.mlir)
### infeed
#### Semantics
Reads data from the infeed and produces `results`.
Semantics of `infeed_config` is implementation-defined.
`results` consist of payload values which come first and a token which comes
last. In the future, we are planning to split the payload and the token into two
separate outputs to improve clarity
([#670](https://github.com/openxla/stablehlo/issues/670)).
#### Inputs
| Label | Name | Type |
|-------|-----------------|---------------------------|
| (I1) | `token` | `token` |
| (I2) | `infeed_config` | constant of type `string` |
#### Outputs
| Name | Type | Constraints |
|-----------|---------------------------------------------------------|-------------|
| `results` | variadic number of tensors, quantized tensors or tokens | (C1-C3) |
#### Constraints
* (C1) `0 < size(results)`.
* (C2) `is_empty(result[:-1])` or `is_tensor(type(results[:-1]))`.
* (C3) `is_token(type(results[-1]))`.
#### Examples
```mlir
// %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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/infeed.mlir)
### iota
#### Semantics
Fills an `output` tensor with values in increasing order starting from zero
along the `iota_dimension` dimension. More formally,
`output[output_index] = constant(is_quantized(output) ?
quantize(output_index[iota_dimension], element_type(output)) :
output_index[iota_dimension], element_type(output))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|------------------|--------|-------------|
| (I1) | `iota_dimension` | `si64` | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-----------------------------------------------------------------------------------|-------------|
| `output` | tensor of integer, floating-point, or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `0 <= iota_dimension < rank(output)`.
#### Examples
```mlir
%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]
// ]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/iota.mlir)
### is_finite
#### Semantics
Performs element-wise check whether the value in `x` is finite (i.e. is neither
+Inf, -Inf, nor NaN) and produces a `y` tensor. Implements the `isFinite`
operation from the IEEE-754 specification. For quantized types, the result is
always `true`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|------|--------------------------------------------------------------|-------------|
| (I1) | `x` | tensor of floating-point type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|------|------------------------|-------------|
| `y` | tensor of boolean type | (C1) |
#### Constraints
* (C1) `shape(x) = shape(y)`.
#### Examples
```mlir
// 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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/is_finite.mlir)
### log
#### Semantics
Performs element-wise logarithm operation on `operand` tensor and produces a
`result` tensor. Depending on the element type, does the following:
* For floats: `log` from IEEE-754.
* For complex numbers: complex logarithm.
* For quantized types: `dequantize_op_quantize(log, operand, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|-------------------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------------------------------------------------|-------------|
| `result` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// %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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/log.mlir)
### log_plus_one
#### Semantics
Performs element-wise logarithm plus one operation on `operand` tensor and
produces a `result` tensor. Depending on the element type, does the following:
* For floats: `logp1` from IEEE-754.
* For complex numbers:
`complex(log(hypot(real(x) + 1, imag(x))), atan2(imag(x), real(x) + 1))`
* For quantized types:
`dequantize_op_quantize(log_plus_one, operand, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|-------------------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------------------------------------------------|-------------|
| `result` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// %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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/log_plus_one.mlir)
### logistic
#### Semantics
Performs element-wise logistic operation on `operand` tensor and produces a
`result` tensor. Depending on the element type, does the following:
* For floats: `division(1, addition(1, exp(-x)))` from IEEE-754.
* For complex numbers: complex logistic.
* For quantized types:
`dequantize_op_quantize(logistic, operand, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|-------------------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------------------------------------------------|-------------|
| `result` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// %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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/logistic.mlir)
### map
> Note: Per [StableHLO v1.0 Cleanup #2283](https://github.com/openxla/stablehlo/pull/2283),
> this op is being explored for deprecation as it appears to be unused by both
> frameworks and compilers. As such, it has limited compatibility guarantees
> (6 months).
#### Semantics
Applies a map function `computation` to `inputs` along the `dimensions` and
produces a `result` tensor.
More formally, `result[result_index] = computation(inputs...[result_index])`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|---------------|------------------------------------------------------------|-------------|
| (I1) | `inputs` | variadic number of tensors or per-tensor quantized tensors | (C1-C4) |
| (I2) | `dimensions` | 1-dimensional tensor constant of type `si64` | (C3) |
| (I3) | `computation` | function | (C4) |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-------------|
| `result` | tensor or per-tensor quantized tensor | (C1), (C4) |
#### Constraints
* (C1) `shape(inputs...) = shape(result)`.
* (C2) `0 < size(inputs) = N`.
* (C3) `dimensions = range(rank(inputs[0]))`.
* (C4) `computation` has type `(tensor<E0>, ..., tensor<EN-1>) -> tensor<E'>`
where `Ei = element_type(inputs[i])` and `E' = element_type(result)`.
#### Examples
```mlir
// %input0: [[0, 1], [2, 3]]
// %input1: [[4, 5], [6, 7]]
%result = "stablehlo.map"(%input0, %input1) ({
^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
%0 = stablehlo.multiply %arg0, %arg1 : tensor<i64>
stablehlo.return %0 : tensor<i64>
}) {
dimensions = array<i64: 0, 1>
} : (tensor<2x2xi64>, tensor<2x2xi64>) -> tensor<2x2xi64>
// %result: [[0, 5], [12, 21]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/map.mlir)
### maximum
#### Semantics
Performs element-wise max operation on tensors `lhs` and `rhs` and produces a
`result` tensor. Depending on the element type, does the following:
* For booleans: logical OR.
* For integers: integer maximum.
* For floats: `maximum` from IEEE-754.
* For complex numbers: lexicographic maximum for the `(real, imaginary)` pair.
Imposing an ordering on complex numbers involves surprising semantics,
so in the future we are planning to remove support for complex numbers
for this operation ([#560](https://github.com/openxla/stablehlo/issues/560)).
* For quantized types:
* `dequantize_op_quantize(maximum, lhs, rhs, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------|---------------------------------------|-------------|
| (I1) | `lhs` | tensor or per-tensor quantized tensor | (C1) |
| (I2) | `rhs` | tensor or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-------------|
| `result` | tensor or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(lhs) = baseline_type(rhs) = baseline_type(result)`.
#### Examples
```mlir
// %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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/maximum.mlir)
### minimum
#### Semantics
Performs element-wise min operation on tensors `lhs` and `rhs` and produces a
`result` tensor. Depending on the element type, does the following:
* For booleans: logical AND.
* For integers: integer minimum.
* For floats: `minimum` from IEEE-754.
* For complex numbers: lexicographic minimum for the `(real, imaginary)` pair.
Imposing an ordering on complex numbers involves surprising semantics,
so in the future we are planning to remove support for complex numbers
for this operation ([#560](https://github.com/openxla/stablehlo/issues/560)).
* For quantized types:
* `dequantize_op_quantize(minimum, lhs, rhs, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------|---------------------------------------|-------------|
| (I1) | `lhs` | tensor or per-tensor quantized tensor | (C1) |
| (I2) | `rhs` | tensor or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-------------|
| `result` | tensor or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(lhs) = baseline_type(rhs) = baseline_type(result)`.
#### Examples
```mlir
// %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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/minimum.mlir)
### multiply
#### Semantics
Performs element-wise product of two tensors `lhs` and `rhs` and produces a
`result` tensor. Depending on the element type, does the following:
* For booleans: logical AND.
* For integers: integer multiplication.
* For floats: `multiplication` from IEEE-754.
* For complex numbers: complex multiplication.
* For quantized types:
* `dequantize_op_quantize(multiply, lhs, rhs, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------|---------------------------------------|-------------|
| (I1) | `lhs` | tensor or per-tensor quantized tensor | (C1) |
| (I2) | `rhs` | tensor or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-------------|
| `result` | tensor or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(lhs) = baseline_type(rhs) = baseline_type(result)`.
#### Examples
```mlir
// %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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/multiply.mlir)
### negate
#### Semantics
Performs element-wise negation of `operand` tensor and produces a `result`
tensor. Depending on the element type, does the following:
* For signed integers: integer negation.
* For unsigned integers: bitcast to signed integer, integer negation, bitcast
back to unsigned integer.
* For floats: `negate` from IEEE-754.
* For complex numbers: complex negation.
* For quantized types:
`dequantize_op_quantize(negate, operand, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|-----------------------------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of integer, floating-point, or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-----------------------------------------------------------------------------------|-------------|
| `result` | tensor of integer, floating-point, or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// 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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/negate.mlir)
### not
#### Semantics
Performs element-wise NOT of tensor `operand` and produces a `result` tensor.
Depending on the element type, does the following:
* For booleans: logical NOT.
* For integers: bitwise NOT.
#### Arguments
| Name | Type | Constraints |
|-----------|-----------------------------------|-------------|
| `operand` | tensor of boolean or integer type | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-----------------------------------|-------------|
| `result` | tensor of boolean or integer type | (C1) |
#### Constraints
* (C1) `type(operand) = type(result)`.
#### Examples
```mlir
// 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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/not.mlir)
### optimization_barrier
#### Semantics
Ensures that the operations that produce the `operand` are executed before any
operations that depend on the `result` and prevents compiler transformations
from moving operations across the barrier. Other than that, the operation is
an identity, i.e. `result = operand`.
#### Arguments
| Name | Type | Constraints |
|-----------|--------------------------------------------------------------------|-------------|
| `operand` | variadic number of tensors, per-tensor quantized tensors or tokens | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|--------------------------------------------------------------------|-------------|
| `result` | variadic number of tensors, per-tensor quantized tensors or tokens | (C1) |
#### Constraints
* (C1) `type(operand...) = type(result...)`.
#### Examples
```mlir
// %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
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/optimization_barrier.mlir)
### or
#### Semantics
Performs element-wise OR of two tensors `lhs` and `rhs` and produces a `result`
tensor. Depending on the element type, does the following:
* For booleans: logical OR.
* For integers: bitwise OR.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------|-----------------------------------|-------------|
| (I1) | `lhs` | tensor of integer or boolean type | (C1) |
| (I2) | `rhs` | tensor of integer or boolean type | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-----------------------------------|-------------|
| `result` | tensor of integer or boolean type | (C1) |
#### Constraints
* (C1) `type(lhs) = type(rhs) = type(result)`.
#### Examples
```mlir
// 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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/or.mlir)
### outfeed
#### Semantics
Writes `inputs` to the outfeed and produces a `result` token.
Semantics of `outfeed_config` is implementation-defined.
#### Inputs
| Label | Name | Type |
|-------|------------------|-------------------------------------------------|
| (I1) | `inputs` | variadic number of tensors or quantized tensors |
| (I2) | `token` | `token` |
| (I3) | `outfeed_config` | constant of type `string` |
#### Outputs
| Name | Type |
|----------|---------|
| `result` | `token` |
#### Examples
```mlir
%result = "stablehlo.outfeed"(%input0, %token) {
outfeed_config = ""
} : (tensor<2x2x2xi64>, !stablehlo.token) -> !stablehlo.token
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/outfeed.mlir)
### pad
#### Semantics
Expands `operand` by padding around the tensor as well as between the elements
of the tensor with the given `padding_value`.
`edge_padding_low` and `edge_padding_high` specify the amount of padding added
at the low-end (next to index 0) and the high-end (next to the highest index) of
each dimension respectively. The amount of padding can be negative, where the
absolute value of negative padding indicates the number of elements to remove
from the specified dimension.
`interior_padding` specifies the amount of padding added between any two
elements in each dimension which may not be negative. Interior padding occurs
before edge padding such that negative edge padding will remove elements from
the interior-padded operand.
More formally, `result[result_index]` is defined as:
* `operand[operand_index]` if
`result_index = edge_padding_low + operand_index * (interior_padding + 1)`.
* `padding_value` otherwise.
#### Inputs
| Label | Name | Type | Constraints |
|-------|---------------------|-----------------------------------------------------|------------------|
| (I1) | `operand` | tensor or per-tensor quantized tensor | (C1), (C2), (C4) |
| (I2) | `padding_value` | 0-dimensional tensor or per-tensor quantized tensor | (C1) |
| (I3) | `edge_padding_low` | 1-dimensional tensor constant of type `si64` | (C1), (C4) |
| (I4) | `edge_padding_high` | 1-dimensional tensor constant of type `si64` | (C1), (C4) |
| (I5) | `interior_padding` | 1-dimensional tensor constant of type `si64` | (C2-C4) |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-------------|
| `result` | tensor or per-tensor quantized tensor | (C3-C6) |
#### Constraints
* (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`.
#### Examples
```mlir
// %operand: [
// [1, 2, 3],
// [4, 5, 6]
// ]
// %padding_value: 0
%result = "stablehlo.pad"(%operand, %padding_value) {
edge_padding_low = array<i64: 0, 1>,
edge_padding_high = array<i64: 2, 1>,
interior_padding = array<i64: 1, 2>
} : (tensor<2x3xi32>, tensor<i32>) -> tensor<5x9xi32>
// %result: [
// [0, 1, 0, 0, 2, 0, 0, 3, 0],
// [0, 0, 0, 0, 0, 0, 0, 0, 0],
// [0, 4, 0, 0, 5, 0, 0, 6, 0],
// [0, 0, 0, 0, 0, 0, 0, 0, 0],
// [0, 0, 0, 0, 0, 0, 0, 0, 0]
// ]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/pad.mlir)
### partition_id
#### Semantics
Produces `partition_id` of the current process.
#### Outputs
| Name | Type |
|----------|-------------------------------------|
| `result` | 0-dimensional tensor of type `ui32` |
#### Examples
```mlir
%result = "stablehlo.partition_id"() : () -> tensor<ui32>
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/partition_id.mlir)
### popcnt
#### Semantics
Performs element-wise count of the number of bits set in the `operand` tensor
and produces a `result` tensor.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|------------------------|-------------|
| (I1) | `operand` | tensor of integer type | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|------------------------|-------------|
| `result` | tensor of integer type | (C1) |
#### Constraints
* (C1) `type(operand) = type(result)`.
#### Examples
```mlir
// %operand: [0, 1, 2, 127]
%result = "stablehlo.popcnt"(%operand) : (tensor<4xi64>) -> tensor<4xi64>
// %result: [0, 1, 1, 7]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/popcnt.mlir)
### power
#### Semantics
Performs element-wise exponentiation of `lhs` tensor by `rhs` tensor and
produces a `result` tensor. Depending on the element type, does the following:
* For integers: integer exponentiation.
* For floats: `pow` from IEEE-754.
* For complex numbers: complex exponentiation.
* For quantized types: `dequantize_op_quantize(power, lhs, rhs, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------|-----------------------------------------------------------------------------------|-------------|
| (I1) | `lhs` | tensor of integer, floating-point, or complex type or per-tensor quantized tensor | (C1) |
| (I2) | `rhs` | tensor of integer, floating-point, or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-----------------------------------------------------------------------------------|-------------|
| `result` | tensor of integer, floating-point, or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(lhs) = baseline_type(rhs) = baseline_type(result)`.
#### Examples
```mlir
// %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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/power.mlir)
### real
#### Semantics
Extracts the real part, element-wise, from the `operand` and produces a `result`
tensor. More formally, for each element `x`:
`real(x) = is_complex(x) ? real_part(x) : x`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point or complex type | (C1), (C2) |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------|-------------|
| `result` | tensor of floating-point type | (C1), (C2) |
#### Constraints
* (C1) `shape(result) = shape(operand)`.
* (C2) `element_type(result)` is defined as:
* `complex_element_type(element_type(operand))` if `is_complex(operand)`.
* `element_type(operand)` otherwise.
#### Examples
```mlir
// %operand: [(1.0, 2.0), (3.0, 4.0)]
%result = "stablehlo.real"(%operand) : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
// %result: [1.0, 3.0]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/real.mlir)
### recv
#### Semantics
Receives data from a channel with `channel_id` and produces `results`.
If `is_host_transfer` is `true`, then the operation transfers data from the
host. Otherwise, it transfers data from another device based on the values of
`source_target_pairs`. This flag duplicates the information provided in
`channel_type`, so in the future we are planning to only keep one of them
([#666](https://github.com/openxla/stablehlo/issues/666)). If `is_host_transfer`
= `false` and `source_target_pairs` is `None` or empty, it is considered
undefined behavior.
`results` consist of payload values which come first and a token which comes
last. In the future, we are planning to split the payload and the token into two
separate outputs to improve clarity
([#670](https://github.com/openxla/stablehlo/issues/670)).
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------------------|-------------------------------------------------|---------------|
| (I1) | `token` | `token` | |
| (I2) | `channel_id` | constant of type `si64` | |
| (I3) | `channel_type` | enum of `DEVICE_TO_DEVICE` and `DEVICE_TO_HOST` | (C5) |
| (I4) | `is_host_transfer` | constant of type `i1` | (C5-C6) |
| (I5) | `source_target_pairs` | 2-dimensional tensor constant of type `si64` | (C1-C4), (C6) |
#### Outputs
| Name | Type | Constraints |
|-----------|---------------------------------------------------------|-------------|
| `results` | variadic number of tensors, quantized tensors or tokens | (C2-C4) |
#### 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) `channel_type` is defined as:
* `DEVICE_TO_HOST` if `is_host_transfer = true`,
* `DEVICE_TO_DEVICE` otherwise.
#### Examples
```mlir
%results0, %results1 = "stablehlo.recv"(%token) {
channel_handle = #stablehlo.channel_handle<handle = 0, type = 1>,
is_host_transfer = false,
source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>
} : (!stablehlo.token) -> (tensor<2x2xi64>, !stablehlo.token)
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/send_recv.mlir)
### reduce
#### Semantics
Applies a reduction function `body` to `inputs` and `init_values` along the
`dimensions` and produces `results` tensors.
The order of reductions is implementation-defined, which means that `body` and
`init_values` must form a monoid to guarantee that the operation produces the
same results for all inputs on all implementations. However, this condition
doesn't hold for many popular reductions. E.g. floating-point addition for
`body` and zero for `init_values` don't actually form a monoid because
floating-point addition is not associative.
More formally, `results...[j0, ..., jR-1] = reduce(input_slices_converted)` where:
* `input_slices = inputs...[j0, ..., :, ..., jR-1]`, where `:` are inserted
at `dimensions`.
* `input_slices_converted = to_destination_type(input_slices...,
type(func_inputs(body)[:len(func_inputs(body))//2])...)`.
* `init_values_converted = to_destination_type(init_values...,
type(func_inputs(body)[len(func_inputs(body))//2:])...)`.
* `reduce(input_slices_converted) = exec(schedule)` for some binary tree
`schedule` where:
* `exec(node) = body(exec(node.left), exec(node.right))`.
* `exec(leaf) = leaf.value`.
* `schedule` is an implementation-defined full binary tree whose in-order
traversal consists of:
* `input_slices_converted...[index]` values, for all `index` in
`index_space(input_slices_converted)` in the ascending lexicographic order
of `index`.
* Interspersed with an implementation-defined amount of
`init_values_converted` at implementation-defined positions.
#### Inputs
| Label | Name | Type | Constraints |
|-------|---------------|--------------------------------------------------------------------------|---------------------|
| (I1) | `inputs` | variadic number of tensors or per-tensor quantized tensors | (C1-C4), (C6), (C7) |
| (I2) | `init_values` | variadic number of 0-dimensional tensors or per-tensor quantized tensors | (C2), (C3) |
| (I3) | `dimensions` | 1-dimensional tensor constant of type `si64` | (C4), (C5), (C7) |
| (I4) | `body` | function | (C6) |
#### Outputs
| Name | Type | Constraints |
|-----------|------------------------------------------------------------|------------------|
| `results` | variadic number of tensors or per-tensor quantized tensors | (C3), (C7), (C8) |
#### Constraints
* (C1) `same(shape(inputs...))`.
* (C2) `element_type(inputs...) = element_type(init_values...)`.
* (C3) `0 < size(inputs) = size(init_values) = size(results) = N`.
* (C4) `0 <= dimensions < rank(inputs[0])`.
* (C5) `is_unique(dimensions)`.
* (C6) `body` has type `(tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ...,`
`tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>)` where
`is_promotable(element_type(inputs[i]), Ei)`.
* (C7) `shape(results...) = shape(inputs...)` except that the dimension
sizes of `inputs...` corresponding to `dimensions` are not included.
* (C8) `element_type(results[i]) = Ei` for all `i` in `[0,N)`.
#### Examples
```mlir
// %input = [[0, 1, 2, 3, 4, 5]]
// %init_value = 0
%result = "stablehlo.reduce"(%input, %init_value) ({
^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
%0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
"stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
dimensions = array<i64: 1>
} : (tensor<1x6xi64>, tensor<i64>) -> tensor<1xi64>
// %result = [15]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/reduce.mlir)
### reduce_precision
#### Semantics
Performs element-wise conversion of `operand` to another floating-point type
that uses `exponent_bits` and `mantissa_bits` and back to the original
floating-point type and produces an `output` tensor.
More formally:
* The mantissa bits of the original value are updated to round the original
value to the nearest value representable with `mantissa_bits` using
`roundToIntegralTiesToEven` semantics.
* Then, if `mantissa_bits` are smaller than the number of mantissa bits of
the original value, the mantissa bits are truncated to `mantissa_bits`.
* Then, if the exponent bits of the intermediate result don't fit into the
range provided by `exponent_bits`, the intermediate result overflows to
infinity using the original sign or underflows to zero using the
original sign.
* For quantized types, performs `dequantize_op_quantize(
lambda operand: reduce_precision(operand, exponent_bits, mantissa_bits),
operand, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------------|--------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point type or per-tensor quantized tensor | (C1) |
| (I2) | `exponent_bits` | constant of type `si32` | (C2) |
| (I3) | `mantissa_bits` | constant of type `si32` | (C3) |
#### Outputs
| Name | Type | Constraints |
|----------|--------------------------------------------------------------|-------------|
| `output` | tensor of floating-point type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(output)`.
* (C2) `1 <= exponent_bits`.
* (C3) `0 <= mantissa_bits`.
#### Examples
```mlir
// 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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/reduce_precision.mlir)
### reduce_scatter
#### Semantics
![reduce_scatter](images/spec/reduce_scatter.svg)
Within each process group in the StableHLO process grid, performs reduction,
using `computations`, over the values of the `operand` tensor from each process,
splits the reduction result along `scatter_dimension` into parts, and scatters
the split parts between the processes to produce the `result`.
The operation splits the StableHLO process grid into `process_groups` which is
defined as follows:
* `cross_replica(replica_groups)`
if `channel_id <= 0 and use_global_device_ids = false`.
* `cross_replica_and_partition(replica_groups)`
if `channel_id > 0 and use_global_device_ids = false`.
* `flattened_ids(replica_groups)`
if `channel_id > 0 and use_global_device_ids = true`.
Afterwards, within each `process_group`:
* `reduced_value = all_reduce(operand, replica_groups, channel_id,
use_global_device_ids, computation)`.
* `parts@sender = split(reduced_value@sender, dim(process_groups, 1),
scatter_dimension)`.
* `result@receiver = parts@sender[receiver_index]` for all `sender` in
`process_group`, where `receiver_index = process_group.index(receiver)`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------------------------|------------------------------------------------------------------------|------------------------|
| (I1) | `operand` | tensor or per-tensor quantized tensor | (C1), (C2), (C7), (C8) |
| (I2) | `scatter_dimension` | constant of type `si64` | (C1), (C2), (C8) |
| (I3) | `replica_groups` | 2-dimensional tensor constant of type `si64` or `ReplicaGroupMeshAxes` | (C3-C5) |
| (I4) | `channel_id` | constant of type `si64` | (C6) |
| (I5) | `use_global_device_ids` | constant of type `i1` | (C6) |
| (I6) | `computation` | function | (C7) |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-------------|
| `result` | tensor or per-tensor quantized tensor | (C8-C9) |
#### Constraints
* (C1) `dim(operand, scatter_dimension) % dim(process_groups, 1) = 0`.
* (C2) `0 <= scatter_dimension < rank(operand)`.
* (C3) `is_unique(replica_groups)`.
* (C4) `size(replica_groups)` is defined as:
* `num_replicas` if `cross_replica` is used.
* `num_replicas` if `cross_replica_and_partition` is used.
* `num_processes` if `flattened_ids` is used.
* (C5) `0 <= replica_groups < size(replica_groups)`.
* (C6) If `use_global_device_ids = true`, then `channel_id > 0`.
* (C7) `computation` has type `(tensor<E>, tensor<E>) -> (tensor<E>)` where
`is_promotable(element_type(operand), E)`.
* (C8) `shape(result) = shape(operand)` except:
* `dim(result, scatter_dimension) = dim(operand, scatter_dimension) /
dim(process_groups, 1)`.
* (C9) `element_type(result) = E`.
#### Examples
```mlir
// 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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/reduce_scatter.mlir)
### reduce_window
#### Semantics
Applies a reduction function `body` to windows of `inputs` and `init_values`
and produces `results`.
The following diagram shows how elements in `results...` are computed from
`inputs...` using a concrete example.
![reduce_window](images/spec/reduce_window.svg)
More formally,
`results...[result_index] = reduce(windows, init_values, axes(inputs...), body)`
(see [reduce](#reduce)) where:
* `padded_inputs = pad(inputs..., init_values..., padding[:, 0], padding[:, 1],
base_dilations - 1)`.
* `window_start = result_index * window_strides`.
* `window_end = window_start + (window_dimensions - 1) * window_dilations + 1`.
* `windows = slice(padded_inputs..., window_start, window_end,
window_dilations)`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|---------------------|--------------------------------------------------------------------------|-------------------------------------------------|
| (I1) | `inputs` | variadic number of tensors or per-tensor quantized tensors | (C1-C4), (C6), (C8), (C10), (C12), (C13), (C15) |
| (I2) | `init_values` | variadic number of 0-dimensional tensors or per-tensor quantized tensors | (C1), (C13) |
| (I3) | `window_dimensions` | 1-dimensional tensor constant of type `si64` | (C4), (C5), (C15) |
| (I4) | `window_strides` | 1-dimensional tensor constant of type `si64` | (C6), (C7), (C15) |
| (I5) | `base_dilations` | 1-dimensional tensor constant of type `si64` | (C8), (C9), (C15) |
| (I6) | `window_dilations` | 1-dimensional tensor constant of type `si64` | (C10), (C11), (C15) |
| (I7) | `padding` | 2-dimensional tensor constant of type `si64` | (C12), (C15) |
| (I8) | `body` | function | (C13) |
#### Outputs
| Name | Type | Constraints |
|-----------|------------------------------------------------------------|-----------------|
| `results` | variadic number of tensors or per-tensor quantized tensors | (C1), (C14-C16) |
#### Constraints
<!-- markdownlint-disable line-length -->
* (C1) `0 < size(inputs) = size(init_values) = size(results) = N`.
* (C2) `same(shape(inputs...))`.
* (C3) `element_type(inputs...) = element_type(init_values...)`.
* (C4) `size(window_dimensions) = rank(inputs[0])`.
* (C5) `0 < window_dimensions`.
* (C6) `size(window_strides) = rank(inputs[0])`.
* (C7) `0 < window_strides`.
* (C8) `size(base_dilations) = rank(inputs[0])`.
* (C9) `0 < base_dilations`.
* (C10) `size(window_dilations) = rank(inputs[0])`.
* (C11) `0 < window_dilations`.
* (C12) `shape(padding) = [rank(inputs[0]), 2]`.
* (C13) `body` has type `(tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ...,`
`tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>)` where
`is_promotable(element_type(inputs[i]), Ei)`.
* (C14) `same(shape(results...))`.
* (C15) `shape(results[0]) = num_windows` where:
* `dilated_input_shape = shape(inputs[0]) = 0 ? 0 : (shape(inputs[0]) - 1) * base_dilations + 1`.
* `padded_input_shape = padding[:, 0] + dilated_input_shape + padding[:, 1]`.
* `dilated_window_shape = (window_dimensions - 1) * window_dilations + 1`.
* `is_empty_window = padded_input_shape = 0 || dilated_window_shape > padded_input_shape`.
* `num_windows = is_empty_window ? 0 : floor((padded_input_shape - dilated_window_shape) / window_strides) + 1`.
* (C16) `element_type(results[i]) = Ei` for all `i` in `[0,N)`.
<!-- markdownlint-enable line-length -->
#### Examples
```mlir
// %input = [[1, 2], [3, 4], [5, 6]]
// %init_value = 0
%result = "stablehlo.reduce_window"(%input, %init_value) ({
^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
%0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
"stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
window_dimensions = array<i64: 2, 1>,
window_strides = array<i64: 4, 1>,
base_dilations = array<i64: 2, 1>,
window_dilations = array<i64: 3, 1>,
padding = dense<[[2, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<3x2xi64>, tensor<i64>) -> tensor<2x2xi64>
// %result = [[0, 0], [3, 4]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/reduce_window.mlir)
### remainder
#### Semantics
Performs element-wise remainder of dividend `lhs` and divisor `rhs` tensors and
produces a `result` tensor.
More formally, the sign of the result is taken from the dividend, and the
absolute value of the result is always less than the divisor's absolute value.
The remainder is calculated as `lhs - d * rhs`, where `d` is given by:
* For integers: `stablehlo.divide(lhs, rhs)`.
* For floats: `division(lhs, rhs)` from IEEE-754 with rounding attribute
`roundTowardZero`.
* For complex numbers: TBD
([#997](https://github.com/openxla/stablehlo/issues/997)).
* For quantized types:
* `dequantize_op_quantize(remainder, lhs, rhs, type(result))`.
For floating-point element types, this operation is in contrast with the
`remainder` operation from IEEE-754 specification where `d` is an integral value
nearest to the exact value of `lhs/rhs` with ties to even.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------|----------------------------------------------------------------------------------|-------------|
| (I1) | `lhs` | tensor of integer, floating-point or complex type or per-tensor quantized tensor | (C1) |
| (I2) | `rhs` | tensor of integer, floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|----------------------------------------------------------------------------------|-------------|
| `result` | tensor of integer, floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(lhs) = baseline_type(rhs) = baseline_type(result)`.
#### Examples
```mlir
// %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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/remainder.mlir)
### replica_id
#### Semantics
Produces `replica_id` of the current process.
#### Outputs
| Name | Type |
|----------|-------------------------------------|
| `result` | 0-dimensional tensor of type `ui32` |
#### Examples
```mlir
%result = "stablehlo.replica_id"() : () -> tensor<ui32>
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/replica_id.mlir)
### reshape
#### Semantics
Performs reshape of `operand` tensor to a `result` tensor. Conceptually, it
amounts to keeping the same canonical representation but potentially changing
the shape, e.g. from `tensor<2x3xf32>` to `tensor<3x2xf32>` or `tensor<6xf32>`.
More formally, `result[result_index] = operand[operand_index]` where
`result_index` and `operand_index` have the same position in the lexicographic
ordering of `index_space(result)` and `index_space(operand)`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|----------------------------|-------------|
| (I1) | `operand` | tensor or quantized tensor | (C1-C3) |
#### Outputs
| Name | Type | Constraints |
|----------|----------------------------|-------------|
| `result` | tensor or quantized tensor | (C1-C3) |
#### 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)` and
`quantization_dimension(result)` may differ, otherwise.
* (C2) `size(operand) = size(result)`.
* (C3) If `is_per_axis_quantized(operand)`:
* `reduce(dims(operand, [0, 1, ..., quantization_dimension(operand) - 1]),
init_values=1, dimensions=[0], body=lambda x, y: x * y) =
reduce(dims(result, [0, 1, ..., quantization_dimension(result) - 1]),
init_values=1, dimensions=[0], body=lambda x, y: x * y)`.
* `dim(operand, quantization_dimension(operand)) =
dim(result, quantization_dimension(result))`.
* `reduce(dims(operand,
[quantization_dimension(operand) + 1, ..., rank(operand) - 1]),
init_values=1, dimensions=[0], body=lambda x, y: x * y) =
reduce(dims(result,
[quantization_dimension(result) + 1, ..., rank(result) - 1]),
init_values=1, dimensions=[0], body=lambda x, y: x * y)`.
#### Examples
```mlir
// %operand: [[1, 2, 3], [4, 5, 6]]
%result = "stablehlo.reshape"(%operand) : (tensor<2x3xi32>) -> tensor<3x2xi32>
// %result: [[1, 2], [3, 4], [5, 6]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/reshape.mlir)
### reverse
#### Semantics
Reverses the order of elements in the `operand` along the specified `dimensions`
and produces a `result` tensor. More formally,
`result[result_index] = operand[operand_index]` where:
* `operand_index[d] = dim(result, d) - result_index[d] - 1`
if `d` in `dimensions`.
* `operand_index[d] = result_index[d]` otherwise.
#### Inputs
| Label | Name | Type | Constraints |
|-------|--------------|----------------------------------------------|-------------|
| (I1) | `operand` | tensor or per-tensor quantized tensor | (C1), (C3) |
| (I2) | `dimensions` | 1-dimensional tensor constant of type `si64` | (C2), (C3) |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-------------|
| `result` | tensor or per-tensor quantized tensor | (C1), (C3) |
#### Constraints
* (C1) `type(operand) = type(result)`.
* (C2) `is_unique(dimensions)`.
* (C3) `0 <= dimensions < rank(result)`.
#### Examples
```mlir
// %operand = [[1, 2], [3, 4], [5, 6]]
%result = "stablehlo.reverse"(%operand) {
dimensions = array<i64: 1>
} : (tensor<3x2xi32>) -> tensor<3x2xi32>
// %result: [[2, 1], [4, 3], [6, 5]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/reverse.mlir)
### rng
> Note: Per [StableHLO v1.0 Cleanup #2283](https://github.com/openxla/stablehlo/pull/2283),
> this op is being explored for deprecation as it appears to be unused by both
> frameworks and compilers. As such, it has limited compatibility guarantees
> (6 months).
#### Semantics
Generates random numbers using the `rng_distribution` algorithm and produces a
`result` tensor of a given shape `shape`.
If `rng_distribution = UNIFORM`, then the random numbers are generated
following the uniform distribution over the interval `[a, b)`. If `a >= b`,
the behavior is undefined.
If `rng_distribution = NORMAL`, then the random numbers are generated
following the normal distribution with mean = `a` and standard deviation = `b`.
If `b < 0`, the behavior is undefined.
The exact way how random numbers are generated is implementation-defined. For
example, they may or may not be deterministic, and they may or may not use
hidden state.
In conversations with many stakeholders, this op has come up as effectively
deprecated, so in the future we are planning to explore removing it
([#597](https://github.com/openxla/stablehlo/issues/597)).
#### Inputs
| Label | Name | Type | Constraints |
|-------|--------------------|------------------------------------------------------------------|-------------|
| (I1) | `a` | 0-dimensional tensor of integer, boolean, or floating-point type | (C1), (C2) |
| (I2) | `b` | 0-dimensional tensor of integer, boolean, or floating-point type | (C1), (C2) |
| (I3) | `shape` | 1-dimensional tensor constant of type `si64` | (C3) |
| (I4) | `rng_distribution` | enum of `UNIFORM` and `NORMAL` | (C2) |
#### Outputs
| Name | Type | Constraints |
|----------|----------------------------------------------------|-------------|
| `result` | tensor of integer, boolean, or floating-point type | (C1-C3) |
#### Constraints
* (C1) `element_type(a) = element_type(b) = element_type(result)`.
* (C2) If `rng_distribution = NORMAL`, then `is_float(a)`.
* (C3) `shape(result) = shape`.
#### Examples
```mlir
// %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
#### Semantics
Returns an `output` filled with uniform random bits and an updated output state
`output_state` using the pseudorandom number generator algorithm `rng_algorithm`
given an initial state `initial_state`. The output is guaranteed to be
deterministic function of `initial_state`, but it is not guaranteed to be
deterministic between implementations.
`rng_algorithm` is one of the following:
* `DEFAULT`: Implementation-defined algorithm.
* `THREE_FRY`: Implementation-defined variant of the Threefry algorithm.*
* `PHILOX`: Implementation-defined variant of the Philox algorithm.*
\* See: [Salmon et al. SC 2011. Parallel random numbers: as easy as 1, 2, 3.
](http://www.thesalmons.org/john/random123/papers/random123sc11.pdf)
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------------|----------------------------------------------|-------------|
| (I1) | `rng_algorithm` | enum of `DEFAULT`, `THREE_FRY`, and `PHILOX` | (C2) |
| (I2) | `initial_state` | 1-dimensional tensor of type `ui64` | (C1), (C2) |
#### Outputs
| Name | Type | Constraints |
|----------------|------------------------------------------|-------------|
| `output_state` | 1-dimensional tensor of type `ui64` | (C1) |
| `output` | tensor of integer or floating-point type | |
#### Constraints
* (C1) `type(initial_state) = type(output_state)`.
* (C2) `size(initial_state)` is defined as:
* implementation-defined if `rng_algorithm = DEFAULT`.
* `2` if `rng_algorithm = THREE_FRY`.
* `2` or `3` if `rng_algorithm = PHILOX`.
#### Examples
```mlir
// %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
#### Semantics
Performs element-wise rounding towards the nearest integer, breaking ties away
from zero, on the `operand` tensor and produces a `result` tensor. Implements
the `roundToIntegralTiesToAway` operation from the IEEE-754 specification. For
quantized types, performs
`dequantize_op_quantize(round_nearest_afz, operand, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|--------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|--------------------------------------------------------------|-------------|
| `result` | tensor of floating-point type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// %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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/round_nearest_afz.mlir)
### round_nearest_even
#### Semantics
Performs element-wise rounding towards the nearest integer, breaking ties
towards the even integer, on the `operand` tensor and produces a `result`
tensor. Implements the `roundToIntegralTiesToEven` operation from the IEEE-754
specification. For quantized types, performs
`dequantize_op_quantize(round_nearest_even, operand, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|--------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|--------------------------------------------------------------|-------------|
| `result` | tensor of floating-point type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// %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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/round_nearest_even.mlir)
### rsqrt
#### Semantics
Performs element-wise reciprocal square root operation on `operand` tensor and
produces a `result` tensor. Depending on the element type, does the following:
* For floats: `rSqrt` from IEEE-754.
* For complex numbers: complex reciprocal square root.
* For quantized types: `dequantize_op_quantize(rsqrt, operand, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|-------------------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------------------------------------------------|-------------|
| `result` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// %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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/rsqrt.mlir)
### scatter
#### Semantics
Produces `results` tensors which are equal to `inputs` tensors except that
several slices specified by `scatter_indices` are updated with the values
`updates` using `update_computation`.
The following diagram shows how elements in `updates...` map on elements in
`results...` using a concrete example. The diagram picks a few example
`updates...` indices and explains in detail which `results...` indices they
correspond to.
![scatter](images/spec/scatter.svg)
More formally, for all `update_index` in `index_space(updates[0])`:
* `update_scatter_dims = [d for d in axes(updates[0]) and d not in
update_window_dims]`.
* `update_scatter_index = update_index[update_scatter_dims...]`.
* `start_index` is defined as:
* `scatter_indices[si0, ..., :, ..., siN]` where `si` are individual
elements in `update_scatter_index` and `:` is inserted at the
`index_vector_dim` index, if `index_vector_dim` <
`rank(scatter_indices)`.
* `[scatter_indices[update_scatter_index]]` otherwise.
* For `d_input` in `axes(inputs[0])`,
* `full_start_index[d_input] = start_index[d_start]` if
`d_input = scatter_dims_to_operand_dims[d_start]`.
* `full_start_index[d_input] = 0` otherwise.
* For `d_input` in `axes(inputs[0])`,
* `full_batching_index[d_input] =
update_scatter_index[d_start - (d_start < index_vector_dim ? 0 : 1)]`
if `d_input = input_batching_dims[i_batching]` and
`d_start = scatter_indices_batching_dims[i_batching]`.
* `full_batching_index[d_input] = 0` otherwise.
* `update_window_index = update_index[update_window_dims...]`.
* `full_window_index = [wi0, ..., 0, ..., wiN]` where `wi` are individual
elements in `update_window_index`, and `0` is inserted at indices from
`inserted_window_dims` and `input_batching_dims`.
* `result_index = full_start_index + full_batching_index + full_window_index`.
Given that, `results = exec(schedule, inputs)`, where:
* `schedule` is an implementation-defined permutation of
`index_space(updates[0])`.
* `exec([update_index, ...], results) = exec([...], updated_results)` where:
* If `result_index` is in bounds for `shape(results...)`
* `updates_converted = to_destination_type(
updates...[update_index], type(func_inputs(update_computation)
[len(func_inputs(update_computation))//2:])... )`
* `updated_values = update_computation(results...[result_index],
updates_converted)`
* `updated_results` is a copy of `results` with `results...[result_index]`
set to `updated_values...`.
* Otherwise
* `updated_results = results`.
* `exec([], results) = results`.
If `indices_are_sorted` is `true` then the implementation can assume that
`scatter_indices` are sorted with respect to `scatter_dims_to_operand_dims`,
otherwise the behavior is undefined. More formally, for all `i1 < i2` from
`indices(result)`, `full_start_index(i1)` <= `full_start_index(i2)`.
If `unique_indices` is `true` then the implementation can assume that all
`result_index` indices being scattered to are unique. If `unique_indices` is
`true` but the indices being scattered to are not unique then the behavior is
undefined.
#### Inputs
| Label | Name | Type | Constraints |
|-------|---------------------------------|------------------------------------------------------------|------------------------------------------------------------|
| (I1) | `inputs` | variadic number of tensors or per-tensor quantized tensors | (C1), (C2), (C4-C6), (C11), (C13), (C18), (C21), (C23-C24) |
| (I2) | `scatter_indices` | tensor of integer type | (C4), (C15), (C19), (C22) |
| (I3) | `updates` | variadic number of tensors or per-tensor quantized tensors | (C3-C6), (C8) |
| (I4) | `update_window_dims` | 1-dimensional tensor constant of type `si64` | (C2), (C4), (C7-C8) |
| (I5) | `inserted_window_dims` | 1-dimensional tensor constant of type `si64` | (C2), (C4), (C9-C11) |
| (I6) | `input_batching_dims` | 1-dimensional tensor constant of type `si64` | (C2), (C4), (C9), (C12-13), (C17-18), (C20) |
| (I7) | `scatter_indices_batching_dims` | 1-dimensional tensor constant of type `si64` | (C14-C18) |
| (I8) | `scatter_dims_to_operand_dims` | 1-dimensional tensor constant of type `si64` | (C19-C21) |
| (I9) | `index_vector_dim` | constant of type `si64` | (C4), (C16), (C19), (C22) |
| (I10) | `indices_are_sorted` | constant of type `i1` | |
| (I11) | `unique_indices` | constant of type `i1` | |
| (I12) | `update_computation` | function | (C23) |
#### Outputs
| Name | Type | Constraints |
|-----------|------------------------------------------------------------|-------------|
| `results` | variadic number of tensors or per-tensor quantized tensors | (C24-C25) |
#### Constraints
<!-- markdownlint-disable line-length -->
* (C1) `same(shape(inputs...))`.
* (C2) `rank(inputs[0]) = size(update_window_dims) + size(inserted_window_dims) + size(input_batching_dims)`.
* (C3) `same(shape(updates...))`.
* (C4) `shape(updates[0]) = combine(update_scatter_dim_sizes,
update_window_dim_sizes)` where:
* `update_scatter_dim_sizes = shape(scatter_indices)` except that
the dimension size of `scatter_indices` corresponding to
`index_vector_dim` is not included.
* `update_window_dim_sizes <= shape(inputs[0])` except that
the dimension sizes in `inputs[0]` corresponding to `inserted_window_dims`
and `input_batching_dims` are not included.
* `combine` puts `update_scatter_dim_sizes` at axes corresponding to
`update_scatter_dims` and `update_window_dim_sizes` at axes corresponding
to `update_window_dims`.
* (C5) `0 < size(inputs) = size(updates) = N`.
* (C6) `element_type(updates...) = element_type(inputs...)`.
* (C7) `is_unique(update_window_dims) and is_sorted(update_window_dims)`.
* (C8) `0 <= update_window_dims < rank(updates[0])`.
* (C9) `is_unique(concatenate(inserted_window_dims, input_batching_dims))`
* (C10) `is_sorted(inserted_window_dims)`.
* (C11) `0 <= inserted_window_dims < rank(inputs[0])`.
* (C12) `is_sorted(input_batching_dims)`.
* (C13) `0 <= input_batching_dims < rank(inputs[0]))`.
* (C14) `is_unique(scatter_indices_batching_dims)`.
* (C15) `0 <= scatter_indices_batching_dims < rank(scatter_indices)`.
* (C16) `index_vector_dim not in scatter_indices_batching_dims`.
* (C17) `size(input_batching_dims) == size(scatter_indices_batching_dims)`.
* (C18) `dim(inputs[0], input_batching_dims...) =
dim(scatter_indices, scatter_indices_batching_dims...)`.
* (C19) `size(scatter_dims_to_operand_dims) =
index_vector_dim < rank(scatter_indices) ?
dim(scatter_indices, index_vector_dim) : 1`.
* (C20) `is_unique(concatenate(scatter_dims_to_operand_dims,
input_batching_dims))`.
* (C21) `0 <= scatter_dims_to_operand_dims < rank(inputs[0])`.
* (C22) `0 <= index_vector_dim <= rank(scatter_indices)`.
* (C23) `update_computation` has type `(tensor<E0>, ..., tensor<EN-1>,
tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>)`,
where `is_promotable(element_type(inputs[i]), Ei)`.
* (C24) `shape(inputs...) = shape(results...)`.
* (C25) `element_type(results[i]) = Ei` for all `i` in `[0,N)`.
<!-- markdownlint-enable line-length -->
#### Examples
```mlir
// %input: [
// [
// [[1, 2], [3, 4], [5, 6], [7, 8]],
// [[9, 10],[11, 12], [13, 14], [15, 16]],
// [[17, 18], [19, 20], [21, 22], [23, 24]]
// ],
// [
// [[25, 26], [27, 28], [29, 30], [31, 32]],
// [[33, 34], [35, 36], [37, 38], [39, 40]],
// [[41, 42], [43, 44], [45, 46], [47, 48]]
// ]
// ]
// %scatter_indices: [
// [
// [[0, 0], [1, 0], [2, 1]],
// [[0, 1], [1, 1], [0, 9]]
// ],
// [
// [[0, 0], [2, 1], [2, 2]],
// [[1, 2], [0, 1], [1, 0]]
// ]
// ]
// %update: [
// [
// [[[1, 1], [1, 1]], [[1, 1], [1, 1]], [[1, 1], [1, 1]]],
// [[[1, 1], [1, 1]], [[1, 1], [1, 1]], [[1, 1], [1, 1]]]
// ],
// [
// [[[1, 1], [1, 1]], [[1, 1], [1, 1]], [[1, 1], [1, 1]]],
// [[[1, 1], [1, 1]], [[1, 1], [1, 1]], [[1, 1], [1, 1]]]
// ]
// ]
%result = "stablehlo.scatter"(%input, %scatter_indices, %update) ({
^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
%0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
"stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
scatter_dimension_numbers = #stablehlo.scatter<
update_window_dims = [3, 4],
inserted_window_dims = [1],
input_batching_dims = [0],
scatter_indices_batching_dims = [1],
scatter_dims_to_operand_dims = [2, 1],
index_vector_dim = 3>,
indices_are_sorted = false,
unique_indices = false
} : (tensor<2x3x4x2xi64>, tensor<2x2x3x2xi64>, tensor<2x2x3x2x2xi64>) -> tensor<2x3x4x2xi64>
// %result: [
// [
// [[3, 4], [6, 7], [6, 7], [7, 8]],
// [[9, 10],[11, 12], [15, 16], [17, 18]],
// [[17, 18], [19, 20], [22, 23], [24, 25]]
// ],
// [
// [[25, 26], [28, 29], [30, 31], [31, 32]],
// [[35, 36], [38, 39], [38, 39], [39, 40]],
// [[41, 42], [44, 45], [46, 47], [47, 48]]
// ]
// ]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/scatter.mlir)
### select
#### Semantics
Produces a `result` tensor where each element is selected from `on_true` or
`on_false` tensor based on the value of the corresponding element of `pred`.
More formally, `result[result_index] = pred_element ? on_true[result_index] :
on_false[result_index]`, where `pred_element = rank(pred) = 0 ? pred[] :
pred[result_index]`. For quantized types, performs
`dequantize_select_quantize(pred, on_true, on_false, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|------------|---------------------------------------|-------------|
| (I1) | `pred` | tensor of type `i1` | (C1) |
| (I2) | `on_true` | tensor or per-tensor quantized tensor | (C1-C2) |
| (I3) | `on_false` | tensor or per-tensor quantized tensor | (C2) |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-------------|
| `result` | tensor or per-tensor quantized tensor | (C2) |
#### Constraints
* (C1) `rank(pred) = 0 or shape(pred) = shape(on_true)`.
* (C2) `baseline_type(on_true) = baseline_type(on_false) = baseline_type(result)`.
#### Examples
```mlir
// %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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/select.mlir)
### select_and_scatter
#### Semantics
Scatters the values from the `source` tensor using `scatter` based on the
outcome of `reduce_window` of the `input` tensor using `select` and produces
a `result` tensor.
The following diagram shows how elements in `result` are computed from
`operand` and `source` using a concrete example.
![select_and_scatter](images/spec/select_and_scatter.svg)
More formally:
* `selected_values = reduce_window_without_init(...)` with the following inputs:
* `inputs = [operand].`
* `window_dimensions`, `window_strides`, and `padding` which are used as is.
* `base_dilations = windows_dilations = 1`.
* `body` is defined as:
```python
def body(arg0: tensor<E>, arg1: tensor<E>) -> tensor<E>:
return select(arg0, arg1) ? arg0 : arg1;
```
where `E = element_type(operand)`, and `reduce_window_without_init` works
exactly like `reduce_window`, except that the `schedule` of the underlying
`reduce` (see [reduce](#reduce)) doesn't include init values. It is currently
unspecified what happens if the corresponding window doesn't have values
([#731](https://github.com/openxla/stablehlo/issues/731)).
* `result[result_index] = reduce([source_values], [init_value], [0], scatter)`
where:
* `source_values = [source[source_index] for source_index in
source_indices]`.
* `selected_index(source_index) = operand_index` if
`selected_values[source_index]` has the `operand` element
from `operand_index`.
* `source_indices = [source_index for source_index in
indices(source) if selected_index(source_index) = result_index]`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|---------------------|-----------------------------------------------------|-------------------------|
| (I1) | `operand` | tensor or per-tensor quantized tensor | (C1-C4), (C6), (C8-C11) |
| (I2) | `source` | tensor or per-tensor quantized tensor | (C1), (C2) |
| (I3) | `init_value` | 0-dimensional tensor or per-tensor quantized tensor | (C3) |
| (I4) | `window_dimensions` | 1-dimensional tensor constant of type `si64` | (C2), (C4), (C5) |
| (I5) | `window_strides` | 1-dimensional tensor constant of type `si64` | (C2), (C6), (C7) |
| (I6) | `padding` | 2-dimensional tensor constant of type `si64` | (C2), (C8) |
| (I7) | `select` | function | (C9) |
| (I8) | `scatter` | function | (C10) |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-------------|
| `result` | tensor or per-tensor quantized tensor | (C11-C12) |
#### Constraints
<!-- markdownlint-disable line-length -->
* (C1) `element_type(operand) = element_type(source)`.
* (C2) `shape(source) = num_windows` where:
* `padded_operand_shape = padding[:, 0] + shape(operand) + padding[:, 1]`.
* `is_empty_window = padded_operand_shape = 0 || window_dimensions > padded_operand_shape`.
* `num_windows = is_empty_window ? 0 : floor((padded_operand_shape - window_dimensions) / window_strides) + 1`.
* (C3) `element_type(init_value) = element_type(operand)`.
* (C4) `size(window_dimensions) = rank(operand)`.
* (C5) `0 < window_dimensions`.
* (C6) `size(window_strides) = rank(operand)`.
* (C7) `0 < window_strides`.
* (C8) `shape(padding) = [rank(operand), 2]`.
* (C9) `select` has type `(tensor<E>, tensor<E>) -> tensor<i1>` where
`E = element_type(operand)`.
* (C10) `scatter` has type `(tensor<E>, tensor<E>) -> tensor<E>` where
`is_promotable(element_type(operand), E)`.
* (C11) `shape(operand) = shape(result)`.
* (C12) `element_type(result) = E`.
<!-- markdownlint-enable line-length -->
#### Examples
```mlir
// %operand: [[1, 5], [2, 5], [3, 6], [4, 4]]
// %source: [[5, 6], [7, 8]]
// %init_value: 0
%result = "stablehlo.select_and_scatter"(%operand, %source, %init_value) ({
^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
%0 = "stablehlo.compare"(%arg0, %arg1) {
comparison_direction = #stablehlo<comparison_direction GE>
} : (tensor<i64>, tensor<i64>) -> tensor<i1>
"stablehlo.return"(%0) : (tensor<i1>) -> ()
}, {
^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
%0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
"stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
window_dimensions = array<i64: 3, 1>,
window_strides = array<i64: 2, 1>,
padding = dense<[[0, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<4x2xi64>, tensor<2x2xi64>, tensor<i64>) -> tensor<4x2xi64>
// %result: [[0, 0], [0, 0], [5, 14], [7, 0]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/select_and_scatter.mlir)
### send
#### Semantics
Sends `inputs` to a channel `channel_id`. Inputs are then sent to other devices
in the order specified by `source_target_pairs`. The operation produces a
`result` token.
If `is_host_transfer` is `true`, then the operation transfers data to the
host. Otherwise, it transfers data to another device based on the values of
`source_target_pairs`. This flag duplicates the information provided in
`channel_type`, so in the future we are planning to only keep one of them
([#666](https://github.com/openxla/stablehlo/issues/666)). If `is_host_transfer`
= `false` and `source_target_pairs` is `None` or empty, it is considered
undefined behavior.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------------------|-------------------------------------------------|---------------|
| (I1) | `inputs` | variadic number of tensors or quantized tensors | |
| (I2) | `token` | `token` | |
| (I3) | `channel_id` | constant of type `si64` | |
| (I4) | `channel_type` | enum of `DEVICE_TO_DEVICE` and `DEVICE_TO_HOST` | (C5) |
| (I5) | `is_host_transfer` | constant of type `i1` | (C5-C6) |
| (I6) | `source_target_pairs` | 2-dimensional tensor constant of type `si64` | (C1-C4), (C6) |
#### Outputs
| Name | Type |
|----------|---------|
| `result` | `token` |
#### 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) `channel_type` is defined as:
* `DEVICE_TO_HOST` if `is_host_transfer = true`,
* `DEVICE_TO_DEVICE` otherwise.
#### Examples
```mlir
%result = "stablehlo.send"(%operand, %token) {
channel_handle = #stablehlo.channel_handle<handle = 0, type = 1>,
is_host_transfer = false,
source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>
} : (tensor<2x2xi64>, !stablehlo.token) -> !stablehlo.token
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/send_recv.mlir)
### shift_left
#### Semantics
Performs element-wise left-shift operation on the `lhs` tensor by `rhs` number
of bits and produces a `result` tensor.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------|------------------------|-------------|
| (I1) | `lhs` | tensor of integer type | (C1) |
| (I2) | `rhs` | tensor of integer type | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|------------------------|-------------|
| `result` | tensor of integer type | (C1) |
#### Constraints
* (C1) `type(lhs) = type(rhs) = type(result)`.
#### Examples
```mlir
// %lhs: [-1, 0, 1]
// %rhs: [1, 2, 3]
%result = "stablehlo.shift_left"(%lhs, %rhs): (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
// %result: [-2, 0, 8]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/shift_left.mlir)
### shift_right_arithmetic
#### Semantics
Performs element-wise arithmetic right-shift operation on the `lhs` tensor by
`rhs` number of bits and produces a `result` tensor.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------|------------------------|-------------|
| (I1) | `lhs` | tensor of integer type | (C1) |
| (I2) | `rhs` | tensor of integer type | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|------------------------|-------------|
| `result` | tensor of integer type | (C1) |
#### Constraints
* (C1) `type(lhs) = type(rhs) = type(result)`.
#### Examples
```mlir
// %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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/shift_right_arithmetic.mlir)
### shift_right_logical
#### Semantics
Performs element-wise logical right-shift operation on the `lhs` tensor by `rhs`
number of bits and produces a `result` tensor.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------|------------------------|-------------|
| (I1) | `lhs` | tensor of integer type | (C1) |
| (I2) | `rhs` | tensor of integer type | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|------------------------|-------------|
| `result` | tensor of integer type | (C1) |
#### Constraints
* (C1) `type(lhs) = type(rhs) = type(result)`.
#### Examples
```mlir
// %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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/shift_right_logical.mlir)
### sign
#### Semantics
Returns the sign of the `operand` element-wise and produces a `result` tensor.
More formally, for each element `x`, the semantics can be expressed using
Python syntax as follows:
```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)))
```
For quantized types, performs
`dequantize_op_quantize(sign, operand, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|------------------------------------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of signed integer, floating-point, or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|------------------------------------------------------------------------------------------|-------------|
| `result` | tensor of signed integer, floating-point, or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// 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]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/sign.mlir)
### sine
#### Semantics
Performs element-wise sine operation on `operand` tensor and produces a `result`
tensor. Depending on the element type, does the following:
* For floats: `sin` from IEEE-754.
* For complex numbers: complex sine.
* For quantized types: `dequantize_op_quantize(sine, operand, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|-------------------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------------------------------------------------|-------------|
| `result` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// %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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/sine.mlir)
### slice
#### Semantics
Extracts a slice from the `operand` using statically-computed starting indices
and produces a `result` tensor. `start_indices` contain the starting indices of
the slice for each dimension, `limit_indices` contain the ending indices
(exclusive) for the slice for each dimension, and `strides` contain the strides
for each dimension.
More formally, `result[result_index] = operand[operand_index]` where
`operand_index = start_indices + result_index * strides`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------------|----------------------------------------------|------------------|
| (I1) | `operand` | tensor or per-tensor quantized tensor | (C1-C3), (C5) |
| (I2) | `start_indices` | 1-dimensional tensor constant of type `si64` | (C2), (C3), (C5) |
| (I3) | `limit_indices` | 1-dimensional tensor constant of type `si64` | (C2), (C3), (C5) |
| (I4) | `strides` | 1-dimensional tensor constant of type `si64` | (C2), (C4) |
#### Outputs
| Name | Type | Constraints |
|----------|---------------------------------------|-------------|
| `result` | tensor or per-tensor quantized tensor | (C1), (C5) |
#### Constraints
* (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)`.
#### Examples
```mlir
// %operand: [
// [0, 0, 0, 0],
// [0, 0, 1, 1],
// [0, 0, 1, 1]
// ]
%result = "stablehlo.slice"(%operand) {
start_indices = array<i64: 1, 2>,
limit_indices = array<i64: 3, 4>,
strides = array<i64: 1, 1>
} : (tensor<3x4xi64>) -> tensor<2x2xi64>
// % result: [
// [1, 1],
// [1, 1]
// ]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/slice.mlir)
### sort
#### Semantics
Sorts 1-dimensional slices of `inputs` along the dimension `dimension` together,
according to a `comparator` and produces `results`.
Unlike similar inputs in other operations, `dimension` allows negative values,
with the semantics described below. In the future, this may be disallowed
for consistency reasons
([#1377](https://github.com/openxla/stablehlo/issues/1377)).
If `is_stable` is true, then the sorting is stable, that is, relative order of
elements considered to be equal by the comparator is preserved. For the case
where there is a single input, two elements `e1` and `e2` are considered to be
equal by the comparator if and only if
`comparator(e1, e2) = comparator(e2, e1) = false`. See the formalization below
for how this generalizes to multiple inputs.
More formally, for all `result_index` in `index_space(results[0])`:
* `adjusted_dimension = dimension >= 0 ? dimension : rank(inputs[0]) + dimension`.
* `result_slice = [ri0, ..., :, ..., riR-1]` where `riN` are individual
elements in `result_index`, and `:` is inserted at `adjusted_dimension`.
* `inputs_together = (inputs[0]..., ..., inputs[N-1]...)`.
* `results_together[result_slice] = sort(inputs_together[result_slice], comparator_together)`.
* where `sort` sorts a 1-dimensional slice in non-descending order expecting
that `comparator_together` returns `true` if the left-hand side argument is
less than the right-hand second argument.
* &#32;
```python
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`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|--------------|------------------------------------------------------------|-------------|
| (I1) | `inputs` | variadic number of tensors or per-tensor quantized tensors | (C1-C5) |
| (I2) | `dimension` | constant of type `si64` | (C4) |
| (I3) | `is_stable` | constant of type `i1` | |
| (I4) | `comparator` | function | (C5) |
#### Outputs
| Name | Type | Constraints |
|-----------|------------------------------------------------------------|-------------|
| `results` | variadic number of tensors or per-tensor quantized tensors | (C2), (C3) |
#### Constraints
* (C1) `0 < size(inputs)`.
* (C2) `type(inputs...) = type(results...)`.
* (C3) `same(shape(inputs...) + shape(results...))`.
* (C4) `-R <= dimension < R`, where `R = rank(inputs[0])`.
* (C5) `comparator` has type
`(tensor<E1>, tensor<E1>, ..., tensor<EN-1>, tensor<EN-1>) -> tensor<i1>`,
where `Ei = element_type(inputs[i])`.
#### Examples
```mlir
// %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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/sort.mlir)
### sqrt
#### Semantics
Performs element-wise square root operation on `operand` tensor and produces a
`result` tensor. Depending on the element type, does the following:
* For floats: `squareRoot` from IEEE-754.
* For complex numbers: complex square root.
* For quantized types: `dequantize_op_quantize(sqrt, operand, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|-------------------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------------------------------------------------|-------------|
| `result` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// %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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/sqrt.mlir)
### subtract
#### Semantics
Performs element-wise subtraction of two tensors `lhs` and `rhs` and produces a
`result` tensor. Depending on the element type, does the following:
* For integers: integer subtraction.
* For floats: `subtraction` from IEEE-754.
* For complex numbers: complex subtraction.
* For quantized types:
* `dequantize_op_quantize(subtract, lhs, rhs, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------|-----------------------------------------------------------------------------------|-------------|
| (I1) | `lhs` | tensor of integer, floating-point, or complex type or per-tensor quantized tensor | (C1) |
| (I2) | `rhs` | tensor of integer, floating-point, or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-----------------------------------------------------------------------------------|-------------|
| `result` | tensor of integer, floating-point, or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(lhs) = baseline_type(rhs) = baseline_type(result)`.
#### Examples
```mlir
// %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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/subtract.mlir)
### tan
#### Semantics
Performs element-wise tangent operation on the `operand` tensor and produces a
`result` tensor. Depending on the element type, does the following:
* For floats: `tan` from IEEE-754.
* For complex numbers: complex tangent.
* For quantized types: `dequantize_op_quantize(tan, operand, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|-------------------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------------------------------------------------|-------------|
| `result` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// %operand: [
// [0.0, 1.57079632], // [0, pi/2]
// [3.14159265, 4.71238898] // [pi, 3pi/2]
// ]
%result = "stablehlo.tan"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
// %result: [
// [0.0, 1.63312e+16],
// [0.0, 5.44375e+15]
// ]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/tan.mlir)
### tanh
#### Semantics
Performs element-wise hyperbolic tangent operation on `operand` tensor and
produces a `result` tensor. Depending on the element type, does the following:
* For floats: `tanh` from IEEE-754.
* For complex numbers: complex hyperbolic tangent.
* For quantized types:
* `dequantize_op_quantize(tanh, operand, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|-------------------------------------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------------------------------------------------|-------------|
| `result` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_type(operand) = baseline_type(result)`.
#### Examples
```mlir
// %operand: [-1.0, 0.0, 1.0]
%result = "stablehlo.tanh"(%operand) : (tensor<3xf32>) -> tensor<3xf32>
// %result: [-0.76159416, 0.0, 0.76159416]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/tanh.mlir)
### transpose
#### Semantics
Permutes the dimensions of `operand` tensor using `permutation` and produces a
`result` tensor. More formally, `result[result_index] = operand[operand_index]`
where `result_index[d] = operand_index[permutation[d]]`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|---------------|----------------------------------------------|-------------|
| (I1) | `operand` | tensor or quantized tensor | (C1-C4) |
| (I2) | `permutation` | 1-dimensional tensor constant of type `si64` | (C2-C4) |
#### Outputs
| Name | Type | Constraints |
|----------|----------------------------|---------------|
| `result` | tensor or quantized tensor | (C1), (C3-C4) |
#### 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)` and
`quantization_dimension(result)` may differ, otherwise.
* (C2) `permutation` is a permutation of `range(rank(operand))`.
* (C3) `shape(result) = dim(operand, permutation...)`.
* (C4) If `is_per_axis_quantized(result)`, then
`quantization_dimension(operand) =
permutation(quantization_dimension(result))`.
#### Examples
```mlir
// %operand: [
// [[1,2], [3,4], [5,6]],
// [[7,8], [9,10], [11,12]]
// ]
%result = "stablehlo.transpose"(%operand) {
permutation = array<i64: 2, 1, 0>
} : (tensor<2x3x2xi32>) -> tensor<2x3x2xi32>
// %result: [
// [[1,7], [3,9], [5,11]],
// [[2,8], [4,10], [6,12]]
// ]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/transpose.mlir)
### triangular_solve
#### Semantics
Solves batches of systems of linear equations with lower or upper triangular
coefficient matrices.
More formally, given `a` and `b`, `result[i0, ..., iR-3, :, :]` is the solution
to `op(a[i0, ..., iR-3, :, :]) * x = b[i0, ..., iR-3, :, :]` when `left_side` is
`true` or `x * op(a[i0, ..., iR-3, :, :]) = b[i0, ..., iR-3, :, :]` when
`left_side` is `false`, solving for the variable `x` where `op(a)` is determined
by `transpose_a`, which can be one of the following:
* `NO_TRANSPOSE`: Perform operation using `a` as-is.
* `TRANSPOSE`: Perform operation on transpose of `a`.
* `ADJOINT`: Perform operation on conjugate transpose of `a`.
Input data is read only from the lower triangle of `a`, if `lower` is `true` or
upper triangle of `a`, otherwise. Output data is returned in the same triangle;
the values in the other triangle are implementation-defined.
If `unit_diagonal` is true, then the implementation can assume that the diagonal
elements of `a` are equal to 1, otherwise the behavior is undefined.
For quantized types, performs
`dequantize_op_quantize(lambda x, y: triangular_solve(x, y, left_side, lower,
unit_diagonal, transpose_a), a, b, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------------|-------------------------------------------------------------------------|-------------|
| (I1) | `a` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1-C3) |
| (I2) | `b` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1-C4) |
| (I3) | `left_side` | constant of type `i1` | (C3) |
| (I4) | `lower` | constant of type `i1` | |
| (I5) | `unit_diagonal` | constant of type `i1` | |
| (I6) | `transpose_a` | enum of `NO_TRANSPOSE`, `TRANSPOSE`, and `ADJOINT` | |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------------------------------------------------|-------------|
| `result` | tensor of floating-point or complex type or per-tensor quantized tensor | (C1) |
#### Constraints
* (C1) `baseline_element_type(a) = baseline_element_type(b)`.
* (C2) `2 <= rank(a) = rank(b) = R`.
* (C3) The relationship between `shape(a)` and `shape(b)` is defined as follows:
* `shape(a)[:-3] = shape(b)[:-3]`.
* `dim(a, -2) = dim(a, -1) = dim(b, left_side ? -2 : -1)`.
* (C4) `baseline_type(b) = baseline_type(result)`.
#### Examples
```mlir
// %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
> Note: Per [StableHLO v1.0 Cleanup #2283](https://github.com/openxla/stablehlo/pull/2283),
> this op is being explored for deprecation as it appears to be unused by both
> frameworks and compilers. As such, it has limited compatibility guarantees
> (6 months).
#### Semantics
Produces a `result` tuple from values `val`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------|---------------------------|-------------|
| (I1) | `val` | variadic number of values | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-------|-------------|
| `result` | tuple | (C1) |
#### Constraints
* (C1) `result` has type `tuple<E0, ..., EN-1>` where `Ei = type(val[i])`.
#### Examples
```mlir
// %val0: memref[1.0, 2.0]
// %val1: (3)
%result = "stablehlo.tuple"(%val0, %val1) : (memref<2xf32>, tuple<tensor<i32>>) -> tuple<memref<2xf32>, tuple<tensor<i32>>>
// %result: (memref[1.0, 2.0], (3))
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/tuple_and_get_tuple_element.mlir)
### uniform_dequantize
#### Semantics
Performs element-wise conversion of quantized tensor `operand` to a
floating-point tensor `result` according to the quantization parameters defined
by the `operand` type.
More formally, `result = dequantize(operand)`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|------------------|-------------|
| (I1) | `operand` | quantized tensor | (C1), (C2) |
#### Outputs
| Name | Type | Constraints |
|----------|-------------------------------|-------------|
| `result` | tensor of floating-point type | (C1), (C2) |
#### Constraints
* (C1) `shape(operand) = shape(result)`.
* (C2) `element_type(result) = expressed_type(operand)`.
#### Examples
```mlir
// %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
#### Semantics
Performs element-wise conversion of floating-point tensor or quantized tensor
`operand` to a quantized tensor `result` according to the quantization
parameters defined by the `result` type.
More formally,
* If `is_float(operand)`:
* `result = quantize(operand, type(result))`.
* If `is_quantized(operand)`:
* `float_result = dequantize(operand)`.
* `result = quantize(float_result, type(result))`.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|--------------------------------------------|-------------|
| (I1) | `operand` | tensor of floating-point or quantized type | (C1), (C2) |
#### Outputs
| Name | Type | Constraints |
|----------|------------------|-------------|
| `result` | quantized tensor | (C1), (C2) |
#### Constraints
* (C1) `shape(operand) = shape(result)`.
* (C2) `expressed_type(result) = is_float(operand) ? element_type(operand) :
expressed_type(operand)`.
#### Examples
```mlir
// %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]
```
### while
#### Semantics
Produces the output from executing `body` function 0 or more times while the
`cond` function outputs `true`. More formally, the semantics can be expressed
using Python syntax as follows:
```python
internal_state = operand
while cond(*internal_state):
internal_state = body(*internal_state)
results = internal_state
```
The behavior of an infinite loop is TBD
([#383](https://github.com/openxla/stablehlo/issues/383)).
#### Inputs
| Label | Name | Type | Constraints |
|-------|-----------|---------------------------|-------------|
| (I1) | `operand` | variadic number of values | (C1-C3) |
| (I2) | `cond` | function | (C1) |
| (I3) | `body` | function | (C2) |
#### Outputs
| Name | Type | Constraints |
|-----------|---------------------------|-------------|
| `results` | variadic number of values | (C3) |
#### Constraints
* (C1) `cond` has type `(T0, ..., TN-1) -> tensor<i1>`, where
`Ti = type(operand[i])`.
* (C2) `body` has type `(T0, ..., TN-1) -> (T0, ..., TN-1)`, where
`Ti = type(operand[i])`.
* (C3) `type(results...) = type(operand...)`.
#### Examples
```mlir
// %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
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/while.mlir)
### xor
#### Semantics
Performs element-wise XOR of two tensors `lhs` and `rhs` and produces a `result`
tensor. Depending on the element type, does the following:
* For booleans: logical XOR.
* For integers: bitwise XOR.
#### Inputs
| Label | Name | Type | Constraints |
|-------|-------|-----------------------------------|-------------|
| (I1) | `lhs` | tensor of boolean or integer type | (C1) |
| (I2) | `rhs` | tensor of boolean or integer type | (C1) |
#### Outputs
| Name | Type | Constraints |
|----------|-----------------------------------|-------------|
| `result` | tensor of boolean or integer type | (C1) |
#### Constraints
* (C1) `type(lhs) = type(rhs) = type(result)`.
#### Examples
```mlir
// 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]]
```
&nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/xor.mlir)
## Dialect Interop
At the moment, StableHLO programs in the wild sometimes contain operations that
are not defined by StableHLO.
### Module, Function, Call and Return
StableHLO uses upstream MLIR operations for ModuleOp, FuncOp, CallOp, and
ReturnOp. This was done for better interop with existing MLIR machinery, as many
useful passes are written targeting FuncOp and ModuleOp, and many compilation
pipelines expect these ops to be present. Full compatibility guarantees are
applied to these ops. If anything ever changes about these ops in an
incompatible way (i.e. removal), StableHLO equivalents will be added to preserve
compatibility.
### CHLO
The CHLO opset contains higher level operations that decompose to StableHLO.
Currently there are no compatibility guarantees for CHLO. For compatibility
guarantees, the [chlo-legalize-to-stablehlo pass](https://github.com/openxla/stablehlo/blob/12fd0a9e7b3c6f3dea3defc513870c962e62726d/stablehlo/transforms/Passes.td#L119)
must be used prior to serialization.
### Shape Operations
It is a common use case in the community to use certain operations from core
MLIR dialects in dynamic StableHLO programs to perform shape computations.
Most commonly, these include [`shape` dialect](https://mlir.llvm.org/docs/Dialects/ShapeDialect/)
ops like `shape_of` or `num_elements`, [`tensor` dialect](https://mlir.llvm.org/docs/Dialects/TensorOps/)
ops like `dim` or `from_elements`, and the builtin `index` type.
The [Dynamism RFC > O2](https://github.com/openxla/stablehlo/blob/main/rfcs/20230704-dynamism-101.md#o2)
denotes these as out of scope, however some support for `index` types is
included for interop purposes. There are no compatibility guarantees for these
ops or types. The [shape-legalize-to-stablehlo](https://github.com/openxla/stablehlo/blob/12fd0a9e7b3c6f3dea3defc513870c962e62726d/stablehlo/transforms/Passes.td#L136)
pass can be used to convert these operations to fully supported StableHLO ops.
## Deprecated Operations
There are several StableHLO operations that were inherited from
[MHLO](https://github.com/openxla/xla/blob/d63deb9250b9c212445290bd08c6effb5b6d0a2b/xla/mlir_hlo/mhlo/IR/hlo_ops.td)
which are deprecated and on the way out of StableHLO. The full details on these
removals can be found in the [StableHLO v1.0 Cleanup #2283](https://github.com/openxla/stablehlo/pull/2283).
The tracker issue for these deprecations is [#2340](https://github.com/openxla/stablehlo/issues/2340).
These operations fall into a few categories:
* "Not in HLO" category of StableHLO operations - they were initially part of
the StableHLO opset but have been later deemed to not fit it well:
`broadcast`, `create_token`, `cross-replica-sum`, `dot`, `einsum`,
`torch_index_select`, `unary_einsum`
([#3](https://github.com/openxla/stablehlo/issues/3)).
* Unused ops - These operations may have been useful at some point, but the ops
were either underdeveloped, or the pipelines using these ops have been
refactored to not require them anymore. This includes `map`, `tuple` ([#598](https://github.com/openxla/stablehlo/issues/598)),
`get_tuple_element`, `rng`, `complex` comparisons [#560](https://github.com/openxla/stablehlo/issues/560),
and convolution `window_reversal` ([#1181](https://github.com/openxla/stablehlo/issues/1181)).
Some of these ops can be removed easily given that they can be expressed using
existing ops (`broadcast`, `create_token`, `cross-replica-sum`, `dot`,
`unary_einsum`) and will be removed after the existing compatibilty window
passes (6 months). Others are still being explored for removal (`einsum`,
`get_tuple_element`, `map`, `rng` `torch_index_select`, `tuple`, `complex`
comparisons, `window_reversal`). Pending community feedback,
these ops will either be removed, or added to the spec with full support. Until
these ops futures are known, they are only guaranteed 6 months of compatibility.
## Execution
### Sequential execution
A StableHLO program is executed by providing input values to the `main` function
and computing output values. Output values of a function are computed by
executing the graph of ops rooted in the corresponding `return` op.
The execution order is implementation-defined as long as it is aligned with
dataflow, i.e. if ops are executed before their uses. In StableHLO, all
side-effecting ops consume one token and produce one token (multiple tokens can
be multiplexed into one token via `after_all`), so the execution order of side
effects is also aligned with dataflow. For example, in the below program
there are two possible execution orders: `%0` → `%1` → `%2` → `return` and
`%1` → `%0` → `%2` → `return`.
```mlir
func.func @main() -> tensor<f64> {
%0 = stablehlo.constant dense<1.0> : tensor<f64>
%1 = stablehlo.constant dense<2.0> : tensor<f64>
%2 = stablehlo.add %0, %1 : tensor<f64>
return %2 : tensor<f64>
}
```
More formally, a **StableHLO process** is a combination of:
1\) a StableHLO program, 2) operation statuses (not executed yet,
already executed), and 3) intermediate values that the process is working on.
The process starts with input values to the `main` function, progresses through
the graph of ops updating operation statuses and intermediate values and
finishes with output values. Further formalization is TBD
([#484](https://github.com/openxla/stablehlo/issues/484)).
### Parallel execution
StableHLO programs can be executed in parallel, organized into a 2D process grid
of `num_replicas` by `num_partitions` which both have type `ui32`.
In the **StableHLO process grid**, `num_replicas * num_partitions` of StableHLO
processes are executing at the same time. Each process has a unique
`process_id = (replica_id, partition_id)`, where
`replica_id` in `replica_ids = range(num_replicas)` and
`partition_id` in `partition_ids = range(num_partitions)` which both have
type `ui32`.
The size of the process grid is known statically for every program (in the
future, we are planning to make it an explicit part of StableHLO programs
[#650](https://github.com/openxla/stablehlo/issues/650)), and the position
within the process grid is known statically for every process. Each process has
access to its position within the process grid via the `replica_id` and
`partition_id` ops.
Within the process grid, the programs can all be the same (in the "Single
Program, Multiple Data" style), can all be different (in the "Multiple Program,
Multiple Data" style) or something in between. In the future, we are planning
to introduce support for other idioms of defining parallel StableHLO programs,
including GSPMD ([#619](https://github.com/openxla/stablehlo/issues/619)).
Within the process grid, the processes are mostly independent from each other -
they have separate operation statuses, separate input/intermediate/output values
and most of the ops are executed separately between processes, with the
exception of a small number of collective ops described below.
Given that execution of most of the ops is only using values from the same
process, it is usually unambiguous to refer to these values by their names.
However, when describing semantics of collective ops, that is insufficient, and
that gives rise to the notation `name@process_id` to refer to the value `name`
within a particular process. (From that perspective, unqualified `name` can be
viewed as a shorthand for `name@(replica_id(), partition_id())`).
The execution order across processes is implementation-defined, except for the
synchronization introduced by point-to-point communication and collective ops
as described below.
### Point-to-point communication
StableHLO processes can communicate with each other through
**StableHLO channels**. A channel is represented by a positive id of type
`si64`. Through various ops, it is possible to send values to channels and
receive them from channels.
Further formalization, e.g. where these channel ids are coming from, how
processes programs become aware of them and what kind of synchronization is
introduced by them, is TBD
([#484](https://github.com/openxla/stablehlo/issues/484)).
### Streaming communication
Every StableHLO process has access to two streaming interfaces:
* **Infeed** that can be read from.
* **Outfeed** that can be written to.
Unlike channels, which are used to communicate between processes and therefore
have processes at both of their ends, infeeds and outfeeds have their other
end implementation-defined.
Further formalization, e.g. how streaming communication influences execution
order and what kind of synchronization is introduced by it, is TBD
([#484](https://github.com/openxla/stablehlo/issues/484)).
### Collective ops
There are six collective ops in StableHLO: `all_gather`, `all_reduce`,
`all_to_all`, `collective_broadcast`, `collective_permute`, `reduce_scatter`
and `collective_reduce`. All these ops split the processes in the StableHLO process
grid into **StableHLO process groups** and execute a joint computation within
each process group, independently from other process groups.
Within each process group, collective ops may introduce a synchronization
barrier. Further formalization, e.g. elaborating on when exactly this
synchronization happens, how exactly the processes arrive at this barrier,
and what happens if they don't, is TBD
([#484](https://github.com/openxla/stablehlo/issues/484)).
If the process group involves cross-partition communication, i.e. there are
processes in the process group whose partition ids are different, then execution
of the collective op needs a channel, and the collective op must provide a
positive `channel_id` of type `si64`. Cross-replica communication doesn't need
channels.
The computations performed by the collective ops are specific to individual ops
and are described in individual op sections above. However, the strategies by
which the process grid is split into process groups are shared between these ops
and are described in this section. More formally, StableHLO supports the
following four strategies.
#### Replica groups
Collective operations use `replica_groups` to define the groups of devices the
computation will be executed on. The `replica_groups` can be either a tensor
constant of type `si64` (representing a list of lists of ids) or a
`ReplicaGroupMeshAxes` object.
```mlir
// Standard replica groups constant (list of lists of ids)
dense<[[0, 1], [2, 3]]> : tensor<2x2xi64>
// Symbol reference mesh (Post-lifting)
#stablehlo.replica_group_mesh_axes<
mesh = @mesh,
axes = [
#stablehlo.axis_ref<name = "x">,
#stablehlo.axis_ref<name = "y">
]
>
// Inlined mesh (Prior to lifting)
#stablehlo.replica_group_mesh_axes<
mesh = #stablehlo.mesh<
axes = [
#stablehlo.mesh_axis<name = "x", size = 2>,
#stablehlo.mesh_axis<name = "y", size = 2>
]
>,
axes = [
#stablehlo.axis_ref<name = "x", sub_axis_info = (1)2>,
#stablehlo.axis_ref<name = "y">
]
>
```
##### StableHLO Mesh
A **StableHLO mesh** defines the topology of the devices the StableHLO process
grid is mapped to. A mesh is defined by `MeshAxis` and `device_ids` attributes.
In addition to a mesh definition, `ReplicaGroupMeshAxes`, `AxisRef`, and
`SubAxisInfo` are used to define the replica groups.
A `Mesh` defines a named grid of devices. `Mesh` fields include:
* `axes`, a list of `MeshAxis` objects representing the dimensions of the
device grid.
* `device_ids`, an optional 1-dimensional tensor constant of type `si64`
representing the mapping of grid coordinates to device IDs. If not
specified, device IDs are assigned in the lexicographical order of
coordinates.
A `MeshAxis` defines a single dimension of a mesh. `MeshAxis` fields include:
* `name`, a constant of type `string` representing the name of this axis.
* `size`, a constant of type `si64` representing the size of this axis.
A `ReplicaGroupMeshAxes` defines the replica groups based on a StableHLO mesh.
`ReplicaGroupMeshAxes` fields include:
* `mesh`, an `attribute` representing the StableHLO mesh to be used (either an
inlined `#stablehlo.mesh` or a symbol reference to a mesh definition).
* `axes`, a list of `AxisRef` elements.
An `AxisRef` defines a reference to either a full axis or a split sub-axis.
`AxisRef` fields include:
* `name`, a constant of type `string` representing the name of this axis.
* `sub_axis_info`, an optional `SubAxisInfo` object that provides additional
information if this is a sub axis.
A `SubAxisInfo` defines how this sub-axis is derived from the full axis.
`SubAxisInfo` fields include:
* `pre_size`, a constant of type `si64` representing the product of sizes of
sub-axes that come before this one in the split.
* `size`, a constant of type `si64` representing the size of this sub-axis.
#### cross_replica
Only cross-replica communications happen within each process group. This
strategy takes `replica_groups` and computes a Cartesian product of
`replica_groups` by `partition_ids`. `replica_groups`
must have unique elements and cover all `replica_ids`. More formally, using
Python syntax:
```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
```
For example, for `replica_groups = [[0, 1], [2, 3]]` and `num_partitions = 2`,
`cross_replica` will produce
`[[(0, 0), (1, 0)], [(0, 1), (1, 1)], [(2, 0), (3, 0)], [(2, 1), (3, 1)]]`.
#### cross_partition
Only cross-partition communications happen within each process group. This
strategy takes `partition_groups` and
computes a Cartesian product of `partition_groups` by `replica_ids`.
`partition_groups` must have unique elements and cover all `partition_ids`.
More formally, using Python syntax:
```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
```
For example, for `partition_groups = [[0, 1]]` and `num_replicas = 4`,
`cross_partition` will produce
`[[(0, 0), (0, 1)], [(1, 0), (1, 1)], [(2, 0), (2, 1)], [(3, 0), (3, 1)]]`.
#### cross_replica_and_partition
Both cross-replica and cross-partition communications may happen within each
process group. This strategy takes `replica_groups` and computes Cartesian
products of each `replica_group` by
`partition_ids`. `replica_groups` must have unique elements and cover all
`replica_ids`. More formally, using Python syntax:
```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
```
For example, for `replica_groups = [[0, 1], [2, 3]]` and `num_partitions = 2`,
`cross_replica_and_partition` will produce
`[[(0, 0), (1, 0), (0, 1), (1, 1)], [(2, 0), (3, 0), (2, 1), (3, 1)]]`.
#### flattened_ids
This strategy takes `flattened_id_groups` - a list of lists of "flattened"
process ids in the form of `replica_id * num_partitions + partition_id` - and
turns them into process ids. `flattened_id_groups` must have unique elements
and cover all `process_ids`. More formally, using Python syntax:
```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
```
For example, for `flattened_id_groups = [[0, 1, 2, 3], [4, 5, 6, 7]]`,
`num_replicas = 4` and `num_partitions = 2`, `flattened_ids` will produce
`[[(0, 0), (0, 1), (1, 0), (1, 1)], [(2, 0), (2, 1), (3, 0), (3, 1)]]`.
### Accuracy
At the moment, StableHLO does not provide guarantees about numerical accuracy,
but this may change in the future
([#1156](https://github.com/openxla/stablehlo/issues/1156)).
### Execution semantics of quantized operation
The interpretation of quantized StableHLO operations may vary depending on the
hardware requirements and capabilities. For instance, some hardware may opt to
interpret quantized operations using a "dequantize, perform floating-point
operation, and finally quantize" strategy. Others may perform the entire
computation with integer arithmetic. Consequently, the interpretation of
quantized StableHLO operations is exclusively determined by the specific
implementation. The interpretation of hybrid quantization
([#1575](https://github.com/openxla/stablehlo/issues/1575)) should be based on
the it's semantics as prescribed in the specification (via
[1792](https://github.com/openxla/stablehlo/pull/1792)).
### Errors
StableHLO programs are validated through an extensive set of constraints for
individual ops, which rules out many classes of errors prior to run time.
However, error conditions are still possible, e.g. through integer overflows,
out-of-bounds accesses, etc. Unless explicitly called out, all these errors
result in implementation-defined behavior, but this may change in the
future ([#1157](https://github.com/openxla/stablehlo/issues/1157)).
#### Floating-point exceptions
As an exception to this rule, floating-point exceptions in StableHLO programs
have well-defined behavior. Operations which result in exceptions defined by the
IEEE-754 standard (invalid operation, division-by-zero, overflow, underflow, or
inexact exceptions) produce default results (as defined in the standard) and
continue execution without raising the corresponding status flag; similar to
`raiseNoFlag` exception handling from the standard. Exceptions for nonstandard
operations (e.g. complex arithmetic and certain transcendental functions) are
implementation-defined.
#### Shape mismatches
StableHLO supports dynamically-shaped tensors. However, shapes have to agree at
runtime, otherwise the behavior is undefined. StableHLO does not explicitly
provide an op that can assert that a tensor has a given shape at runtime.
Generating correct code is the responsibility of the producer.
As a specific example, the below program is valid. However, at runtime, the
exact shapes of `%arg0` and `%arg1` will have to be the same, otherwise the
behavior of the program is undefined:
```mlir
func.func @foo(%arg0: tensor<?xi32>, %arg1: tensor<?xi32>) -> tensor<?xi32> {
%0 = stablehlo.add %arg0, %arg1 : tensor<?xi32>
return %0 : tensor<?xi32>
}
```
## Notation
For describing syntax, this document is using the modified ISO flavor of EBNF
syntax ([ISO/IEC 14977:1996](https://www.iso.org/standard/26153.html),
[Wikipedia](https://en.wikipedia.org/wiki/Extended_Backus%E2%80%93Naur_form)),
with two modifications: 1) rules are defined using `::=` rather than `=`,
2) concatenation is expressed using juxtaposition rather than `,`.
For describing semantics (i.e. within "Types", "Constants" and "Ops" sections),
we are using formulas which are based on Python syntax extended with support
for concisely expressing array operations as described below. This works well
for small snippets of code, but in rare cases when larger snippets of code are
needed, we use vanilla Python syntax which is always introduced explicitly.
### Formulas
Let's explore how formulas work based on an example from the `dot_general`
specification. One of the constraints for this operation looks as follows:
`dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...)`.
The names used in this formula come from two sources: 1) global functions,
i.e. `dim`, 2) member definitions of the corresponding program element, i.e.
`lhs`, `lhs_batching_dimensions`, `rhs` and `rhs_batching_dimensions` inputs
defined in the "Inputs" section of `dot_general`.
As mentioned above, the syntax of this formula is Python-based with some
conciseness-oriented extensions. To make sense of the formula, let's transform
it into vanilla Python syntax.
A) In these formulas, we are using `=` to represent equality, so the first step
towards obtaining Python syntax is replacing `=` with `==`, as follows:
`dim(lhs, lhs_batching_dimensions...) == dim(rhs, rhs_batching_dimensions...)`.
B) Also, these formulas support ellipses (`...`) which turn scalar expressions
into tensor expressions. In a nutshell, `f(xs...)` roughly means "for each
scalar `x` in the tensor `xs`, compute a scalar `f(x)` and then return all
these scalar results together as a tensor result". In vanilla Python syntax,
our example formula turns into:
`[dim(lhs, dim1) for dim1 in lhs_batching_dimensions] ==
[dim(rhs, dim2) for dim2 in rhs_batching_dimensions]`.
Thanks to ellipses, it is often possible to avoid working at the level of
individual scalars. However, in some tricky cases, lower-level semi-informal
syntax may be used like in the `start_indices[bi0, ..., :, ..., biN]` formula
from the `gather` specification. In the service of conciseness, we don't
provide an exact formalism for translating such syntax to vanilla Python, in
hopes that it is still intuitively understandable on case-by-case basis.
Please let us know if some specific formulas look opaque, and we'll try to
improve them.
Also, you will notice that formulas use ellipses to expand all sorts of lists,
including tensors, lists of tensors (which e.g. can arise from a variadic
number of tensors), etc. This is another area where we don't provide an exact
formalism (e.g. lists are not even part of the StableHLO type system) and
instead rely on intuitive understandability.
C) The final noteworthy notational vehicle that we employ is implicit
broadcasting. While the StableHLO opset doesn't support implicit broadcasting,
the formulas do, also in the service of conciseness. In a nutshell, if a scalar
is used in a context where a tensor is expected, the scalar is broadcasted to
the expected shape.
To continue the `dot_general` example, here's another constraint:
`0 <= lhs_batching_dimensions < rank(lhs)`. As defined in the `dot_general`
specification, `lhs_batching_dimensions` is a tensor, however both `0` and
`rank(lhs)` are scalars. After we apply implicit broadcasting, the formula will
become `[0, ..., 0] <= lhs_batching_dimensions < [rank(lhs), ..., rank(lhs)]`.
When applied to a particular `dot_general` operation, this formula will
evaluate to a tensor of booleans. When formulas are used as constraints, the
constraint holds if the formula evaluates to either `true` or to a tensor which
only has `true` elements.
### Names
In formulas, lexical scope includes: 1) global functions, 2) member definitions,
3) local definitions. The list of global functions is provided below. The list
of element definitions depends on the program element that the notation is
applied to:
* For operations, member definitions include names introduced in "Inputs" and
"Outputs" sections.
* For everything else, member definitions include structural parts of the
program element, named after the corresponding EBNF non-terminals. Most of
the time, the names of these structural parts are obtained by converting the
names of the non-terminals to snake case (e.g. `IntegerLiteral` =>
`integer_literal`), but sometimes names get abbreviated in the process (e.g.
`QuantizationStorageType` => `storage_type`) in which case the names are
introduced explicitly similarly to "Inputs" / "Outputs" sections in operation
specifications.
* Additionally, member definitions always include `self` to refer to the
corresponding program element.
### Values
When formulas are evaluated, they work with the following types of values:
1\) `Value` (actual values, e.g. `dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>`;
they always know their types),
2\) `Placeholder` (future values, e.g. `lhs`, `rhs` or `result`; their actual
values are not known yet, only their types are known),
3\) `Type` (types as defined in the "Types" section),
4\) `Function` (global functions as defined in the "Functions" section).
Depending on the context, names may be referring to different values. More
specifically, the "Semantics" section for ops (and equivalents for other program
elements) defines runtime logic, so all inputs are available as `Value`.
In contrast, the "Constraints" section for ops (and equivalents) defines
"compile-time" logic, i.e. something that is typically executed before runtime,
so only constant inputs are available as `Value` and other inputs are
available only as `Placeholder`.
| Names | In "Semantics" | In "Constraints" |
|---------------------|---------------------------|---------------------------|
| Global functions | `Function` | `Function` |
| Constant inputs | `Value` | `Value` |
| Non-constant inputs | `Value` | `Placeholder` |
| Outputs | `Value` | `Placeholder` |
| Local definitions | Depends on the definition | Depends on the definition |
Let's consider an example `transpose` operation:
```mlir
%result = "stablehlo.transpose"(%operand) {
permutation = dense<[2, 1, 0]> : tensor<3xi64>
} : (tensor<2x3x2xi32>) -> tensor<2x3x2xi32>
```
For this operation, `permutation` is a constant, so it's available as a `Value`
in both semantics and constraints. In contrast, `operand` and `result` are
available as a `Value` in semantics but only as a `Placeholder` in constraints.
### Functions
#### Construction of types
There are no functions that can be used to construct types. Instead, we directly
use type syntax because it's typically more concise. E.g.
`(tensor<E>, tensor<E>) -> (tensor<E>)` rather than `function_type(
[tensor_type([], E), tensor_type([], E)], [tensor_type([], E)])`.
#### Functions on types
* `element_type` is defined on tensor types and quantized tensor types and
returns, respectively, the `TensorElementType` or `QuantizedTensorElementType`
part of the corresponding `TensorType` or `QuantizedTensorType`.
```python
def element_type(x: Value | Placeholder | Type):
if type(x) == TensorType:
return tensor_element_type(x)
if type(x) == QuantizedTensorType:
return quantized_tensor_element_type(x)
if type(x) is not Type:
return element_type(type(x))
```
* `is_per_axis_quantized(x: Value | Placeholder | Type) -> Value` is a shortcut
for `is_quantized(x) and quantization_dimension(x) is not None`.
* `is_per_tensor_quantized(x: Value | Placeholder | Type) -> Value` is a
shortcut for `is_quantized(x) and quantization_dimension(x) is None`.
* `is_promotable(x: Type, y: Type) -> bool` checks if type `x` can be promoted
to type `y`. When `x` and `y` are `QuantizedTensorElementType`s, the promotion
is applied only to the `storage_type`. This specific version of promotion is
currently used in context of reduction computation (refer to
[RFC](https://github.com/openxla/stablehlo/pull/1664) for more details).
```python
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_fp8(x: Value | Placeholder | Type) -> Value` returns `true` if `x` is one
of `Float8E4M3Type`, `Float8E4M3FNType`, `Float8E4M3B11FNUZType`,
`Float8E4M3FNUZType`, `Float8E5M2Type`, `Float8E5M2FNUZType`, `Float8E3M4Type`,
or `Float8E8M0FNUType`. If `x` is a value or placeholder, this function is a
shortcut for `is_fp8(type(x))`.
* `is_quantized(x: Value | Placeholder | Type) -> Value` is a shortcut for
`is_quantized_tensor_element_type(x)`.
* `is_type_name(x: Value | Placeholder | Type) -> Value`. Available for all
types. For example, `is_float(x)` returns `true` if `x` is a `FloatType`.
If `x` is a value or placeholder, this function is a shortcut for
`is_type_name(type(x))`.
* `max_value(x: Type) -> Value` returns the maximum value of an
`TensorElementType`. If `x` is not an `TensorElementType`, returns `None`.
* `min_value(x: Type) -> Value` returns the minimum possible value of an
`TensorElementType`. If `x` is not an `TensorElementType`, returns `None`.
* `member_name(x: Value | Placeholder | Type) -> Any`. Available for all member
definitions `member_name` of all types. For example, `tensor_element_type(x)`
returns the `TensorElementType` part of a corresponding `TensorType`.
If `x` is a value or placeholder, this function is a shortcut for
`member_name(type(x))`. If `x` is not a type that has an appropriate member, or
a value or a placeholder of such a type, returns `None`.
* `is_empty_algorithm(*args: Type)` checks if all dot algorithm fields are set
to `None`. This is needed since dot algorithms have implementation defined
default behaviors, so specifying a default value would be incorrect.
#### Construction of values
* `operation_name(*xs: Value | Type) -> Value`. Available for all operations.
For example, `add(lhs, rhs)` takes two tensor values `lhs` and `rhs` and
returns the output of evaluating the `add` operation with these inputs.
For some operations e.g. `broadcast_in_dim`, types of their outputs are
"load-bearing", i.e. needed to evaluate an operation. In this case, the function
takes these types as arguments.
#### Functions on values
* All Python's operators and functions are available. E.g. both
[subscription](https://docs.python.org/3/reference/expressions.html#subscriptions)
and [slicing](https://docs.python.org/3/reference/expressions.html#slicings)
notations from Python are available to index into tensors, quantized tensors
and tuples.
* `to_destination_type(x: Value, destination_type: Type) -> Value` is defined on
tensors and returns the converted value of `x` based on the `type(x)` and
`destination_type` as follows:
```python
def to_destination_type(x: Value, destination_type: Type) -> Value:
if type(x) == destination_type:
return x
if is_quantized(destination_type):
if is_quantized(type(x)):
return quantize(x, destination_type)
assert is_float(type(x))
return quantize(x, destination_type)
if is_quantized(type(x)):
assert destination_type = expressed_type(type(x))
return dequantize(type(x))
return convert(x, destination_type)
```
There is early discussion on merging `convert`, `uniform_quantize` and
`uniform_dequantize` operations ([#1576](https://github.com/openxla/stablehlo/issues/1576)).
After the merge we do not need the above function and can use the operation name
for `convert` instead.
* `is_nan(x: Value) -> Value` is defined on tensors and returns `true` if
all elements of `x` are `NaN` or `false` otherwise. If `x` is not a tensor,
returns `None`.
* `is_sorted(x: Value) -> Value` is defined on tensors and returns `true` if
elements of `x` are sorted in ascending order with respect to the ascending
lexicographical order of their indices or `false` otherwise. If `x` is not a
tensor, returns `None`.
* `is_unique(x: Value) -> Value` is defined on tensors and returns `true` if `x`
doesn't have duplicate elements or `false` otherwise. If `x` is not a tensor,
returns `None`.
* `member_name(x: Value) -> Any` is defined for all member definitions
`member_name` of all values. For example, `real_part(x)` returns the `RealPart`
part of a corresponding `ComplexConstant`. If `x` is not a value that has an
appropriate member, returns `None`.
* `same(x: Value) -> Value` is defined on tensors and returns `true` if
elements of `x` are all equal to each other or `false` otherwise. If the tensor
doesn't have elements, that counts as "all equal to each other", i.e. the
function returns `true`. If `x` is not a tensor, returns `None`.
* `split(x: Value, num_results: Value, axis: Value) -> Value` is defined on
tensors and returns `num_results` slices of `x` along the axis `axis`.
If `x` is not a tensor or `dim(x, axis) % num_results != 0`, returns `None`.
* `is_defined_in_parent_scope(x: Value) -> Value` is defined on strings
and returns `true` if `x` is the name of a function defined in the same scope
as the parent function of the relevant op.
* `is_namespaced_op_name(x: Value) -> Value` is defined on strings and returns
`true` if `x` is a valid op name, that is it respects the following regular
expression: `[a-zA-Z][a-zA-Z0-9_]*([.][a-zA-Z0-9_$]+)+`
#### Shape computations
* `axes(x: Value | Placeholder | Type) -> Value` is a shortcut for
`range(rank(x))`.
* `dim(x: Value | Placeholder | Type, axis: Value) -> Value` is a shortcut for
`shape(x)[axis]`.
* `dims(x: Value | Placeholder | Type, axes: List) -> List` is a shortcut for
`list(map(lambda axis: dim(x, axis), axes))`.
* `index_space(x: Value | Placeholder | Type) -> Value` is defined on tensors
and returns `size(x)` indices for the corresponding `TensorType` sorted in
ascending lexicographical order, i.e. `[0, ..., 0]`, `[0, ..., 1]`, ...,
`shape(x) - 1`. If `x` is not a tensor type, a quantized tensor type, or a value
or a placeholder of one of these types, returns `None`.
* `rank(x: Value | Placeholder | Type) -> Value` is a shortcut for
`size(shape(x))`.
* `shape(x: Value | Placeholder | Type) -> Value` is defined in the "Functions
on types" section via `member_name`.
* `size(x: Value | Placeholder | Type) -> Value` is a shortcut for
`reduce(lambda x, y: x * y, shape(x))`.
#### Quantization computations
* `def baseline_element_type(x: Value | Placeholder | Type) -> Type` is a
shortcut for `element_type(baseline_type(x))`.
* `baseline_type` is defined on tensor types and quantized tensor types and
transforms them to a "baseline", i.e. a type with the same shape but with the
quantization parameters of the element type reset to default values. This is
used as a handy trick to compare both tensor and quantized tensor types
uniformly, which is needed quite often. For quantized types, this enables
comparing types ignoring the quantization parameters, that is, `shape`,
`storage_type`, `expressed_type`, `storage_min`, `storage_max`, and
`quantization_dimension` (for per-axis quantized type) must all match, but
`scales` and `zero points` may differ.
```python
def baseline_type(x: Value | Placeholder | Type) -> Type:
if type(x) == TensorType:
return x
if type(x) == QuantizedTensorType:
element_type = quantized_tensor_element_type(x)
baseline_element_type = QuantizedTensorElementType(
storage_type = storage_type(element_type),
storage_min = storage_min(element_type),
storage_max = storage_max(element_type),
expressed_type = expressed_type(element_type),
quantization_dimension = quantization_dimension(element_type),
scales = [constant(1.0, expressed_type(element_type))] * dim(x, quantization_dimension(element_type)),
zero_points = [constant(0, storage_type(element_type))] * dim(x, quantization_dimension(element_type)))
return QuantizedTensorType(shape(x), baseline_element_type)
if type(x) is not Type:
return baseline_element_type(type(x))
```
* `dequantize` is defined on quantized tensor types and turns them into
floating-point tensor types. This happens via converting quantized elements
which represent integer values of the storage type into corresponding
floating-point values of the expressed type using the zero point and scale
associated with the quantized element type.
```python
def compute_zero_points(quantized_type, result_type):
if is_per_tensor_quantized(quantized_type):
return broadcast_in_dim(constant(zero_point(quantized_type), storage_type(quantized_type)), [], result_type)
if is_per_axis_quantized(quantized_type):
for i in index_space(result_type):
d = quantization_dimension(quantized_type)
zero_points[i] = zero_points(quantized_type)[i[d]]
return zero_points
def compute_scales(quantized_type, result_type):
if is_per_tensor_quantized(quantized_type):
return broadcast_in_dim(constant(scale(quantized_type), expressed_type(quantized_type)), [],
type(result_type))
if is_per_axis_quantized(quantized_type):
for i in index_space(result_type):
d = quantization_dimension(quantized_type)
scales[i] = scales(quantized_type)[i[d]]
return scales
def dequantize(x: Value) -> Value:
assert is_quantized(x)
x_storage = bitcast_convert(x, storage_type(x))
x_storage_sub = x_storage - compute_zero_points(type(x), type(x_storage))
x_expressed_sub = convert(x_storage_sub, expressed_type(x))
return x_expressed_sub * compute_scales(type(x), type(x_expressed_sub))
```
* `quantize` is defined on floating-point tensor types and turns them into
quantized tensor types. This happens via converting floating-point values
of the expressed type into corresponding integer values of the storage type
using the zero point and scale associated with the quantized element type.
```python
def quantize(x: Value, result_type: Type) -> Value:
assert is_float(x) and is_quantized(result_type)
zero_points = compute_zero_points(result_type, TensorType(shape(x), storage_type(result_type)))
converted_zero_points = convert(zero_points, expressed_type(result_type))
converted_min = convert(storage_min(result_type), expressed_type(result_type))
converted_max = convert(storage_max(result_type), expressed_type(result_type))
x_scaled = x / compute_scales(result_type, type(x))
x_scaled_add_zp = x_scaled + converted_zero_points
x_clamped = clamp(converted_min, x_scaled_add_zp, converted_max)
x_rounded = round_nearest_even(x_clamped)
return convert(x_rounded, result_type)
```
* `dequantize_op_quantize` is used to specify element-wise computations on
quantized tensors. It dequantizes, i.e. turns quantized elements into their
expressed types, then performs an operation, and then quantizes, i.e. turns
the results back into their storage types. At the moment, this function only
works for per-tensor quantization. Per-axis quantization is work in progress
([#1574](https://github.com/openxla/stablehlo/issues/1574)).
```python
def dequantize_op_quantize(op, *inputs_and_output_type):
inputs = inputs_and_output_type[:-1]
output_type = inputs_and_output_type[-1]
float_inputs = map(dequantize, inputs)
float_result = op(*float_inputs)
return quantize(float_result, output_type)
def dequantize_batch_norm_grad_or_training_quantize(op, *inputs_and_output_types):
inputs = inputs_and_output_type[:-3]
float_inputs = map(dequantize, inputs)
float_results = op(*float_inputs)
return map(quantize, float_results, inputs_and_output_type[-3:])
def dequantize_compare(lhs, rhs, comparison_direction):
float_lhs = dequantize(lhs)
float_rhs = dequantize(rhs)
return compare(float_lhs, float_rhs, comparison_direction, FLOAT)
def dequantize_select_quantize(pred, on_true, on_false, output_type):
float_on_true = dequantize(on_true)
float_on_false = dequantize(on_false)
float_result = select(pred, float_on_true, float_on_false)
return quantize(float_result, output_type)
```
* `hybrid_dequantize_then_op` is used to specify weight-only quantization for
hybrid op which accepts lhs in floating-point and rhs in quantized types. It
dequantizes quantized inputs into their expressed types and performs computation
in float. Element type of float lhs tensor and expressed type of quantized rhs
tensor should be identical.
```python
def hybrid_dequantize_then_op(op, lhs, rhs):
assert(is_float(lhs) and is_quantized(rhs) and element_type(lhs) == expressed_type(rhs))
return op(lhs, dequantize(rhs))
```
#### Grid computations
* `cross_partition(replica_groups: Value) -> Value`. See the "cross_replica"
section above.
* `cross_replica(replica_groups: Value) -> Value`. See the "cross_replica"
section above.
* `cross_replica_and_partition(replica_groups: Value) -> Value`. See the
"cross_replica_and_partition" section above.
* `flattened_ids(replica_groups: Value) -> Value`. See the "flattened_ids"
section above.
## Dynamism
StableHLO values can have dynamic dimension sizes, e.g. `tensor<?xi64>`.
However, StableHLO values cannot have a dynamic number of dimensions (unranked
dynamism, e.g. `tensor<*xi64>`). Operands and results are allowed to use dynamic
dimension sizes, even if there are constraints on the sizes. Constraints will be
verified statically if possible, otherwise they are deferred to runtime and
mismatches will result in undefined behavior. See below for examples.
### Shape mismatches for unary elementwise operations
Consider the following toy program:
```mlir
func.func @foo(%arg0: tensor<?xf64>) {
%0 = stablehlo.abs %arg0 : (tensor<?xf64>) -> tensor<2xf64>
return
}
```
Such a program is unusual, because it is not common to know the shape of the
result but not the shape of the input. Nonetheless, this is a valid StableHLO
program. It is not possible to statically validate the `abs` operation in this
program, because the exact shape of the operand is unknown. However, the shapes
are certainly compatible, and this can be checked statically: `?` could turn out
to be `2` at runtime, and there would be no issue. However, `?` could
also turn out to be some other integer, in which case the behavior is undefined.
Note that if a dimension size is dynamic in the result, there cannot be
undefined behavior. Indeed, there is no "expected" size, so there cannot be a
mismatch.
### Shape mismatches for binary elementwise operations
Consider the following toy program:
```mlir
func.func @foo(%arg0: tensor<?xf64>, %arg1: tensor<?xf64>) {
%0 = stablehlo.add %arg0, %arg0 : (tensor<?xf64>, tensor<?xf64>) -> tensor<?xf64>
return
}
```
When it comes to binary elementwise operations, the shapes of the inputs and the
result must agree at runtime. At compile time, static dimensions must be equal,
otherwise they merely need to be compatible.
If *any* dimension is dynamic in the inputs, then there could be undefined
behavior at runtime, because the dynamic size may not match the corresponding
size in the other operand (be it static or dynamic). If all the inputs are
static, then whether the result is dynamic or not does not matter: statically
known dimensions will be checked statically, and dynamic dimensions do not
impose any constraints.
### Shape mismatches for ops that take their output shape as an operand
Consider the following toy program:
```mlir
func.func @foo(%arg0: tensor<2xi32>) {
%0 = stablehlo.dynamic_iota %arg0, dim = 0 : (tensor<2xi32>) -> tensor<3x4xi64>
return
}
```
The values in the shape operand at runtime must match the shape of the result,
otherwise the behavior is undefined. That is, at runtime `%arg0` must have a
value of `dense<[3, 4]> : tensor<2xi32>`. If the shape operand is constant, this
can be verified statically. If the result shape is fully dynamic, then there
cannot be a mismatch.