'mhlo' Dialect

Operations

mhlo.abs (mhlo::AbsOp)

Abs operation

Syntax:

operation ::= `mhlo.abs` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise abs operation on operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#abs

Example:

%result = mhlo.abs %operand : tensor<3xi32>

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of 2/4/8/16/32/64-bit signless integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
result ranked tensor of 2/4/8/16/32/64-bit signless integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.add (mhlo::AddOp)

Add operation

Syntax:

operation ::= `mhlo.add` $lhs `,` $rhs attr-dict
              `:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))

Performs element-wise addition of two tensors lhs and rhs and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#add

Example:

%result = mhlo.add %lhs, %rhs : tensor<2x2xi32>

Traits: AlwaysSpeculatableImplTrait, Commutative, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
lhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
rhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.add_dependency (mhlo::AddDependencyOp)

AddDependency operation

Syntax:

operation ::= `mhlo.add_dependency` operands attr-dict `:` functional-type(operands, results)

This operation is private to the XLA compiler, so it is does not yet have a specification.

Informally, this operation two operands: a data operand and a token. The output of the operation is the data operand. When used with AfterAll this operation enables ordering non-side-effecting operations (those that do not produce token values).

Example:

%1 = mhlo.add_dependency %arg0, %0 : (tensor<3x4xf32>, !mhlo.token) -> tensor<3x4xf32>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token
token token

Results:

Result Description
output ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token

mhlo.after_all (mhlo::AfterAllOp)

AfterAll operation

Syntax:

operation ::= `mhlo.after_all` $inputs attr-dict
              `:` custom<VariadicSameOperandsAndResultType>(ref($inputs), type($inputs), type($result))

Ensures that the operations producing the inputs are executed before any operations that depend on result.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#after_all

Example:

%result = mhlo.after_all %input0, %input1 : !mhlo.token

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
inputs variadic of token

Results:

Result Description
result token

mhlo.all_gather (mhlo::AllGatherOp)

AllGather operation

Within each process group in the process grid, concatenates the values of the operand tensor from each process along all_gather_dim and produces a result tensor. The computation is applied separately for each operand in operands, producing one result per operand.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#all_gather

Example:

%result = "mhlo.all_gather"(%operand) {
  all_gather_dim = 1 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
  // channel_id = 0
  channel_handle = #mhlo.channel_handle<handle = 0, type = 0>,
  // use_global_device_ids = false
} : (tensor<2x2xf32>) -> tensor<2x4xf32>

Traits: SameOperandsAndResultElementType

Attributes:

AttributeMLIR TypeDescription
all_gather_dim::mlir::IntegerAttr64-bit signless integer attribute whose value is non-negative
replica_groups::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
channel_handle::mlir::mhlo::ChannelHandleAttrtwo 64-bit integers 'handle' and 'type'
use_global_device_ids::mlir::UnitAttrunit attribute

Operands:

Operand Description
operands variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.all_reduce (mhlo::AllReduceOp)

AllReduce operation

Within each process group in the process grid, applies a reduction function computation to the values of an operand tensor from each process and produces a result tensor. The computation is applied separately for each operand in operands, producing one result per operand.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#all_reduce

Example:

%result = "mhlo.all_reduce"(%operand) ({
  ^bb0(%arg0: tensor<f32>, %arg1: tensor<f32>):
    %0 = mhlo.add %arg1, %arg2 : tensor<f32>
    mhlo.return %0 : tensor<f32>
}) {
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
  // channel_id = 0
  channel_handle = #mhlo.channel_handle<handle = 0, type = 0>
  // use_global_device_ids = false
} : (tensor<4xf32>) -> tensor<4xf32>

Traits: InferTensorType, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface

Attributes:

AttributeMLIR TypeDescription
replica_groups::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
channel_handle::mlir::mhlo::ChannelHandleAttrtwo 64-bit integers 'handle' and 'type'
use_global_device_ids::mlir::UnitAttrunit attribute

Operands:

Operand Description
operands variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.all_to_all (mhlo::AllToAllOp)

AllToAll operation

Within each process group in the process grid, splits the values of the operand tensor along split_dimension into parts, scatters the split parts between the processes, concatenates the scattered parts along concat_dimension and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#all_to_all

Example:

%result = "mhlo.all_to_all"(%operand) {
  split_dimension = 1 : i64,
  concat_dimension = 0 : i64,
  split_count = 2 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
} : (tensor<2x4xf32>) -> tensor<4x2xf32>

Traits: AlwaysSpeculatableImplTrait, InferTensorType, SameOperandsElementType, SameOperandsShape, SameVariadicOperandSize

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
split_dimension::mlir::IntegerAttr64-bit signless integer attribute whose value is non-negative
concat_dimension::mlir::IntegerAttr64-bit signless integer attribute whose value is non-negative
split_count::mlir::IntegerAttr64-bit signless integer attribute whose value is positive
replica_groups::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
channel_handle::mlir::mhlo::ChannelHandleAttrtwo 64-bit integers 'handle' and 'type'

Operands:

Operand Description
operand variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.and (mhlo::AndOp)

And operation

Syntax:

operation ::= `mhlo.and` $lhs `,` $rhs attr-dict
              `:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))

Performs element-wise AND of two tensors lhs and rhs and produces a result tensor

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#and

Example:

%result = mhlo.and %lhs, %rhs : tensor<2x2xi32>

Traits: AlwaysSpeculatableImplTrait, Commutative, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
lhs ranked tensor of pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values
rhs ranked tensor of pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.async_done (mhlo::AsyncDoneOp)

AsyncDone operation

This operation is private to the XLA compiler, so it is does not yet have a specification.

Informally, this operation blocks until the end of an asynchronous computation. It returns the final result of the asynchronous computation.

See the documentation for AsyncStart for more information.

Interfaces: InferTypeOpInterface

Operands:

Operand Description
bundle async_bundle with any combination of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values

Results:

Result Description
«unnamed» variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token or nested tuple with any combination of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values

mhlo.async_start (mhlo::AsyncStartOp)

AsyncStart operation

This operation is private to the XLA compiler, so it is does not yet have a specification.

Informally, this operation kicks off an asynchronous computation.

This is used when there are functions that contain both asynchronous waits (such as DMAs) and on-thread computation. For example, a function might consist of a computation, a DMA, another computation, a second DMA, and a final computation. This would be represented as an async_start followed by and async_update and an async_done. The async_start would do the first computation on-thread and then start the DMA. The async_update would wait for the DMA to complete if it wasn't yet done, then execute the second computation in the function, and start the second DMA. Finally, the async_done would wait on this last DMA, and then run the last computation that needs to be run on-thread and return the result of that final computation.

operands are passed to the computation directly called_computation is the function that will be run asynchronously execution_thread is the name of the thread in which it will be run. The main thread is called "main". All threads have names.

This returns all the state needed between async ops. After buffer assignment, the return values represents the space needed to hold the input, results, and any scratchpads needed or edited by the async op.

Attributes:

AttributeMLIR TypeDescription
called_computation::mlir::FlatSymbolRefAttrflat symbol reference attribute
execution_thread::mlir::StringAttrstring attribute

Operands:

Operand Description
inputs variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token or nested tuple with any combination of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values

Results:

Result Description
«unnamed» async_bundle with any combination of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values

mhlo.async_update (mhlo::AsyncUpdateOp)

AsyncUpdate operation

This operation is private to the XLA compiler, so it is does not yet have a specification.

Informally, this operation blocks on an asynchronous computation until a sync barrier. This returns bundle after operating on it.

See the documentation for AsyncStart for more information.

Interfaces: InferTypeOpInterface

Operands:

Operand Description
bundle async_bundle with any combination of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values

Results:

Result Description
«unnamed» async_bundle with any combination of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values

mhlo.atan2 (mhlo::Atan2Op)

Atan2 operation

Syntax:

operation ::= `mhlo.atan2` $lhs `,` $rhs attr-dict
              `:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))

Performs element-wise atan2 operation on lhs and rhs tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#atan2

Example:

%result = mhlo.atan2 %lhs, %rhs : tensor<3xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
lhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values
rhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

mhlo.batch_norm_grad (mhlo::BatchNormGradOp)

BatchNormGrad operation

Computes gradients of several inputs of BatchNormTrainingOp backpropagating from grad_output, and produces grad_operand, grad_scale and grad_offset tensors.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#batch_norm_grad

Example:

%grad_operand, %grad_scale, %grad_offset =
"mhlo.batch_norm_grad"(%operand, %scale, %mean, %variance, %grad_output) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
} : (tensor<2x2x2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>,
    tensor<2x2x2xf32>) -> (tensor<2x2x2xf32>, tensor<2xf32>, tensor<2xf32>)

Traits: AlwaysSpeculatableImplTrait, InferTensorType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
epsilon::mlir::FloatAttr32-bit float attribute
feature_index::mlir::IntegerAttr64-bit signless integer attribute whose value is non-negative

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
scale 1D tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
mean 1D tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
variance 1D tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
grad_output ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

Results:

Result Description
grad_operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
grad_scale 1D tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
grad_offset 1D tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

mhlo.batch_norm_inference (mhlo::BatchNormInferenceOp)

BatchNormInference operation

Normalizes the operand tensor across all dimensions except for the feature_index dimension and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#batch_norm_inference

Example:

%result = "mhlo.batch_norm_inference"(%operand, %scale, %offset, %mean, %variance) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
} : (tensor<2x2x2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>) -> tensor<2x2x2xf32>

Traits: AlwaysSpeculatableImplTrait, InferTensorType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
epsilon::mlir::FloatAttr32-bit float attribute
feature_index::mlir::IntegerAttr64-bit signless integer attribute whose value is non-negative

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
scale 1D tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
offset 1D tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
mean 1D tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
variance 1D tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

mhlo.batch_norm_training (mhlo::BatchNormTrainingOp)

BatchNormTraining operation

Computes mean and variance across batch and spatial dimensions and normalizes the operand tensor, for each feature in the feature_index dimension and produces output, batch_mean and batch_var tensors.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#batch_norm_training

Example:

%output, %batch_mean, %batch_var = "mhlo.batch_norm_training"(%operand, %scale, %offset) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
} : (tensor<2x2x2xf32>, tensor<2xf32>, tensor<2xf32>) -> (tensor<2x2x2xf32>, tensor<2xf32>, tensor<2xf32>)

Traits: AlwaysSpeculatableImplTrait, InferTensorType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
epsilon::mlir::FloatAttr32-bit float attribute
feature_index::mlir::IntegerAttr64-bit signless integer attribute whose value is non-negative

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
scale 1D tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
offset 1D tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

Results:

Result Description
output ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
batch_mean 1D tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
batch_var 1D tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

mhlo.bitcast (mhlo::BitcastOp)

Bitcast operation

Syntax:

operation ::= `mhlo.bitcast` operands attr-dict `:` functional-type(operands, results)

This operation is private to the XLA compiler, so it is does not yet have a specification.

Informally, this operation changes the shape of the input in the way that the physical arrangement of elements are unchanged.

This operation needs layout information to make sense of "physical arrangement of elements", and layout support in MHLO is currently a work in progress.

Example:

%0 = mhlo.bitcast %arg0 : (tensor<3x4xf32>) -> tensor<3x4x1xf32>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.bitcast_convert (mhlo::BitcastConvertOp)

BitcastConvert operation

Syntax:

operation ::= `mhlo.bitcast_convert` operands attr-dict `:` functional-type(operands, results)

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.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#bitcast_convert

Example:

%result = mhlo.bitcast_convert %operand : (tensor<2xf32>) -> tensor<2x4xi8>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.broadcast (mhlo::BroadcastOp)

Broadcast operation

This operation is on its way out of StableHLO, so it is not included in the specification: https://github.com/openxla/stablehlo/issues/3

Informally, this operation does the same thing as XLA's Broadcast: https://www.tensorflow.org/xla/operation_semantics#broadcast

Example:

%result = mhlo.broadcast %operand, sizes = [1, 2] : (tensor<3xi32>) -> tensor<1x2x3xi32>

Traits: AlwaysSpeculatableImplTrait, InferTensorType, SameOperandsAndResultElementType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
broadcast_sizes::mlir::DenseIntElementsAttr64-bit signless integer elements attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.broadcast_in_dim (mhlo::BroadcastInDimOp)

BroadcastInDim operation

Expands the dimensions and/or rank of an input tensor by duplicating the data in the operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#broadcast_in_dim

Example:

%result = mhlo.broadcast_in_dim %operand, dims = [2, 1] : (tensor<1x3xi32>) -> tensor<2x3x2xi32>

Traits: AlwaysSpeculatableImplTrait, HLO_CompatibleOperandsAndResultElementType

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
broadcast_dimensions::mlir::DenseIntElementsAttr64-bit signless integer elements attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» statically shaped tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.case (mhlo::CaseOp)

Case operation

Produces the output from executing exactly one function from branches depending on the value of index.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#case

Example:

%result0, %result1 = "mhlo.case"(%index) ({
  mhlo.return %result_branch0, %result_branch0 : tensor<2xi64>, tensor<2xi64>
}, {
  mhlo.return %result_branch1, %result_branch1 : tensor<2xi64>, tensor<2xi64>
}) : (tensor<i32>) -> (tensor<2xi64>, tensor<2xi64>)

Traits: RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfaces: InferTypeOpInterface

Operands:

Operand Description
index tensor of 32-bit signless integer values

Results:

Result Description
«unnamed» variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token

mhlo.cbrt (mhlo::CbrtOp)

Cbrt operation

Syntax:

operation ::= `mhlo.cbrt` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise cubic root operation on operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#cbrt

Example:

%result = mhlo.cbrt %operand : tensor<4xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

mhlo.ceil (mhlo::CeilOp)

Ceil operation

Syntax:

operation ::= `mhlo.ceil` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise ceil of operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#ceil

Example:

%result = mhlo.ceil %operand : tensor<5xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

mhlo.cholesky (mhlo::CholeskyOp)

Cholesky operation

Computes the Cholesky decomposition of a batch of matrices.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#cholesky

Example:

%result = mhlo.cholesky %a, lower = true : tensor<3x3xf32>

Traits: AlwaysSpeculatableImplTrait, InferTensorType, SameOperandsAndResultElementType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
lower::mlir::BoolAttrbool attribute

Operands:

Operand Description
a ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements values

mhlo.clamp (mhlo::ClampOp)

Clamp operation

Syntax:

operation ::= `mhlo.clamp` $min `,` $operand `,` $max attr-dict
              `:` custom<SameOperandsAndResultType>(type($min), type($operand), type($max), type($result))

Clamps every element of the operand tensor between a minimum and maximum value and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#clamp

Example:

%result = mhlo.clamp %min, %operand, %max : tensor<3xi32>

Traits: AlwaysSpeculatableImplTrait, HLO_BroadcastingElementwise, InferTensorType, SameOperandsAndResultElementType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
min ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
max ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.collective_broadcast (mhlo::CollectiveBroadcastOp)

CollectiveBroadcast operation

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#collective_broadcast

Example:

%result = "mhlo.collective_broadcast"(%operand) {
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  channel_handle = #mhlo.channel_handle<handle = 0, type = 0>
} : (tensor<1x2xi64>) -> tensor<1x2xi64>

Traits: CompatibleOperandsAndResultType

Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface

Attributes:

AttributeMLIR TypeDescription
replica_groups::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
channel_handle::mlir::mhlo::ChannelHandleAttrtwo 64-bit integers 'handle' and 'type'

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.collective_permute (mhlo::CollectivePermuteOp)

CollectivePermute operation

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#collective_permute

Example:

%result = "mhlo.collective_permute"(%operand) {
  source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>,
  // channel_id = 0
  channel_handle = #mhlo.channel_handle<handle = 0, type = 0>
} : (tensor<4x2xf32>) -> tensor<4x2xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
source_target_pairs::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
channel_handle::mlir::mhlo::ChannelHandleAttrtwo 64-bit integers 'handle' and 'type'

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.compare (mhlo::CompareOp)

Compare operation

Syntax:

operation ::= `mhlo.compare` $comparison_direction `,` $lhs `,` $rhs (`,` $compare_type^)?
              attr-dict `:` functional-type(operands, results)

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#compare

Example:

%result = mhlo.compare LT, %lhs, %rhs, FLOAT : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xi1>

Traits: AlwaysSpeculatableImplTrait, Elementwise, InferTensorType, SameOperandsAndResultShape, SameOperandsElementType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
comparison_direction::mlir::mhlo::ComparisonDirectionAttrWhich comparison operation to perform.
compare_type::mlir::mhlo::ComparisonTypeAttrWhich comparison type to use.

Operands:

Operand Description
lhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
rhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of pred (AKA boolean or 1-bit integer) values

mhlo.complex (mhlo::ComplexOp)

Complex operation

Syntax:

operation ::= `mhlo.complex` operands attr-dict
              `:` custom<ComplexOpType>(type($lhs), type($rhs), type($result))

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#complex

Example:

%result = mhlo.complex %lhs, %rhs : tensor<2xcomplex<f32>>

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultShape, SameOperandsElementType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
lhs ranked tensor of 32-bit float or 64-bit float values
rhs ranked tensor of 32-bit float or 64-bit float values

Results:

Result Description
result ranked tensor of complex type with 32-bit float or 64-bit float elements values

mhlo.composite (mhlo::CompositeOp)

Composite operation

Syntax:

operation ::= `mhlo.composite` $name $inputs attr-dict `:` functional-type(operands, results)

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

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#composite

Example:

%results = mhlo.composite "my.op" %arg0, %arg1 {
  decomposition = @my_op,
  composite_attributes = { my_attribute = "my_value" },
  version = 1 : i32
} : (tensor<f32>, tensor<f32>) -> tensor<f32>

Interfaces: SymbolUserOpInterface

Attributes:

AttributeMLIR TypeDescription
name::mlir::StringAttrstring attribute
composite_attributes::mlir::DictionaryAttrdictionary of named attribute values
decomposition::mlir::FlatSymbolRefAttrflat symbol reference attribute
version::mlir::IntegerAttr32-bit signless integer attribute

Operands:

Operand Description
inputs variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token or nested tuple with any combination of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values

Results:

Result Description
«unnamed» variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token or nested tuple with any combination of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values

mhlo.concatenate (mhlo::ConcatenateOp)

Concatenate operation

Concatenates a variadic number of tensors in inputs along dimension dimension in the same order as the given arguments and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#concatenate

Example:

%result = mhlo.concatenate %input0, %input1, dim = 0 : (tensor<3x2xi64>, tensor<1x2xi64>) -> tensor<4x2xi64>

Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
dimension::mlir::IntegerAttr64-bit signless integer attribute whose value is non-negative

Operands:

Operand Description
val variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.constant (mhlo::ConstantOp)

Constant operation

Produces an output tensor from a constant value.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#constant

Example:

%output = mhlo.constant dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32>

Traits: AlwaysSpeculatableImplTrait, ConstantLike

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
value::mlir::ElementsAttrconstant vector/tensor attribute

Results:

Result Description
output statically shaped tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.convert (mhlo::ConvertOp)

Convert operation

Syntax:

operation ::= `mhlo.convert` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#convert

Example:

%result = mhlo.convert %operand : (tensor<3xi32>) -> tensor<3xcomplex<f32>>

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.convolution (mhlo::ConvolutionOp)

Convolution operation

Syntax:

operation ::= `mhlo.convolution` `(`operands`)`
              `dim_numbers` `=` custom<ConvolutionDimensions>($dimension_numbers) `,`
              `window` `=` `{` custom<WindowAttributes>($window_strides, $padding,
              $lhs_dilation, $rhs_dilation,
              $window_reversal) `}`
              attr-dict `:` functional-type(operands, results)

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#convolution

Example:

%result = "mhlo.convolution"(%lhs, %rhs) {
  window_strides = dense<4> : tensor<2xi64>,
  padding = dense<0> : tensor<2x2xi64>,
  lhs_dilation = dense<2> : tensor<2xi64>,
  rhs_dilation = dense<1> : tensor<2xi64>,
  window_reversal = dense<false> : tensor<2xi1>,
  dimension_numbers = #mhlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
  feature_group_count = 1 : i64,
  batch_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi32>, tensor<3x3x1x1xi32>) -> tensor<1x2x2x1xi32>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
window_strides::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
padding::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
lhs_dilation::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
rhs_dilation::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
window_reversal::mlir::DenseElementsAttrconstant boolean vector/tensor attribute
dimension_numbers::mlir::mhlo::ConvDimensionNumbersAttrStructure of dimension information for conv op
feature_group_count::mlir::IntegerAttr64-bit signless integer attribute whose value is positive
batch_group_count::mlir::IntegerAttr64-bit signless integer attribute whose value is positive
precision_config::mlir::ArrayAttrPrecision Config attribute

Operands:

Operand Description
lhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
rhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.copy (mhlo::CopyOp)

Copy operation

Syntax:

operation ::= `mhlo.copy` operands attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

This operation is private to the XLA compiler, so it is does not yet have a specification.

Informally, this operation a copy of operand. Depending on the metadata attached to the operation, it can behave quite differently from a no-op.

Example:

%0 = mhlo.copy %arg0 : tensor<f32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
cross_program_prefetch_index::mlir::IntegerAttr32-bit signless integer attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token or nested tuple with any combination of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token or nested tuple with any combination of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values

mhlo.cosine (mhlo::CosineOp)

Cosine operation

Syntax:

operation ::= `mhlo.cosine` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise cosine operation on operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#cosine

Example:

%result = mhlo.cosine %operand : tensor<2xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

mhlo.count_leading_zeros (mhlo::ClzOp)

Clz operation

Syntax:

operation ::= `mhlo.count_leading_zeros` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#count_leading_zeros

Example:

%result = mhlo.count_leading_zeros %operand : tensor<2x2xi8>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
result ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

mhlo.create_token (mhlo::CreateTokenOp)

CreateToken operation

Syntax:

operation ::= `mhlo.create_token` attr-dict `:` type(results)

This operation is on its way out of StableHLO, so it is not included in the specification: https://github.com/openxla/stablehlo/issues/3

Informally, this operation does the same thing as AfterAllOp with 0 inputs: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#after_all

Example:

%output = mhlo.create_token : !mhlo.token

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Results:

Result Description
output token

mhlo.cross-replica-sum (mhlo::CrossReplicaSumOp)

CrossReplicaSum operation

This operation is on its way out of StableHLO, so it is not included in the specification: https://github.com/openxla/stablehlo/issues/3

Informally, this operation does the same thing as AllReduceOp with channel_id = 0, use_global_device_ids = false and computation implementing addition: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#all_reduce

Example:

%result = "mhlo.cross-replica-sum"(%operand) {
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
} : (tensor<4xf32>) -> tensor<4xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
replica_groups::mlir::DenseIntElementsAttr64-bit signless integer elements attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.custom_call (mhlo::CustomCallOp)

CustomCall operation

Syntax:

operation ::= `mhlo.custom_call` custom<CustomCallTarget>($call_target_name) `(` $inputs `)`
              attr-dict `:` functional-type(operands, results)

Encapsulates an implementation-defined operation call_target_name that takes inputs and called_computations and produces results.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#custom_call

Example:

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

A custom call invokes code external to XLA. The `inputs` are passed to the
external code, and the external code is expected to produce a result of the
given type. The exact mechanism is backend-specific. For example, in the CPU
backend, a call instruction is emitted which targets a symbol with the name
`call_target_name`.

If XLA runtime is enabled for a backend, then custom calls use the runtime
custom call calling convention to call into the external functions. This
calling convention defines an ABI for encoding arguments, attributes and
results.

Depending on the API version there are two ways to pass extra bits of static
information to the external function:

1. For `API_VERSION_TYPED_FFI` custom calls `backend_config` must be a
   dictionary attribute, that will be encoded according to the custom call
   calling convention and passed to the external function as the attributes
   argument. External code is expected to use declarative bindings (see
   `xla/runtime/custom_call.h`) to decode them at run time. These custom
   calls are only supported if XLA uses XLA runtime.

2. For previous API versions it is the user responsibility to encode extra
   bits of static information as a string `backend_config` attribute, and
   decode it at run time.

Interfaces: MemoryEffectOpInterface

Attributes:

AttributeMLIR TypeDescription
call_target_name::mlir::StringAttrstring attribute
has_side_effect::mlir::BoolAttrbool attribute
backend_config::mlir::Attributestring attribute or dictionary of named attribute values
api_version::mlir::mhlo::CustomCallApiVersionAttrCustom call API version
called_computations::mlir::ArrayAttrflat symbol ref array attribute
custom_call_schedule::mlir::mhlo::CustomCallScheduleAttrSpecifies the desired schedule for the custom-call.
operand_layouts::mlir::ArrayAttrArray of layout (1D tensor of index type) attributes
result_layouts::mlir::ArrayAttrArray of layout (1D tensor of index type) attributes
output_operand_aliases::mlir::ArrayAttrAliasing attribute for outputs and operands of CustomCall

Operands:

Operand Description
inputs variadic of tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token or nested tuple with any combination of tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values

Results:

Result Description
«unnamed» variadic of tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token or nested tuple with any combination of tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values

mhlo.divide (mhlo::DivOp)

Div operation

Syntax:

operation ::= `mhlo.divide` $lhs `,` $rhs attr-dict
              `:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))

Performs element-wise division of dividend lhs and divisor rhs tensors and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#divide

Example:

%result = mhlo.divide %lhs, %rhs : tensor<4xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
lhs ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
rhs ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
result ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.domain (mhlo::DomainOp)

Domain operation

This operation is private to the XLA compiler, so it is does not yet have a specification.

Informally, these operations are used to group instructions with the same DomainMetadata property. ShardingMetadata is the main use case today to group instructions on the same device. Domain instructions provide two major benefits:

  • Prevent unintentionally optimizing instructions across domains.
  • Automatically assign the metadata of the instructions created in the domain. Without domain instructions, each HLO optimization pass would have to check and propagate the metadata, which would be easy to miss and also adds complexity to the compiler. Since domain instructions connect two different domains, each domain instruction is associated with two DomainMetadata -- one on the operand side and one on the user side of the domain.

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
kind::mlir::mhlo::DomainKindAttrKind of domain metatdata attached to an HLO domain.
entry_metadata::mlir::StringAttrstring attribute
exit_metadata::mlir::StringAttrstring attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token

mhlo.dot (mhlo::DotOp)

Dot operation

This operation is on its way out of StableHLO, so it is not included in the specification: https://github.com/openxla/stablehlo/issues/3

Informally, this operation does the same thing as XLA's Dot: https://www.tensorflow.org/xla/operation_semantics#dot

Example:

%0 = mhlo.dot %arg0, %arg1 : (tensor<1x2xi32>, tensor<2x1xi32>) -> tensor<1x1xi32>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
precision_config::mlir::ArrayAttrPrecision Config attribute

Operands:

Operand Description
lhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
rhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.dot_general (mhlo::DotGeneralOp)

DotGeneral operation

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dot_general

Example:

%result = "mhlo.dot_general"(%lhs, %rhs) {
  dot_dimension_numbers = #mhlo.dot<
    lhs_batching_dimensions = [0],
    rhs_batching_dimensions = [0],
    lhs_contracting_dimensions = [2],
    rhs_contracting_dimensions = [1]
  >,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<2x2x2xi32>, tensor<2x2x2xi32>) -> tensor<2x2x2xi32>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
dot_dimension_numbers::mlir::mhlo::DotDimensionNumbersAttrAttribute that models the dimension information for dot.
precision_config::mlir::ArrayAttrPrecision Config attribute
algorithm::mlir::mhlo::DotAlgorithmAttrAttribute that models the algorithm constraints to use for computing dot.

Operands:

Operand Description
lhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
rhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.dynamic_broadcast_in_dim (mhlo::DynamicBroadcastInDimOp)

DynamicBroadcastInDim operation

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

It also accepts optional attributes to express static knowledge about the expanding behavior of dimensions. If not specified, all dimensions are assumed to be possibly expanding. The sets of dimensions that are known to be expanding and the set of dimensions that are known to be non-expanding must be disjoint and they must be a subset of the operand's dimensions.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dynamic_broadcast_in_dim

Example:

%operand = mhlo.constant dense<[[1, 2, 3]]> : tensor<1x3xi64>
%output_dimensions = mhlo.constant dense<[2, 3, 2]> : tensor<3xi64>
%result = "mhlo.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>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
broadcast_dimensions::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
known_expanding_dimensions::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
known_nonexpanding_dimensions::mlir::DenseIntElementsAttr64-bit signless integer elements attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
output_dimensions 1D tensor of index or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.dynamic_conv (mhlo::DynamicConvOp)

DynamicConv operation

This operation is a work in progress, so it is not yet included in the specification: https://github.com/openxla/stablehlo/issues/8

Informally, this operation does the same thing as ConvolutionOp except that padding is specified dynamically via d_padding: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#convolution

Example:

%result = "mhlo.dynamic_conv"(%lhs, %rhs, %d_padding) {
  window_strides = dense<4> : tensor<2xi64>,
  lhs_dilation = dense<2> : tensor<2xi64>,
  rhs_dilation = dense<1> : tensor<2xi64>,
  window_reversal = dense<false> : tensor<2xi1>,
  dimension_numbers = #mhlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
  feature_group_count = 1 : i64,
  batch_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi32>, tensor<3x3x1x1xi32>, tensor<2x2xi64>) -> tensor<1x2x2x1xi32>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
window_strides::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
padding::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
lhs_dilation::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
rhs_dilation::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
window_reversal::mlir::DenseElementsAttrconstant boolean vector/tensor attribute
dimension_numbers::mlir::mhlo::ConvDimensionNumbersAttrStructure of dimension information for conv op
feature_group_count::mlir::IntegerAttr64-bit signless integer attribute whose value is positive
batch_group_count::mlir::IntegerAttr64-bit signless integer attribute whose value is positive
precision_config::mlir::ArrayAttrPrecision Config attribute

Operands:

Operand Description
lhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
rhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
d_padding ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.dynamic_gather (mhlo::DynamicGatherOp)

DynamicGather operation

This operation is a work in progress, so it is not yet included in the specification: https://github.com/openxla/stablehlo/issues/8

Informally, this operation does the same thing as GatherOp except that slice_sizes are specified dynamically: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#gather

Example:

%result = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
  dimension_numbers = #mhlo.gather<
    offset_dims = [2, 3],
    collapsed_slice_dims = [0],
    start_index_map = [0, 2],
    index_vector_dim = 2>,
  indices_are_sorted = false
} : (tensor<3x4x2xi32>, tensor<2x3x2xi64>, tensor<3xi64>) -> tensor<2x3x2x2xi32>

Traits: AlwaysSpeculatableImplTrait, InferTensorType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
dimension_numbers::mlir::mhlo::GatherDimensionNumbersAttrAttribute that models the dimension information for gather
indices_are_sorted::mlir::BoolAttrbool attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
start_indices ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values
slice_sizes statically shaped 1-dimensional integer tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.dynamic_iota (mhlo::DynamicIotaOp)

DynamicIota operation

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dynamic_iota

Example:

%0 = mhlo.dynamic_iota %arg0, dim = 0 : (tensor<1xindex>) -> tensor<4xi32>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
iota_dimension::mlir::IntegerAttr64-bit signless integer attribute whose value is non-negative

Operands:

Operand Description
output_shape 1D tensor of index or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.dynamic_pad (mhlo::DynamicPadOp)

DynamicPad operation

Syntax:

operation ::= `mhlo.dynamic_pad` operands attr-dict `:` functional-type(operands, results)

Dynamically Pads the operand, with amount of padding added at low-end/high-end/interior is passed through input tensors.

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
padding_value ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
edge_padding_low 1D tensor of index or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values
edge_padding_high 1D tensor of index or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values
interior_padding 1D tensor of index or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.dynamic_reshape (mhlo::DynamicReshapeOp)

DynamicReshape operation

Syntax:

operation ::= `mhlo.dynamic_reshape` operands attr-dict `:` functional-type(operands, results)

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dynamic_reshape

Example:

%output_shape = mhlo.constant dense<[3, 2]> : tensor<2xi64>
%result = mhlo.dynamic_reshape %operand, %output_shape : (tensor<2x3xi64>, tensor<2xi64>) -> tensor<3x2xi64>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
output_shape 1D tensor of index or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
result tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.dynamic_slice (mhlo::DynamicSliceOp)

DynamicSlice operation

Extracts a slice from the operand using dynamically-computed starting indices and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dynamic_slice

Example:

%result = mhlo.dynamic_slice %operand, %start_indices0, %start_indices1, sizes = [2, 2]
  : (tensor<4x4xi32>, tensor<i64>, tensor<i64>) -> tensor<2x2xi32>

Traits: AlwaysSpeculatableImplTrait, InferTensorType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
slice_sizes::mlir::DenseIntElementsAttr64-bit signless integer elements attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
start_indices variadic of 0D tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.dynamic_update_slice (mhlo::DynamicUpdateSliceOp)

DynamicUpdateSlice operation

Syntax:

operation ::= `mhlo.dynamic_update_slice` operands attr-dict `:` functional-type(operands, results)

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.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dynamic_update_slice

Example:

%result = mhlo.dynamic_update_slice %operand, %update, %start_indices0, %start_indices1
  : (tensor<4x4xi32>, tensor<2x2xi32>, tensor<i64>, tensor<i64>) -> tensor<4x4xi32>

Traits: AlwaysSpeculatableImplTrait, InferTensorType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
update ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
start_indices variadic of 0D tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.einsum (mhlo::EinsumOp)

Einsum operation

This operation is on its way out of StableHLO, so it is not included in the specification: https://github.com/openxla/stablehlo/issues/3

Informally, this operation does the same thing as TF's einsum: https://www.tensorflow.org/api_docs/python/tf/einsum

Example:

%result = "mhlo.einsum"(%lhs, %rhs) {
  einsum_config = "ab,bc->ac"
} : (tensor<4x16xf32>, tensor<16x4xf32>) -> tensor<4x4xf32>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
einsum_config::mlir::StringAttrstring attribute

Operands:

Operand Description
lhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
rhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.erf (mhlo::ErfOp)

Erf operation

Syntax:

operation ::= `mhlo.erf` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise erf operation on operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#erf

Example:

%result = mhlo.erf %operand : tensor<2x2xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

mhlo.exponential (mhlo::ExpOp)

Exp operation

Syntax:

operation ::= `mhlo.exponential` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise exponential operation on operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#exponential

Example:

%result = mhlo.exponential %operand : tensor<2x2xf64>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

mhlo.exponential_minus_one (mhlo::Expm1Op)

Expm1 operation

Syntax:

operation ::= `mhlo.exponential_minus_one` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise exponential minus one operation on operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#exponential_minus_one

Example:

%result = mhlo.exponential_minus_one %operand : tensor<2xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

mhlo.fft (mhlo::FftOp)

Fft operation

Performs the forward and inverse Fourier transforms for real and complex inputs/outputs.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#fft

Example:

%result = mhlo.fft %operand, type = FFT, length = [4] : (tensor<4xcomplex<f32>>) -> tensor<4xcomplex<f32>>

Traits: AlwaysSpeculatableImplTrait, InferTensorType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
fft_type::mlir::mhlo::FftTypeAttrXLA fast fourier transform type.
fft_length::mlir::DenseIntElementsAttr64-bit signless integer elements attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.floor (mhlo::FloorOp)

Floor operation

Syntax:

operation ::= `mhlo.floor` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise floor of operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#floor

Example:

%result = mhlo.floor %operand : tensor<2xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

mhlo.fusion (mhlo::FusionOp)

Fusion operation

This operation is private to the XLA compiler, so it is does not yet have a specification.

Informally, this operation consists of a group of basic ops (represented as a region attached to it). It serves as a hint to the backend that it is beneficial to emit the contained ops into a single loop nest or kernel.

Attributes:

AttributeMLIR TypeDescription
fusion_kind::mlir::mhlo::FusionKindAttrfusion kind
output_operand_aliases::mlir::ArrayAttrAliasing attribute for outputs and operands of Fusion

Operands:

Operand Description
inputs variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token

Results:

Result Description
results variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or nested tuple with any combination of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values

mhlo.gather (mhlo::GatherOp)

Gather operation

Gathers slices from operand tensor from offsets specified in start_indices and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#gather

Example:

%result = "mhlo.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 = dense<[0, 2, 2]> : tensor<3xi64>,
  indices_are_sorted = false
} : (tensor<2x3x4x2xi64>, tensor<2x2x3x2xi64>) -> tensor<2x2x3x2x2xi64>

Traits: AlwaysSpeculatableImplTrait, InferTensorType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
dimension_numbers::mlir::mhlo::GatherDimensionNumbersAttrAttribute that models the dimension information for gather
slice_sizes::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
indices_are_sorted::mlir::BoolAttrbool attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
start_indices ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.get_dimension_size (mhlo::GetDimensionSizeOp)

GetDimensionSize operation

Produces the size of the given dimension of the operand.

See https://github.com/openxla/stablehlo/blob/main/docs/spec.md#get_dimension_size

Example:

%result = mhlo.get_dimension_size %operand, dim = 1 : (tensor<2x3xf32>) -> tensor<i32>

Traits: AlwaysSpeculatableImplTrait, InferTensorType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
dimension::mlir::IntegerAttr64-bit signless integer attribute whose value is non-negative

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» tensor of 32-bit signless integer values

mhlo.get_tuple_element (mhlo::GetTupleElementOp)

GetTupleElement operation

Syntax:

operation ::= `mhlo.get_tuple_element` $operand `[` $index `]` attr-dict `:` functional-type(operands, results)

Extracts element at index position of the operand tuple and produces a result.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#get_tuple_element

Example:

%result = mhlo.get_tuple_element %operand[0] : (tuple<tensor<2xf32>, tuple<tensor<i32>>>) -> tensor<2xf32>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
index::mlir::IntegerAttr32-bit signless integer attribute whose value is non-negative

Operands:

Operand Description
operand nested tuple with any combination of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token or nested tuple with any combination of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values

mhlo.if (mhlo::IfOp)

If operation

Produces the output from executing exactly one branch from true_branch or false_branch depending on the value of pred.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#if

Example: %result = "mhlo.if"(%pred) ({ "mhlo.return"(%result_true_branch) : (tensor) -> () }, { "mhlo.return"(%result_false_branch) : (tensor) -> () }) : (tensor) -> tensor

Traits: RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfaces: InferTypeOpInterface

Operands:

Operand Description
pred ranked tensor of pred (AKA boolean or 1-bit integer) values

Results:

Result Description
«unnamed» variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token

mhlo.imag (mhlo::ImagOp)

Imag operation

Syntax:

operation ::= `mhlo.imag` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Extracts the imaginary part, element-wise, from the operand and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#imag

Example:

%result = mhlo.imag %operand : (tensor<2xcomplex<f32>>) -> tensor<2xf32>

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

mhlo.infeed (mhlo::InfeedOp)

Infeed operation

Reads data from the infeed and produces results.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#infeed

Example:

%results:2 = "mhlo.infeed"(%token) {
  infeed_config = ""
} : (!mhlo.token) -> (tensor<3x3x3xi32>, !mhlo.token)

Attributes:

AttributeMLIR TypeDescription
infeed_config::mlir::StringAttrstring attribute
layout::mlir::ArrayAttrarray attribute

Operands:

Operand Description
token token

Results:

Result Description
«unnamed» variadic of statically shaped tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or statically shaped tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token

mhlo.iota (mhlo::IotaOp)

Iota operation

Fills an output tensor with values in increasing order starting from zero along the iota_dimension dimension.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#iota

Example:

%output = mhlo.iota dim = 0 : tensor<4x5xi32>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
iota_dimension::mlir::IntegerAttr64-bit signless integer attribute whose value is non-negative

Results:

Result Description
output statically shaped tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements values

mhlo.is_finite (mhlo::IsFiniteOp)

IsFinite operation

Syntax:

operation ::= `mhlo.is_finite` $x attr-dict `:` functional-type(operands, results)

Performs element-wise check whether the value in x is finite (i.e. is neither +Inf, -Inf, nor NaN) and produces a y tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#is_finite

Example:

%y = mhlo.is_finite %x : (tensor<7xf32>) -> tensor<7xi1>

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
x ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

Results:

Result Description
y ranked tensor of pred (AKA boolean or 1-bit integer) values

mhlo.log (mhlo::LogOp)

Log operation

Syntax:

operation ::= `mhlo.log` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise logarithm operation on operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#log

Example:

%result = mhlo.log %operand : tensor<2x2xf64>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

mhlo.log_plus_one (mhlo::Log1pOp)

Log1p operation

Syntax:

operation ::= `mhlo.log_plus_one` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise logarithm plus one operation on operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#log_plus_one

Example:

%result = mhlo.log_plus_one %operand : tensor<6xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

mhlo.logistic (mhlo::LogisticOp)

Logistic operation

Syntax:

operation ::= `mhlo.logistic` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise logistic operation on operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#logistic

Example:

%result = mhlo.logistic %operand : tensor<2x2xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

mhlo.map (mhlo::MapOp)

Map operation

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#map

Example:

%result = "mhlo.map"(%input0, %input1) ({
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = mhlo.multiply %arg0, %arg1 : tensor<i32>
    mhlo.return %0 : tensor<i32>
}) {
  dimensions = dense<[0, 1]> : tensor<2xi64>
} : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>

Traits: InferTensorType, RecursiveMemoryEffects, SameOperandsAndResultShape, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface

Attributes:

AttributeMLIR TypeDescription
dimensions::mlir::DenseIntElementsAttr64-bit signless integer elements attribute

Operands:

Operand Description
inputs variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.maximum (mhlo::MaxOp)

Max operation

Syntax:

operation ::= `mhlo.maximum` $lhs `,` $rhs attr-dict
              `:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))

Performs element-wise max operation on tensors lhs and rhs and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#maximum

Example:

%result = mhlo.maximum %lhs, %rhs : tensor<4xf32>

Traits: AlwaysSpeculatableImplTrait, Commutative, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
lhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
rhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.minimum (mhlo::MinOp)

Min operation

Syntax:

operation ::= `mhlo.minimum` $lhs `,` $rhs attr-dict
              `:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))

Performs element-wise min operation on tensors lhs and rhs and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#minimum

Example:

%result = mhlo.minimum %lhs, %rhs : tensor<4xf32>

Traits: AlwaysSpeculatableImplTrait, Commutative, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
lhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
rhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.minimum_broadcast_shapes (mhlo::MinimumBroadcastShapesOp)

Minimizes the rank of two or more shapes to be broadcasted

Syntax:

operation ::= `mhlo.minimum_broadcast_shapes` $shapes attr-dict `:` type($shapes) `->` type($results)

Given two or more 1D tensors representing shapes, returns one 1D tensor for each operand, where operand i corresponds to output i.

The returned tensors have the property that they specify a shape which is a reshape of the corresponding input shape, and the broadcasted output shape (using shape::BroadcastOp) of the returned shapes is a reshape of the broadcasted output shape of the input shapes. Among all possibilities with this property, the one is chosen which minimizes the rank of each returned shape.

The general idea of this op is that it can be used for ops which have a broadcasting semantic to operate on shapes with a possibly smaller rank while preserving equivalence of the computed values. After computing the result of the op using reshaped operands, the result can be reshaped to the result that would have been originally computed.

Here is an example with two input shapes:

mhlo.minimum_broadcast_shapes [1, 2, 3, 1, 2, 1],
                                 [1, 1, 1, 2, 3] -> [6, 2, 1], [2, 3]

The broadcasted output shape of the operands is [1, 2, 3, 1, 2, 3], the broadcasted output shape of the outputs is [6, 2, 3]. These two shapes are reshapes of each other, and also each output is a reshape of the corresponding input.

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
shapes variadic of 1D tensor of index values

Results:

Result Description
results variadic of 1D tensor of index values

mhlo.multiply (mhlo::MulOp)

Mul operation

Syntax:

operation ::= `mhlo.multiply` $lhs `,` $rhs attr-dict
              `:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))

Performs element-wise product of two tensors lhs and rhs and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#multiply

Example:

%result = mhlo.multiply %lhs, %rhs : tensor<2xi32>

Traits: AlwaysSpeculatableImplTrait, Commutative, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
lhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
rhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.negate (mhlo::NegOp)

Neg operation

Syntax:

operation ::= `mhlo.negate` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise negation of operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#negate

Example:

%result = mhlo.negate %operand : tensor<2x3xi32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
result ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.not (mhlo::NotOp)

Not operation

Syntax:

operation ::= `mhlo.not` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise NOT of tensor operand of type integer and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#not

Example:

%result = mhlo.not %operand : tensor<5x3x1xi1>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
result ranked tensor of pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

mhlo.optimization_barrier (mhlo::OptimizationBarrierOp)

OptimizationBarrier operation

Syntax:

operation ::= `mhlo.optimization_barrier` attr-dict ($operand^ `:` custom<PairwiseOpType>(type($operand), type($result))):(`(` `)`)?

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.

See https://github.com/openxla/stablehlo/blob/main/docs/spec.md#optimization_barrier

Example:

%result0, %result1 = mhlo.optimization_barrier %operand0, %operand1 : tensor<f32>, tensor<f32>

Traits: AlwaysSpeculatableImplTrait, HLO_PairwiseSameOperandAndResultType

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token

Results:

Result Description
result variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token

mhlo.or (mhlo::OrOp)

Or operation

Syntax:

operation ::= `mhlo.or` $lhs `,` $rhs attr-dict
              `:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))

Performs element-wise OR of two tensors lhs and rhs and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#or

Example:

%result = mhlo.or %lhs, %rhs : tensor<2xi1>

Traits: AlwaysSpeculatableImplTrait, Commutative, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
lhs ranked tensor of pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values
rhs ranked tensor of pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.outfeed (mhlo::OutfeedOp)

Outfeed operation

Writes inputs to the outfeed and produces a result token.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#outfeed

Example:

%result = "mhlo.outfeed"(%input0, %token) {
  outfeed_config = ""
} : (tensor<3x3x3xi32>, !mhlo.token) -> !mhlo.token

Interfaces: InferTypeOpInterface

Attributes:

AttributeMLIR TypeDescription
outfeed_config::mlir::StringAttrstring attribute

Operands:

Operand Description
inputs variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
token token

Results:

Result Description
«unnamed» token

mhlo.pad (mhlo::PadOp)

Pad operation

Expands operand by padding around the tensor as well as between the elements of the tensor with the given padding_value.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#pad

Example:

%0 = mhlo.pad %arg0, %arg1, low = [0, 1], high = [2, 1], interior = [1, 2]
  : (tensor<2x3xi32>, tensor<i32>) -> tensor<5x9xi32>

Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
edge_padding_low::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
edge_padding_high::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
interior_padding::mlir::DenseIntElementsAttr64-bit signless integer elements attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
padding_value ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.partition_id (mhlo::PartitionIdOp)

PartitionId operation

Syntax:

operation ::= `mhlo.partition_id` attr-dict `:` type(results)

Produces partition_id of the current process.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#partition_id

Example:

%result = mhlo.partition_id : tensor<ui32>

Interfaces: InferTypeOpInterface

Results:

Result Description
«unnamed» ranked tensor of 32-bit unsigned integer values

mhlo.popcnt (mhlo::PopulationCountOp)

PopulationCount operation

Syntax:

operation ::= `mhlo.popcnt` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#popcnt

Example:

%result = mhlo.popcnt %operand : tensor<4xi8>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
result ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

mhlo.power (mhlo::PowOp)

Pow operation

Syntax:

operation ::= `mhlo.power` $lhs `,` $rhs attr-dict
              `:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))

Performs element-wise exponentiation of lhs tensor by rhs tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#power

Example:

%result = mhlo.power %lhs, %rhs : tensor<6xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
lhs ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
rhs ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
result ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.real (mhlo::RealOp)

Real operation

Syntax:

operation ::= `mhlo.real` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Extracts the real part, element-wise, from the operand and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#real

Example:

%result = mhlo.real %operand : (tensor<2xcomplex<f32>>) -> tensor<2xf32>

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

mhlo.real_dynamic_slice (mhlo::RealDynamicSliceOp)

RealDynamicSlice operation

Syntax:

operation ::= `mhlo.real_dynamic_slice` operands attr-dict `:` functional-type(operands, results)

This operation is a work in progress, so it is not yet included in the specification: https://github.com/openxla/stablehlo/issues/8

Informally, this operation does the same thing as SliceOp except that start_indices, limit_indices and strides are specified dynamically: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#slice

Example:

%result = mhlo.real_dynamic_slice %operand,
            %start_indices, %limit_indices, %strides
       : (tensor<256x?xf32>, tensor<2xindex>, tensor<2xindex>, tensor<2xindex>) -> tensor<256x?xf32>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
start_indices 1D tensor of index or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values
limit_indices 1D tensor of index or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values
strides 1D tensor of index or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.recv (mhlo::RecvOp)

Recv operation

Receives data from a channel with channel_id and produces results.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#recv

Example:

%results:2 = "mhlo.recv"(%token) {
  // channel_id = 5 : i64,
  // channel_type = #stablehlo<channel_type HOST_TO_DEVICE>,
  channel_handle = #mhlo.channel_handle<handle = 5, type = 3>,
  is_host_transfer = true
} : (!mhlo.token) -> (tensor<3x4xi32>, !mhlo.token)

Attributes:

AttributeMLIR TypeDescription
channel_handle::mlir::mhlo::ChannelHandleAttrtwo 64-bit integers 'handle' and 'type'
is_host_transfer::mlir::BoolAttrbool attribute

Operands:

Operand Description
token token

Results:

Result Description
«unnamed» variadic of statically shaped tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or statically shaped tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token

mhlo.reduce (mhlo::ReduceOp)

Reduce operation

Applies a reduction function body to inputs and init_values along the dimensions and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reduce

Example:

%result = "mhlo.reduce"(%input, %init_value) ({
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = "mhlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
    "mhlo.return"(%0) : (tensor<i32>) -> ()
}) {
  dimensions = dense<1> : tensor<1xi64>
} : (tensor<1x6xi32>, tensor<i32>) -> tensor<1xi32>

Traits: InferTensorType, RecursiveMemoryEffects, SameVariadicOperandSize, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface

Attributes:

AttributeMLIR TypeDescription
dimensions::mlir::DenseIntElementsAttr64-bit signless integer elements attribute

Operands:

Operand Description
inputs variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
init_values variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.reduce_precision (mhlo::ReducePrecisionOp)

ReducePrecision operation

Syntax:

operation ::= `mhlo.reduce_precision` $operand `,` `format` `=` custom<ExponentMantissa>($exponent_bits, $mantissa_bits)
              attr-dict `:` custom<SameOperandsAndResultType>(type($operand), type($output))

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.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reduce_precision

Example:

%output = mhlo.reduce_precision %operand, format = e5m2 : tensor<6xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
exponent_bits::mlir::IntegerAttr32-bit signless integer attribute whose value is positive
mantissa_bits::mlir::IntegerAttr32-bit signless integer attribute whose value is non-negative

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

Results:

Result Description
output ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

mhlo.reduce_scatter (mhlo::ReduceScatterOp)

ReduceScatter operation

Within each process group in the 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.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reduce_scatter

Example:

%result = "mhlo.reduce_scatter"(%operand) ({
  ^bb0(%arg0: tensor<f32>, %arg1: tensor<f32>):
  %0 = mhlo.add %arg0, %arg1 : tensor<f32>
  mhlo.return %0 : tensor<f32>
}) {
  scatter_dimension = 1 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  // channel_id = 0
  channel_handle = #mhlo.channel_handle<handle = 0, type = 0>
  // use_global_device_ids = false
} : (tensor<2x4xf32>) -> tensor<2x2xf32>

Attributes:

AttributeMLIR TypeDescription
scatter_dimension::mlir::IntegerAttr64-bit signless integer attribute whose value is non-negative
replica_groups::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
channel_handle::mlir::mhlo::ChannelHandleAttrtwo 64-bit integers 'handle' and 'type'
use_global_device_ids::mlir::UnitAttrunit attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.reduce_window (mhlo::ReduceWindowOp)

ReduceWindow operation

Applies a reduction function body to windows of inputs and init_values and produces results.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reduce_window

Example:

%result = "mhlo.reduce_window"(%input, %init_value) ({
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = mhlo.add %arg0, %arg1 : tensor<i32>
    mhlo.return %0 : tensor<i32>
}) {
  window_dimensions = dense<[2, 1]> : tensor<2xi64>,
  window_strides = dense<[4, 1]> : tensor<2xi64>,
  base_dilations = dense<[2, 1]> : tensor<2xi64>,
  window_dilations = dense<[3, 1]> : tensor<2xi64>,
  padding = dense<[[2, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<3x2xi32>, tensor<i32>) -> tensor<2x2xi32>

Traits: InferTensorType, RecursiveMemoryEffects, SameVariadicOperandSize, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface

Attributes:

AttributeMLIR TypeDescription
window_dimensions::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
window_strides::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
base_dilations::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
window_dilations::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
padding::mlir::DenseIntElementsAttr64-bit signless integer elements attribute

Operands:

Operand Description
inputs variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
init_values variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.remainder (mhlo::RemOp)

Rem operation

Syntax:

operation ::= `mhlo.remainder` $lhs `,` $rhs attr-dict
              `:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))

Performs element-wise remainder of dividend lhs and divisor rhs tensors and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#remainder

Example:

%result = mhlo.remainder %lhs, %rhs : tensor<4xi64>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
lhs ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
rhs ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
result ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.replica_id (mhlo::ReplicaIdOp)

ReplicaId operation

Syntax:

operation ::= `mhlo.replica_id` attr-dict `:` type(results)

Produces replica_id of the current process.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#replica_id

Example:

%result = mhlo.replica_id : tensor<ui32>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Results:

Result Description
«unnamed» ranked tensor of 32-bit unsigned integer values

mhlo.reshape (mhlo::ReshapeOp)

Reshape operation

Syntax:

operation ::= `mhlo.reshape` operands attr-dict `:` functional-type(operands, results)

Performs reshape of operand tensor to a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reshape

Example:

%result = mhlo.reshape %operand : (tensor<2xf32>) -> tensor<1x2xf32>

Traits: AlwaysSpeculatableImplTrait, HLO_CompatibleOperandsAndResultElementType

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» statically shaped tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.return (mhlo::ReturnOp)

_This operation is a work in progress, so it is not yet included in the specification: https://github.com/openxla/stablehlo/issues/425

Informally, this operation serves as a terminator for regions defined by
the StableHLO ops. Non-StableHLO ops, e.g. `func.func`, have their own
terminators, e.g. `func.return`.

Example:

    ```mlir
    %result = "mhlo.reduce"(%input, %init_value) ({
      ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
        %0 = "mhlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
        "mhlo.return"(%0) : (tensor<i32>) -> ()
    }) {
      dimensions = dense<1> : tensor<1xi64>
    } : (tensor<1x6xi32>, tensor<i32>) -> tensor<1xi32>
    ```_


Syntax:

```

operation ::= mhlo.return $results attr-dict (: type($results)^)?



Traits: `AlwaysSpeculatableImplTrait`, `Terminator`

Interfaces: `ConditionallySpeculatable`, `NoMemoryEffect (MemoryEffectOpInterface)`

Effects: `MemoryEffects::Effect{}`

#### Operands:

| Operand | Description |
| :-----: | ----------- |
| `results` | variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token or nested tuple with any combination of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values


### `mhlo.reverse` (mhlo::ReverseOp)

_Reverse operation_

Reverses the order of elements in the `operand` along the specified
`dimensions` and produces a `result` tensor.

See:
<a href="https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reverse">https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reverse</a>

Example:
```mlir
%result = mhlo.reverse %operand, dims = [1] : tensor<3x2xi32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
dimensions::mlir::DenseIntElementsAttr64-bit signless integer elements attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.rng (mhlo::RngOp)

Rng operation

Generates random numbers using the rng_distribution algorithm and produces a result tensor of a given shape shape.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#rng

Example:

%result = mhlo.rng %a, %b, %shape, distribution = NORMAL : (tensor<i32>, tensor<i32>, tensor<2xi64>) -> tensor<3x3xi32>

Traits: InferTensorType

Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface

Attributes:

AttributeMLIR TypeDescription
rng_distribution::mlir::mhlo::RngDistributionAttrXLA PRNG distribution to be used.

Operands:

Operand Description
a 0D tensor of pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
b 0D tensor of pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
shape 1D tensor of index or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
result ranked tensor of pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

mhlo.rng_bit_generator (mhlo::RngBitGeneratorOp)

RngBitGenerator operation

Returns an output filled with uniform random data and an updated output state output_state given an initial state initial_state using the pseudorandom number generator algorithm rng_algorithm.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#rng_bit_generator

Example:

%output_state, %output = mhlo.rng_bit_generator %initial_state, algorithm = THREE_FRY : (tensor<2xui64>) -> (tensor<2xui64>, tensor<2x2xui64>)

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
rng_algorithm::mlir::mhlo::RngAlgorithmAttrXLA PRNG algorithm to be used.

Operands:

Operand Description
initial_state ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

Results:

Result Description
output_state ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
output statically shaped tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

mhlo.round_nearest_afz (mhlo::RoundOp)

Round operation

Syntax:

operation ::= `mhlo.round_nearest_afz` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise rounding towards the nearest integer, breaking ties away from zero, on the operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#round_nearest_afz

Example:

%result = mhlo.round_nearest_afz %operand : tensor<5xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

mhlo.round_nearest_even (mhlo::RoundNearestEvenOp)

RoundNearestEven operation

Syntax:

operation ::= `mhlo.round_nearest_even` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise rounding towards the nearest integer, breaking ties towards the even integer, on the operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#round_nearest_even

Example:

%result = mhlo.round_nearest_even %operand : tensor<5xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

mhlo.rsqrt (mhlo::RsqrtOp)

Rsqrt operation

Syntax:

operation ::= `mhlo.rsqrt` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise reciprocal square root operation on operand tensor and produces a result tensor, implementing the rSqrt operation from the IEEE-754 specification.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#rsqrt

Example:

%result = mhlo.rsqrt %operand : tensor<2x2xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

mhlo.scatter (mhlo::ScatterOp)

Scatter operation

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.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#scatter

Example:

%result = "mhlo.scatter"(%input, %scatter_indices, %update) ({
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = mhlo.add %arg0, %arg1 : tensor<i32>
    mhlo.return %0 : tensor<i32>
}) {
  scatter_dimension_numbers = #mhlo.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>

Traits: RecursiveMemoryEffects, SameVariadicOperandSize

Interfaces: InferTypeOpInterface

Attributes:

AttributeMLIR TypeDescription
scatter_dimension_numbers::mlir::mhlo::ScatterDimensionNumbersAttrAttribute that models the dimension information for scatter
indices_are_sorted::mlir::BoolAttrbool attribute
unique_indices::mlir::BoolAttrbool attribute

Operands:

Operand Description
inputs variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
scatter_indices ranked tensor of integer or index values
updates variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.select (mhlo::SelectOp)

Select operation

Syntax:

operation ::= `mhlo.select` operands attr-dict `:`
              custom<SelectOpType>(type($pred), type($on_true), type($on_false), type($result))

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.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#select

Example:

%result = mhlo.select %pred, %on_true, %on_false : tensor<2x2xi1>, tensor<2x2xi32>

Traits: AlwaysSpeculatableImplTrait, HLO_BroadcastingElementwise, InferTensorType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
pred ranked tensor of pred (AKA boolean or 1-bit integer) values
on_true ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
on_false ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.select_and_scatter (mhlo::SelectAndScatterOp)

SelectAndScatter operation

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.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#select_and_scatter

Example:

%result = "mhlo.select_and_scatter"(%operand, %source, %init_value) ({
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = "mhlo.compare"(%arg0, %arg1) {
      comparison_direction = #stablehlo<comparison_direction GE>
    } : (tensor<i32>, tensor<i32>) -> tensor<i1>
    "mhlo.return"(%0) : (tensor<i1>) -> ()
}, {
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = "mhlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
    "mhlo.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>

Traits: RecursiveMemoryEffects

Interfaces: InferTypeOpInterface

Attributes:

AttributeMLIR TypeDescription
window_dimensions::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
window_strides::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
padding::mlir::DenseIntElementsAttr64-bit signless integer elements attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
source ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
init_value ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.send (mhlo::SendOp)

Send operation

Sends inputs to a channel channel_id and produces a result token.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#send

Example:

%result = "mhlo.send"(%operand, %token) {
  // channel_id = 5 : i64,
  // channel_type = #stablehlo<channel_type DEVICE_TO_HOST>,
  channel_handle = #mhlo.channel_handle<handle = 5, type = 2>,
  is_host_transfer = true
} : (tensor<3x4xi32>, !mhlo.token) -> !mhlo.token

Interfaces: InferTypeOpInterface

Attributes:

AttributeMLIR TypeDescription
channel_handle::mlir::mhlo::ChannelHandleAttrtwo 64-bit integers 'handle' and 'type'
is_host_transfer::mlir::BoolAttrbool attribute

Operands:

Operand Description
inputs variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
token token

Results:

Result Description
«unnamed» token

mhlo.set_dimension_size (mhlo::SetDimensionSizeOp)

SetDimensionSize operation

This operation is a work in progress, so it is not yet included in the specification: https://github.com/openxla/stablehlo/issues/8

Informally, this operation does the same thing as XLA's SetDimensionSize: https://www.tensorflow.org/xla/operation_semantics#setdimensionsize

Example:

%0 = mhlo.set_dimension_size %arg0, %arg1, dim = 1 : (tensor<4x2xf32>, tensor<i32>) -> tensor<4x2xf32>

Traits: AlwaysSpeculatableImplTrait, InferTensorType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
dimension::mlir::IntegerAttr64-bit signless integer attribute whose value is non-negative

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
size tensor of 32-bit signless integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.shift_left (mhlo::ShiftLeftOp)

ShiftLeft operation

Syntax:

operation ::= `mhlo.shift_left` $lhs `,` $rhs attr-dict
              `:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#shift_left

Example:

%result = mhlo.shift_left %lhs, %rhs : tensor<6xi8>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
lhs ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values
rhs ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
result ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

mhlo.shift_right_arithmetic (mhlo::ShiftRightArithmeticOp)

ShiftRightArithmetic operation

Syntax:

operation ::= `mhlo.shift_right_arithmetic` $lhs `,` $rhs attr-dict
              `:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#shift_right_arithmetic

Example:

%result = mhlo.shift_right_arithmetic %lhs, %rhs : tensor<6xi8>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
lhs ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values
rhs ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
result ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

mhlo.shift_right_logical (mhlo::ShiftRightLogicalOp)

ShiftRightLogical operation

Syntax:

operation ::= `mhlo.shift_right_logical` $lhs `,` $rhs attr-dict
              `:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#shift_right_logical

Example:

%result = mhlo.shift_right_logical %lhs, %rhs : tensor<6xi8>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
lhs ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values
rhs ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
result ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

mhlo.sign (mhlo::SignOp)

Sign operation

Syntax:

operation ::= `mhlo.sign` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Returns the sign of the operand element-wise and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#sign

Example:

%result = mhlo.sign %operand : tensor<7xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of 2/4/8/16/32/64-bit signless integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

Results:

Result Description
result ranked tensor of 2/4/8/16/32/64-bit signless integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

mhlo.sine (mhlo::SineOp)

Sine operation

Syntax:

operation ::= `mhlo.sine` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise sine operation on operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#sine

Example:

%result = mhlo.sine %operand : tensor<2xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

mhlo.slice (mhlo::SliceOp)

Slice operation

Extracts a slice from the operand using statically-computed starting indices and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#slice

Example:

%result = "mhlo.slice" (%operand) {
  start_indices = dense<[1, 2]> : tensor<2xi64>,
  limit_indices = dense<[3, 4]> : tensor<2xi64>,
  strides = dense<1> : tensor<2xi64>
} : (tensor<3x4xi64>) -> tensor<2x2xi64>

Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
start_indices::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
limit_indices::mlir::DenseIntElementsAttr64-bit signless integer elements attribute
strides::mlir::DenseIntElementsAttr64-bit signless integer elements attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.sort (mhlo::SortOp)

Sort operation

Sorts a variadic number of tensors in inputs together, according to a custom comparator, along the given dimension and produces a variadic number of tensors as results.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#sort

Example:

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

Traits: InferTensorType, RecursiveMemoryEffects, SameOperandsAndResultShape

Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface

Attributes:

AttributeMLIR TypeDescription
dimension::mlir::IntegerAttr64-bit signless integer attribute
is_stable::mlir::BoolAttrbool attribute

Operands:

Operand Description
inputs variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.sparse_dot (mhlo::SparseDotOp)

Sparse dot operation

Similar to dot_general operation, with one or both of the operands being sparse. An additional argument provides sparsity meta information. Disclaimer: this op is experimental / a work in progress.

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
lhs_sparsity::mlir::mhlo::SparsityDescriptorAttrDescribes structured (N:M) sparsity configuration
rhs_sparsity::mlir::mhlo::SparsityDescriptorAttrDescribes structured (N:M) sparsity configuration
dot_dimension_numbers::mlir::mhlo::DotDimensionNumbersAttrAttribute that models the dimension information for dot.
precision_config::mlir::ArrayAttrPrecision Config attribute

Operands:

Operand Description
lhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
rhs ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
meta variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.sqrt (mhlo::SqrtOp)

Sqrt operation

Syntax:

operation ::= `mhlo.sqrt` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise square root operation on operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#sqrt

Example:

%result = mhlo.sqrt %operand : tensor<2x2xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

mhlo.stochastic_convert (mhlo::StochasticConvertOp)

StochasticConvert operation

This operation is a work in progress, so it is not yet included in the specification: https://github.com/openxla/stablehlo/issues/295

Informally, this operation performs element-wise conversion of values from a bigger type to a smaller one with stochastic rounding using the random number passed in.

Traits: AlwaysSpeculatableImplTrait, Elementwise

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values
random ranked tensor of 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.subtract (mhlo::SubtractOp)

Subtract operation

Syntax:

operation ::= `mhlo.subtract` $lhs `,` $rhs attr-dict
              `:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))

Performs element-wise subtraction of two tensors lhs and rhs and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#subtract

Example:

%result = mhlo.subtract %lhs, %rhs : tensor<2xi32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
lhs ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
rhs ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
result ranked tensor of 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.tan (mhlo::TanOp)

Tan operation

Syntax:

operation ::= `mhlo.tan` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

This operation is a work in progress, so it is not yet included in the specification: https://github.com/openxla/stablehlo/issues/954

Informally, this operation returns Tan(operand) element-wise.

Example:

%0 = mhlo.tan %arg0 : tensor<2xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements values

mhlo.tanh (mhlo::TanhOp)

Tanh operation

Syntax:

operation ::= `mhlo.tanh` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

Performs element-wise hyperbolic tangent operation on operand tensor and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#tanh

Example:

%result = mhlo.tanh %operand : tensor<2xf32>

Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values

mhlo.topk (mhlo::TopKOp)

TopK operation

Syntax:

operation ::= `mhlo.topk` `(`$operand `,` `k` `=` $k (`,` `largest` `=` $largest^)? `)` attr-dict `:`
              type($operand) `->` `(`type($values)`,` type($indices)`)`

Returns top k values and their indices, along the last dimension of the operand if largest=true or the bottom k values if largest=false.

See: https://www.tensorflow.org/xla/operation_semantics#top-k

Example:

%values, %indices = mhlo.topk(%operand, k=5, largest=true)
  : tensor<100xf32> -> (tensor<5xf32>, tensor<5xi32>)

Traits: InferTensorType, RecursiveMemoryEffects

Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface

Attributes:

AttributeMLIR TypeDescription
k::mlir::IntegerAttr64-bit signless integer attribute
largest::mlir::BoolAttrbool attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
values ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
indices ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.torch_index_select (mhlo::TorchIndexSelectOp)

TorchIndexSelect operation

This operation is on its way out of StableHLO, so it is not included in the specification: https://github.com/openxla/stablehlo/issues/3

Informally, this operation does the same thing as PyTorch's index_select, augmented with support for batch dimensions: https://pytorch.org/docs/stable/generated/torch.index_select.html

The batch_dims attribute specifies the number of major batch dimensions (0 or more) that act like a multidimensional loop over both the operand and the index.

Example:

%result = "mhlo.torch_index_select"(%operand, %index) {
  dim = 2 : i64,
  batch_dims = 1 : i64
} : (tensor<8x128x3072x64xf32>, tensor<8x16x1024xi32>) -> tensor<8x128x16x1024x64xf32>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
dim::mlir::IntegerAttr64-bit signless integer attribute
batch_dims::mlir::IntegerAttr64-bit signless integer attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values
index ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.trace (mhlo::TraceOp)

Trace operation

Syntax:

operation ::= `mhlo.trace` $operand `,` $tag attr-dict `:` type($operand)

This operation is on its way out of StableHLO, so it is not included in the specification: https://github.com/openxla/stablehlo/issues/604

It is not used by JAX, PyTorch or TensorFlow, so it looks like we should've classified it as "Private to XLA" and not included it in StableHLO in the first place. With that in mind, its semantics will not be documented here.

Example:

mhlo.trace %arg0, "In test code." : tensor<5x1x5xi32>

Attributes:

AttributeMLIR TypeDescription
tag::mlir::StringAttrstring attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.transpose (mhlo::TransposeOp)

Transpose operation

Permutes the dimensions of operand tensor using permutation and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#transpose

Example:

%0 = mhlo.transpose %arg0, dims = [2, 1, 0] : (tensor<1x2x3xi32>) -> tensor<3x2x1xi32>

Traits: AlwaysSpeculatableImplTrait, HLO_CompatibleOperandsAndResultElementType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
permutation::mlir::DenseIntElementsAttr64-bit signless integer elements attribute

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.triangular_solve (mhlo::TriangularSolveOp)

TriangularSolve operation

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#triangular_solve

Example:

%result = "mhlo.triangular_solve"(%a, %b) {
  left_side = true,
  lower = true,
  unit_diagonal = false,
  transpose_a = #stablehlo<transpose NO_TRANSPOSE>
} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>

Traits: AlwaysSpeculatableImplTrait, InferTensorType, SameOperandsAndResultElementType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
left_side::mlir::BoolAttrbool attribute
lower::mlir::BoolAttrbool attribute
unit_diagonal::mlir::BoolAttrbool attribute
transpose_a::mlir::mhlo::TransposeAttrTranspose options

Operands:

Operand Description
a ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements values
b ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements values

Results:

Result Description
«unnamed» ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements values

mhlo.tuple (mhlo::TupleOp)

Tuple operation

Syntax:

operation ::= `mhlo.tuple` $val attr-dict `:` custom<TupleOpType>(type($val), type($result))

Produces a result tuple from values val.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#tuple

Example:

%result = mhlo.tuple %val0, %val1 : tuple<tensor<2xf32>, tuple<tensor<i32>>>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
val variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token or nested tuple with any combination of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values

Results:

Result Description
result nested tuple with any combination of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token values

mhlo.uniform_dequantize (mhlo::UniformDequantizeOp)

UniformDequantize operation

Syntax:

operation ::= `mhlo.uniform_dequantize` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

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

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#uniform_dequantize

Example:

%result = mhlo.uniform_dequantize %operand : (tensor<16x16x!quant.uniform<i8:f32, 34.0:16>>) -> tensor<16x16xf32>

Traits: AlwaysSpeculatableImplTrait, Elementwise, InferTensorType, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type values

mhlo.uniform_quantize (mhlo::UniformQuantizeOp)

UniformQuantize operation

Syntax:

operation ::= `mhlo.uniform_quantize` $operand attr-dict
              `:` custom<SameOperandsAndResultType>(type($operand), type($result))

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.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#uniform_quantize

Example:

%result = mhlo.uniform_quantize %operand : (tensor<16x16xf32>) -> tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
operand ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Results:

Result Description
result ranked tensor of 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

mhlo.while (mhlo::WhileOp)

While operation

Produces the output from executing body function 0 or more times while the cond function outputs true.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#while

Example:

%results0, %results1 = "mhlo.while"(%operand0, %operand1) ({
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = "mhlo.compare"(%arg0, %arg1) {
      comparison_direction = #stablehlo<comparison_direction LT>
    } : (tensor<i32>, tensor<i32>) -> tensor<i1>
    "mhlo.return"(%0) : (tensor<i1>) -> ()
}, {
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = "mhlo.add"(%arg0, %constant0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
    "mhlo.return"(%0, %arg1) : (tensor<i32>, tensor<i32>) -> ()
}) : (tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)

Traits: RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfaces: InferTypeOpInterface, OpAsmOpInterface

Operands:

Operand Description
operand variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token

Results:

Result Description
«unnamed» variadic of ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer values or ranked tensor of 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values or token

mhlo.xla.rng_get_and_update_state (mhlo::XlaRngGetAndUpdateStateOp)

XlaRngGetAndUpdateState operation

Syntax:

operation ::= `mhlo.xla.rng_get_and_update_state` attr-dict

This operation is private to the XLA compiler, so it is does not yet have a specification.

Informally, this operation represents the change of the global random number generator state for rng instructions. The global state is incremented by delta and the old state is returned.

The output is currently defined for a single output type. If this changes in the future to support multiple types, lowering to use of a global memref must ensure that a single memref is still used and updated appropriately.

Interfaces: InferTypeOpInterface

Attributes:

AttributeMLIR TypeDescription
delta::mlir::IntegerAttr64-bit signless integer attribute

Results:

Result Description
«unnamed» statically shaped tensor of 64-bit unsigned integer values

mhlo.xor (mhlo::XorOp)

Xor operation

Syntax:

operation ::= `mhlo.xor` $lhs `,` $rhs attr-dict
              `:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))

Performs element-wise XOR of two tensors lhs and rhs and produces a result tensor.

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#xor

Example:

%result = mhlo.xor %lhs, %rhs : tensor<2xi32>

Traits: AlwaysSpeculatableImplTrait, Commutative, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand Description
lhs ranked tensor of pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values
rhs ranked tensor of pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer values

Results:

Result Description
result ranked tensor of f8E3M4 type or f8E4M3 type or f8E4M3FN type or f8E4M3FNUZ type or f8E4M3B11FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or pred (AKA boolean or 1-bit integer) or 2/4/8/16/32/64-bit signless integer or 2/4/8/16/32/64-bit unsigned integer or complex type with 32-bit float or 64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values

Attributes

ArgResultAliasAttr

Attribute that models the alias relationship of entry function argument

This attribute captures the alias relationship of an MHLO main function argument to one of the results, denoted by resultIndex. The argTupleIndices and resultTupleIndices are used to index into nested tuples in operand and result respectively. If isMustAlias is true then the operand-result pair must alias.

This is meant to be used as an attribute on a function argument in MHLO. For example, in the following code it expresses that %arg1 may alias 0-th result.

func @main(%arg0: tensor<2xf32>, %arg1: tensor<3xf32> {mhlo.result_alias =
    mhlo.result_alias<result_index = [2], ...>}
  ) -> tensor<2xf32>, tensor<3xf32> {
  // function body ...
}

Parameters:

Parameter C++ type Description
argTupleIndices ::llvm::ArrayRef<int64_t> Dimension
resultIndex int64_t
resultTupleIndices ::llvm::ArrayRef<int64_t> Dimension
isMustAlias bool

ChannelHandleAttr

two 64-bit integers 'handle' and 'type'

Syntax:

#mhlo.channel_handle<
  int64_t,   # handle
  int64_t   # type
>

Parameters:

Parameter C++ type Description
handle int64_t
type int64_t

ComparisonDirectionAttr

Which comparison operation to perform.

Syntax:

#mhlo.comparison_direction<
  ::mlir::mhlo::ComparisonDirection   # value
>

Enum cases:

  • EQ (EQ)
  • NE (NE)
  • GE (GE)
  • GT (GT)
  • LE (LE)
  • LT (LT) #### Parameters:
Parameter C++ type Description
value ::mlir::mhlo::ComparisonDirection an enum of type ComparisonDirection

ComparisonTypeAttr

Which comparison type to use.

Syntax:

#mhlo.comparison_type<
  ::mlir::mhlo::ComparisonType   # value
>

Enum cases:

  • NOTYPE (NOTYPE)
  • FLOAT (FLOAT)
  • TOTALORDER (TOTALORDER)
  • SIGNED (SIGNED)
  • UNSIGNED (UNSIGNED) #### Parameters:
Parameter C++ type Description
value ::mlir::mhlo::ComparisonType an enum of type ComparisonType

ConvDimensionNumbersAttr

Structure of dimension information for conv op

Parameters:

Parameter C++ type Description
inputBatchDimension int64_t
inputFeatureDimension int64_t
inputSpatialDimensions ::llvm::ArrayRef<int64_t> Dimension
kernelInputFeatureDimension int64_t
kernelOutputFeatureDimension int64_t
kernelSpatialDimensions ::llvm::ArrayRef<int64_t> Dimension
outputBatchDimension int64_t
outputFeatureDimension int64_t
outputSpatialDimensions ::llvm::ArrayRef<int64_t> Dimension

CrossProgramPrefetchAttr

Argument that is prefetched from another program

Syntax:

#mhlo.cross_program_prefetch<
  int64_t,   # parameter
  ::llvm::ArrayRef<int64_t>,   # indices
  std::optional<int64_t>   # offset
>

This attribute captures an argument that is prefetched from another program. For a given CrossProgramPrefetchAttr, parameter tells us which argument of the main function of the module is prefetched, and indices is a shape index telling us what subshape of that argument is prefetched.

A shape has a subshape iff it is a tuple. In that case, the subshape of the tuple by indices is the shape achieved after indexing by each element of indices in turn. For example, the [1,0] subshape of tuple<tuple<token, token>, tuple<tensor<i32>, token>> is tensor<i32>.

An empty value for indices means the whole shape is prefetched.

For example,

module attributes { mhlo.cross_program_prefetch = [ #mhlo.cross_program_prefetch< parameter = 0, indices = [0]> ]} {
  func.func @copy(%arg0 : tuple<tensor<2x3xi32>, tensor<i32>>) -> tuple<tensor<2x3xi32>, tensor<i32>> {
    %0 = "mhlo.copy"(%arg0) {is_cross_program_prefetch}
    return %0 : tuple<tensor<2x3xi32>, tensor<i32>>
  }
  func.func @main(%arg0 : tuple<tensor<2x3xi32>, tensor<i32>>) -> tuple<tensor<2x3xi32>, tensor<i32>> {
    %1 = "mhlo.async_start"(%arg0) {called_computation=@copy}
    %2 = "mhlo.async_done"(%1) {called_computation=@copy}
    return %2 : tuple<tensor<2x3xi32>, tensor<i32>>
  }
}

The parameter = 0 tells us that the async copy of the 0th parameter is a cross_program_prefetch, while the index of [0] tells us that the 0th element of the tuple is prefetched while the other element of the tuple is not.

Parameters:

Parameter C++ type Description
parameter int64_t
indices ::llvm::ArrayRef<int64_t> Dimension
offset std::optional<int64_t>

CustomCallScheduleAttr

Specifies the desired schedule for the custom-call.

Syntax:

#mhlo.custom_call_schedule<
  ::mlir::mhlo::CustomCallSchedule   # value
>

Enum cases:

  • NONE (NONE)
  • LATEST (LATEST)
  • EARLIEST (EARLIEST) #### Parameters:
Parameter C++ type Description
value ::mlir::mhlo::CustomCallSchedule an enum of type CustomCallSchedule

DequantizeModeAttr

Dequantization mode. Only MIN_COMBINED is supported.

Syntax:

#mhlo.dequantize_mode<
  ::mlir::mhlo::DequantizeMode   # value
>

Enum cases:

  • MIN_COMBINED (MIN_COMBINED) #### Parameters:
Parameter C++ type Description
value ::mlir::mhlo::DequantizeMode an enum of type DequantizeMode

DomainKindAttr

Kind of domain metatdata attached to an HLO domain.

Syntax:

#mhlo.kind<
  ::mlir::mhlo::DomainKind   # value
>

Enum cases:

  • sharding (sharding) #### Parameters:
Parameter C++ type Description
value ::mlir::mhlo::DomainKind an enum of type DomainKind

DotAlgorithmAttr

Attribute that models the algorithm constraints to use for computing dot.

Syntax:

#mhlo.dot_algorithm<
  Type,   # lhsPrecisionType
  Type,   # rhsPrecisionType
  Type,   # accumulationType
  int64_t,   # lhsComponentCount
  int64_t,   # rhsComponentCount
  int64_t,   # numPrimitiveOperations
  bool   # allowImpreciseAccumulation
>

Parameters:

Parameter C++ type Description
lhsPrecisionType Type
rhsPrecisionType Type
accumulationType Type
lhsComponentCount int64_t
rhsComponentCount int64_t
numPrimitiveOperations int64_t
allowImpreciseAccumulation bool

DotDimensionNumbersAttr

Attribute that models the dimension information for dot.

Parameters:

Parameter C++ type Description
lhsBatchingDimensions ::llvm::ArrayRef<int64_t> Dimension
rhsBatchingDimensions ::llvm::ArrayRef<int64_t> Dimension
lhsContractingDimensions ::llvm::ArrayRef<int64_t> Dimension
rhsContractingDimensions ::llvm::ArrayRef<int64_t> Dimension

FftTypeAttr

XLA fast fourier transform type.

Syntax:

#mhlo.fft_type<
  ::mlir::mhlo::FftType   # value
>

Enum cases:

  • FFT (FFT)
  • IFFT (IFFT)
  • RFFT (RFFT)
  • IRFFT (IRFFT) #### Parameters:
Parameter C++ type Description
value ::mlir::mhlo::FftType an enum of type FftType

FusionKindAttr

fusion kind

Syntax:

#mhlo.fusion_kind<
  ::mlir::mhlo::FusionKind   # value
>

Enum cases:

  • kLoop (kLoop)
  • kInput (kInput)
  • kOutput (kOutput)
  • kCustom (kCustom) #### Parameters:
Parameter C++ type Description
value ::mlir::mhlo::FusionKind an enum of type FusionKind

GatherDimensionNumbersAttr

Attribute that models the dimension information for gather

Parameters:

Parameter C++ type Description
offsetDims ::llvm::ArrayRef<int64_t> Dimension
collapsedSliceDims ::llvm::ArrayRef<int64_t> Dimension
operandBatchingDims ::llvm::ArrayRef<int64_t> Dimension
startIndicesBatchingDims ::llvm::ArrayRef<int64_t> Dimension
startIndexMap ::llvm::ArrayRef<int64_t> Dimension
indexVectorDim int64_t

OutputOperandAliasAttr

Attribute that models the alias relationship of output and operand of a CustomCall op

Syntax:

#mhlo.output_operand_alias<
  ::llvm::ArrayRef<int64_t>,   # outputTupleIndices
  int64_t,   # operandIndex
  ::llvm::ArrayRef<int64_t>   # operandTupleIndices
>

This attribute captures the alias relationship of the output to one of the operands for a CustomCall op, denoted by operand_index. The output_tuple_indices and operand_tuple_indices are used to index into output and operand types. These indices lists are empty if the corresponding types are not tuple types, and can be arbitrarily long in case of arbitrarily nested tuple types.

See https://www.tensorflow.org/xla/aliasing

Example when used as array with in mhlo.custom-call:

%0 = "mhlo.custom_call"(%arg0, %arg1) {
  // other attributes
  output_operand_alias = [
    #mhlo.output_operand_alias<output_tuple_indices = [0],
                               operand_index = 0,
                               operand_tuple_indices = [1]>
  ]
} : (tuple<tensor<1x1xf32>, tensor<2x3xf32>>, tensor<5x5xf32>) -> tuple<tensor<2x3xf32>>

The output and the 0th operand are both tuples. The aliasing shows the
relationship between the 0th element in output tuple with the 1st element in
the 0th operand. And both of them are of the same type: tensor<2x3xf32>.

Parameters:

Parameter C++ type Description
outputTupleIndices ::llvm::ArrayRef<int64_t> Dimension
operandIndex int64_t
operandTupleIndices ::llvm::ArrayRef<int64_t> Dimension

PrecisionAttr

XLA precision for an operand. Has backend specific meaning.

Syntax:

#mhlo.precision<
  ::mlir::mhlo::Precision   # value
>

Enum cases:

  • DEFAULT (DEFAULT)
  • HIGH (HIGH)
  • HIGHEST (HIGHEST)
  • PACKED_NIBBLE (PACKED_NIBBLE) #### Parameters:
Parameter C++ type Description
value ::mlir::mhlo::Precision an enum of type Precision

RngAlgorithmAttr

XLA PRNG algorithm to be used.

Syntax:

#mhlo.rng_algorithm<
  ::mlir::mhlo::RngAlgorithm   # value
>

Enum cases:

  • DEFAULT (DEFAULT)
  • THREE_FRY (THREE_FRY)
  • PHILOX (PHILOX) #### Parameters:
Parameter C++ type Description
value ::mlir::mhlo::RngAlgorithm an enum of type RngAlgorithm

RngDistributionAttr

XLA PRNG distribution to be used.

Syntax:

#mhlo.rng_distribution<
  ::mlir::mhlo::RngDistribution   # value
>

Enum cases:

  • UNIFORM (UNIFORM)
  • NORMAL (NORMAL) #### Parameters:
Parameter C++ type Description
value ::mlir::mhlo::RngDistribution an enum of type RngDistribution

ScatterDimensionNumbersAttr

Attribute that models the dimension information for scatter

Parameters:

Parameter C++ type Description
updateWindowDims ::llvm::ArrayRef<int64_t> Dimension
insertedWindowDims ::llvm::ArrayRef<int64_t> Dimension
inputBatchingDims ::llvm::ArrayRef<int64_t> Dimension
scatterIndicesBatchingDims ::llvm::ArrayRef<int64_t> Dimension
scatterDimsToOperandDims ::llvm::ArrayRef<int64_t> Dimension
indexVectorDim int64_t

SparsityDescriptorAttr

Describes structured (N:M) sparsity configuration

Syntax:

#mhlo.sparsity<
  int64_t,   # dimension
  int64_t,   # n
  int64_t   # m
>

This attribute is defined for a sparse dot operation with a structured sparse input tensor. With (N=2,M=4), every 4 consecutive logical elements have exactly 2 non-zero physical elements in the input tensor.

$dimension defines the index of the contracting dimension that is sparse (it has to be the most minor dimension). The additional metadata operand in the sparse dot operation defines which logical elements are zeroed out.

Parameters:

Parameter C++ type Description
dimension int64_t
n int64_t
m int64_t

TransposeAttr

Transpose options

Syntax:

#mhlo.transpose<
  ::mlir::mhlo::Transpose   # value
>

Enum cases:

  • TRANSPOSE_INVALID (TRANSPOSE_INVALID)
  • NO_TRANSPOSE (NO_TRANSPOSE)
  • TRANSPOSE (TRANSPOSE)
  • ADJOINT (ADJOINT) #### Parameters:
Parameter C++ type Description
value ::mlir::mhlo::Transpose an enum of type Transpose

TypeExtensionsAttr

Attribute that extends tensor type with MHLO type properties.

Syntax:

#mhlo.type_extensions<
  ::llvm::ArrayRef<int64_t>   # bounds
>

This attribute is used to extend MLIR tensor type with MHLO tensor specific properties. These properties aren't modeled in the MLIR type. This attribute is set in the encoding field of the tensor type.

See HLO_BoundedAttrInterface for documentation for bounds.

Parameters:

Parameter C++ type Description
bounds ::llvm::ArrayRef<int64_t>

Types

AsyncBundleType

Opaque collection of other types

Syntax:

!mhlo.async_bundle<
  ::llvm::ArrayRef<Type>   # types
>

Parameters:

Parameter C++ type Description
types ::llvm::ArrayRef<Type>

Enums

ComparisonDirection

Which comparison operation to perform.

Cases:

Symbol Value String
EQ 0 EQ
NE 1 NE
GE 2 GE
GT 3 GT
LE 4 LE
LT 5 LT

ComparisonType

Which comparison type to use.

Cases:

Symbol Value String
NOTYPE 0 NOTYPE
FLOAT 1 FLOAT
TOTALORDER 2 TOTALORDER
SIGNED 3 SIGNED
UNSIGNED 4 UNSIGNED

CustomCallApiVersion

Custom call API version

Cases:

Symbol Value String
API_VERSION_UNSPECIFIED 0 API_VERSION_UNSPECIFIED
API_VERSION_ORIGINAL 1 API_VERSION_ORIGINAL
API_VERSION_STATUS_RETURNING 2 API_VERSION_STATUS_RETURNING
API_VERSION_STATUS_RETURNING_UNIFIED 3 API_VERSION_STATUS_RETURNING_UNIFIED
API_VERSION_TYPED_FFI 4 API_VERSION_TYPED_FFI

CustomCallSchedule

Specifies the desired schedule for the custom-call.

Cases:

Symbol Value String
NONE 0 NONE
LATEST 1 LATEST
EARLIEST 2 EARLIEST

DequantizeMode

Dequantization mode. Only MIN_COMBINED is supported.

Cases:

Symbol Value String
MIN_COMBINED 0 MIN_COMBINED

DomainKind

Kind of domain metatdata attached to an HLO domain.

Cases:

Symbol Value String
sharding 0 sharding

FftType

XLA fast fourier transform type.

Cases:

Symbol Value String
FFT 0 FFT
IFFT 1 IFFT
RFFT 2 RFFT
IRFFT 3 IRFFT

FusionKind

fusion kind

Cases:

Symbol Value String
kLoop 0 kLoop
kInput 1 kInput
kOutput 2 kOutput
kCustom 3 kCustom

Precision

XLA precision for an operand. Has backend specific meaning.

Cases:

Symbol Value String
DEFAULT 0 DEFAULT
HIGH 1 HIGH
HIGHEST 2 HIGHEST
PACKED_NIBBLE 3 PACKED_NIBBLE

RngAlgorithm

XLA PRNG algorithm to be used.

Cases:

Symbol Value String
DEFAULT 0 DEFAULT
THREE_FRY 1 THREE_FRY
PHILOX 2 PHILOX

RngDistribution

XLA PRNG distribution to be used.

Cases:

Symbol Value String
UNIFORM 1 UNIFORM
NORMAL 2 NORMAL

Transpose

Transpose options

Cases:

Symbol Value String
TRANSPOSE_INVALID 0 TRANSPOSE_INVALID
NO_TRANSPOSE 1 NO_TRANSPOSE
TRANSPOSE 2 TRANSPOSE
ADJOINT 3 ADJOINT