'mhlo' Dialect

Stay organized with collections Save and categorize content based on your preferences.

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:

Results:

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:

Results:

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:

Results:

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:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Results:

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:

Results:

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:

Operands:

Results:

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:

Results:

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:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Results:

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:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Results:

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{}

Attributes:

Operands:

Results:

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:

Results:

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:

Operands:

Results:

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:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Results:

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:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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{}

Attributes:

Operands:

Results:

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:

Results:

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:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Results:

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:

Results:

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:

Operands:

Results:

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:

Results:

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:

Operands:

Results:

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:

Results:

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{}

Attributes:

Operands:

Results:

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{}

Attributes:

Operands:

Results:

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:

Operands:

Results:

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:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Operands:

Results:

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:

Results:

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:

Results:

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:

Operands:

Results:

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:

Results:

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:

Results:

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>