Operation semantics

The following describes the semantics of operations defined in the XlaBuilder interface. Typically, these operations map one-to-one to operations defined in the RPC interface in xla_data.proto.

A note on nomenclature: the generalized data type XLA deals with is an N-dimensional array holding elements of some uniform type (such as 32-bit float). Throughout the documentation, array is used to denote an arbitrary-dimensional array. For convenience, special cases have more specific and familiar names; for example a vector is a 1-dimensional array and a matrix is a 2-dimensional array.

Abs

Add

AddDependency

AddDependency may appear in HLO dumps, but they are not intended to be constructed manually by end users.

Note: AddDependency is only found in HLO. It is not found in StableHLO.

AfterAll

Arguments	Type	Semantics
`tokens`	vector of `XlaOp`	variadic number of tokens

AllGather

Arguments	Type	Semantics
`operand`	`XlaOp`	Array to concatenate across
: : : replicas :
`all_gather_dimension`	`int64`	Concatenation dimension
`shard_count`	`int64`	The size of each replica
: : : group :
`replica_groups`	vector of vectors of	Groups between which the
: : `int64` : concatenation is performed :
`channel_id`	optional	Optional channel ID for
: : `ChannelHandle` : cross-module communication :
`layout`	optional `Layout`	Creates a layout pattern
: : : that will capture the :
: : : matched layout in the :
: : : argument :
`use_global_device_ids`	optional `bool`	Returns true if the ids in
: : : the ReplicaGroup config :
: : : represent a global id :

AllReduce

Arguments	Type	Semantics
`operand`	`XlaOp`	Array or a non-empty tuple
: : : of arrays to reduce across :
: : : replicas :
`computation`	`XlaComputation`	Reduction computation
`replica_groups`	`ReplicaGroup` vector	Groups between which the
: : : reductions are performed :
`channel_id`	optional	Optional channel ID for
: : `ChannelHandle` : cross-module communication :
`shape_with_layout`	optional `Shape`	Defines the layout of the
: : : data transferred :
`use_global_device_ids`	optional `bool`	Returns true if the ids in
: : : the ReplicaGroup config :
: : : represent a global id :

CrossReplicaSum

Arguments	Type	Semantics
`operand`	XlaOp	Array or a non-empty tuple of
: : : arrays to reduce across :
: : : replicas :
`replica_groups`	vector of vectors of	Groups between which the
: : `int64` : reductions are performed :

AllToAll

AllToAll - Example 1.

XlaBuilder b("alltoall");
auto x = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {4, 16}), "x");
AllToAll(
    x,
    /*split_dimension=*/ 1,
    /*concat_dimension=*/ 0,
    /*split_count=*/ 4);

In the above example, there are 4 cores participating in the Alltoall. On each core, the operand is split into 4 parts along dimension 1, so each part has shape f32[4,4]. The 4 parts are scattered to all cores. Then each core concatenates the received parts along dimension 0, in the order of core 0-4. So the output on each core has shape f32[16,4].

AllToAll - Example 2 - StableHLO

An example of AllToAll dataflow for StableHLO

In the above example, there are 2 replicas participating in the AllToAll. On each replica, the operand has shape f32[2,4]. The operand is split into 2 parts along dimension 1, so each part has shape f32[2,2]. The 2 parts are then exchanged across the replicas according to their position in the replica group. Each replica collects its corresponding part from both operands and concatenates them along dimension 0. As a result, the output on each replica has shape f32[4,2].

RaggedAllToAll

Arguments	Type	Semantics
`input`	`XlaOp`	N array of type T
`input_offsets`	`XlaOp`	N array of type T
`send_sizes`	`XlaOp`	N array of type T
`output`	`XlaOp`	N array of type T
`output_offsets`	`XlaOp`	N array of type T
`recv_sizes`	`XlaOp`	N array of type T
`replica_groups`	`ReplicaGroup` vector	Each group contains a list of
: : : replica ids. :
`channel_id`	optional `ChannelHandle`	unique identifier for each
: : : send/recv pair :

And

Async

AsyncDone, AsyncStart, and AsyncUpdate are internal HLO instructions used for Asynchronous operations and serve as primitives in HLO. These ops may appear in HLO dumps but they are not intended to be constructed manually by end users.

Note: AsyncStart, AsyncUpdate, and AsyncDone are only found in HLO. They are not found in StableHLO.

Atan2

BatchNormGrad

See also XlaBuilder::BatchNormGrad and the original batch normalization paper for a detailed description of the algorithm.

Calculates gradients of batch norm.

BatchNormGrad(operand, scale, batch_mean, batch_var, grad_output, epsilon, feature_index)

Arguments	Type	Semantics
`operand`	XlaOp	n dimensional array to be normalized (x)
`scale`	XlaOp	1 dimensional array ($\gamma$)
`batch_mean`	XlaOp	1 dimensional array ($\mu$)
`batch_var`	XlaOp	1 dimensional array ($\sigma^2$)
`grad_output`	XlaOp	Gradients passed to `BatchNormTraining` ($\nabla y$)
`epsilon`	`float`	Epsilon value ($\epsilon$)
`feature_index`	`int64`	Index to feature dimension in `operand`

For each feature in the feature dimension (feature_index is the index for the feature dimension in operand), the operation calculates the gradients with respect to operand, offset, and scale across all the other dimensions. The feature_index must be a valid index for the feature dimension in operand.

The three gradients are defined by the following formulas (assuming a 4-dimensional array as operand and with feature dimension index l, batch size m and spatial sizes w and h):

$$ \begin{split} c_l&= \frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \left( \nabla y_{ijkl} \frac{x_{ijkl} - \mu_l}{\sigma^2_l+\epsilon} \right) \\ d_l&= \frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \nabla y_{ijkl} \\ \nabla x_{ijkl} &= \frac{\gamma_{l}}{\sqrt{\sigma^2_{l}+\epsilon}} \left( \nabla y_{ijkl} - d_l - c_l (x_{ijkl} - \mu_{l}) \right) \\ \nabla \gamma_l &= \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \left( \nabla y_{ijkl} \frac{x_{ijkl} - \mu_l}{\sqrt{\sigma^2_{l}+\epsilon}} \right) \\
\nabla \beta_l &= \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \nabla y_{ijkl} \end{split} $$

The inputs batch_mean and batch_var represent moments values across batch and spatial dimensions.

The output type is a tuple of three handles:

Outputs	Type	Semantics
`grad_operand`	XlaOp	gradient with respect to
: : : input `operand` :
: : : ($\nabla x$) :
`grad_scale`	XlaOp	gradient with respect to
: : : input `scale` :
: : : ($\nabla\gamma$) :
`grad_offset`	XlaOp	gradient with respect to
: : : input :
: : : `offset`($\nabla\beta$) :

For StableHLO information see StableHLO - batch_norm_grad.

BatchNormInference

See also XlaBuilder::BatchNormInference and the original batch normalization paper for a detailed description of the algorithm.

Normalizes an array across batch and spatial dimensions.

BatchNormInference(operand, scale, offset, mean, variance, epsilon, feature_index)

Arguments	Type	Semantics
`operand`	XlaOp	n dimensional array to be normalized
`scale`	XlaOp	1 dimensional array
`offset`	XlaOp	1 dimensional array
`mean`	XlaOp	1 dimensional array
`variance`	XlaOp	1 dimensional array
`epsilon`	`float`	Epsilon value
`feature_index`	`int64`	Index to feature dimension in `operand`

For each feature in the feature dimension (feature_index is the index for the feature dimension in operand), the operation calculates the mean and variance across all the other dimensions and uses the mean and variance to normalize each element in operand. The feature_index must be a valid index for the feature dimension in operand.

BatchNormInference is equivalent to calling BatchNormTraining without computing mean and variance for each batch. It uses the input mean and variance instead as estimated values. The purpose of this op is to reduce latency in inference, hence the name BatchNormInference.

The output is an n-dimensional, normalized array with the same shape as input operand.

For StableHLO information see StableHLO - batch_norm_inference.

BatchNormTraining

See also XlaBuilder::BatchNormTraining and the original batch normalization paper for a detailed description of the algorithm.

Normalizes an array across batch and spatial dimensions.

BatchNormTraining(operand, scale, offset, epsilon, feature_index)

Arguments	Type	Semantics
`operand`	`XlaOp`	n dimensional array to be normalized (x)
`scale`	`XlaOp`	1 dimensional array ($\gamma$)
`offset`	`XlaOp`	1 dimensional array ($\beta$)
`epsilon`	`float`	Epsilon value ($\epsilon$)
`feature_index`	`int64`	Index to feature dimension in `operand`

The algorithm goes as follows for each batch in operand $x$ that contains m elements with w and h as the size of spatial dimensions (assuming operand is a 4 dimensional array):

Calculates batch mean $\mu_l$ for each feature l in feature dimension: $\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}$
Calculates batch variance $\sigma^2_l$: $\sigma^2_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h (x_{ijkl} - \mu_l)^2$
Normalizes, scales and shifts: $y_{ijkl}=\frac{\gamma_l(x_{ijkl}-\mu_l)}{\sqrt[2]{\sigma^2_l+\epsilon}}+\beta_l$

The epsilon value, usually a small number, is added to avoid divide-by-zero errors.

The output type is a tuple of three XlaOps:

Outputs	Type	Semantics
`output`	`XlaOp`	n dimensional array with the same shape as input
: : : `operand` (y) :
`batch_mean`	`XlaOp`	1 dimensional array ($\mu$)
`batch_var`	`XlaOp`	1 dimensional array ($\sigma^2$)

The batch_mean and batch_var are moments calculated across the batch and spatial dimensions using the formulas above.

For StableHLO information see StableHLO - batch_norm_training.

Bitcast

Bitcast may appear in HLO dumps, but they are not intended to be constructed manually by end users.

Note: Bitcast is only found in HLO. It is not found in StableHLO.

BitcastConvertType

Similar to a tf.bitcast in TensorFlow, performs an element-wise bitcast operation from a data shape to a target shape. The input and output size must match: e.g. s32 elements become f32 elements via bitcast routine, and one s32 element will become four s8 elements. Bitcast is implemented as a low-level cast, so machines with different floating-point representations will give different results.

BitcastConvertType(operand, new_element_type)

Arguments	Type	Semantics
`operand`	`XlaOp`	array of type T with dims D
`new_element_type`	`PrimitiveType`	type U

The dimensions of the operand and the target shape must match, apart from the last dimension which will change by the ratio of the primitive size before and after the conversion.

The source and destination element types must not be tuples.

For StableHLO information see StableHLO - bitcast_convert.

Bitcast-converting to primitive type of different width

BitcastConvert HLO instruction supports the case where the size of the output element type T' is not equal to the size of the input element T. As the whole operation is conceptually a bitcast and does not change the underlying bytes, the shape of the output element has to change. For B = sizeof(T), B' = sizeof(T'), there are two possible cases.

First, when B > B', the output shape gets a new minor-most dimension of size B/B'. For example:

  f16[10,2]{1,0} %output = f16[10,2]{1,0} bitcast-convert(f32[10]{0} %input)

The rule remains the same for effective scalars:

  f16[2]{0} %output = f16[2]{0} bitcast-convert(f32[] %input)

Alternatively, for B' > B the instruction requires the last logical dimension of the input shape to be equal to B'/B, and this dimension is dropped during the conversion:

  f32[10]{0} %output = f32[10]{0} bitcast-convert(f16[10,2]{1,0} %input)

Note that conversions between different bitwidths are not elementwise.

Broadcast

Arguments	Type	Semantics
`operand`	`XlaOp`	The array to duplicate
`broadcast_sizes`	`ArraySlice<int64>`	The sizes of the new dimensions

BroadcastInDim

Arguments	Type	Semantics
`operand`	`XlaOp`	The array to duplicate
`out_dim_size`	`ArraySlice<int64>`	The sizes of the dimensions
: : : of the target shape :
`broadcast_dimensions`	`ArraySlice<int64>`	Which dimension in the target
: : : shape each dimension of the :
: : : operand shape corresponds to :

Call

Arguments	Type	Semantics
`computation`	`XlaComputation`	computation of type `T_0, T_1, ...,
: : : T_{N-1} -> S` with N parameters of :
: : : arbitrary type :
`operands`	sequence of N `XlaOp`s	N arguments of arbitrary type

CompositeCall

Arguments	Type	Semantics
`computation`	`XlaComputation`	computation of type `T_0, T_1, ...,
: : : T_{N-1} -> S` with N parameters of :
: : : arbitrary type :
`operands`	sequence of N `XlaOp`s	variadic number of values
`name`	`string`	name of the composite
`attributes`	optional `string`	optional stringified dictionary of
: : : attributes :
`version`	optional `int64`	number to version updates to
: : : semantics of the composite op :

Cbrt

Ceil

Cholesky

Arguments	Type	Semantics
`a`	`XlaOp`	an array of a complex or floating-point type with > 2
: : : dimensions. :
`lower`	`bool`	whether to use the upper or lower triangle of `a`.

Clamp

Arguments	Type	Semantics
`min`	`XlaOp`	array of type T
`operand`	`XlaOp`	array of type T
`max`	`XlaOp`	array of type T

Collapse

See also XlaBuilder::Collapse. and the tf.reshape operation.

Collapses dimensions of an array into one dimension.

Collapse(operand, dimensions)

Arguments	Type	Semantics
`operand`	`XlaOp`	array of type T
`dimensions`	`int64` vector	in-order, consecutive subset of T's dimensions.

Collapse replaces the given subset of the operand‘s dimensions by a single dimension. The input arguments are an arbitrary array of type T and a compile-time-constant vector of dimension indices. The dimension indices must be an in-order (low to high dimension numbers), consecutive subset of T’s dimensions. Thus, {0, 1, 2}, {0, 1}, or {1, 2} are all valid dimension sets, but {1, 0} or {0, 2} are not. They are replaced by a single new dimension, in the same position in the dimension sequence as those they replace, with the new dimension size equal to the product of original dimension sizes. The lowest dimension number in dimensions is the slowest varying dimension (most major) in the loop nest which collapses these dimensions, and the highest dimension number is fastest varying (most minor). See the tf.reshape operator if more general collapse ordering is needed.

For example, let v be an array of 24 elements:

let v = f32[4x2x3] {{{10, 11, 12}, {15, 16, 17}},
{{20, 21, 22}, {25, 26, 27}},
{{30, 31, 32}, {35, 36, 37}},
{{40, 41, 42}, {45, 46, 47}}};

// Collapse to a single dimension, leaving one dimension.
let v012 = Collapse(v, {0,1,2});
then v012 == f32[24] {10, 11, 12, 15, 16, 17,
20, 21, 22, 25, 26, 27,
30, 31, 32, 35, 36, 37,
40, 41, 42, 45, 46, 47};

// Collapse the two lower dimensions, leaving two dimensions.
let v01 = Collapse(v, {0,1});
then v01 == f32[4x6] {{10, 11, 12, 15, 16, 17},
{20, 21, 22, 25, 26, 27},
{30, 31, 32, 35, 36, 37},
{40, 41, 42, 45, 46, 47}};

// Collapse the two higher dimensions, leaving two dimensions.
let v12 = Collapse(v, {1,2});
then v12 == f32[8x3] {{10, 11, 12},
{15, 16, 17},
{20, 21, 22},
{25, 26, 27},
{30, 31, 32},
{35, 36, 37},
{40, 41, 42},
{45, 46, 47}};

Note: For further information on Collapse see HLO - Reshape

Clz

CollectiveBroadcast

Broadcasts data across replicas. Data is sent from the first replica id in each group to the other ids in the same group. If a replica id is not in any replica group, the output on that replica is a tensor consisting of 0(s) in shape.

CollectiveBroadcast(operand, replica_groups, channel_id)

Arguments	Type	Semantics
`operand`	`XlaOp`	The operand to the function
`replica_groups`	`ReplicaGroup`vector	Each group contains a list of
: : : replica ids :
`channel_id`	optional `ChannelHandle`	unique identifier for each
: : : send/recv pair :

For StableHLO information see StableHLO - collective_broadcast.

CollectivePermute

CollectivePermute is a collective operation that sends and receives data across replicas.

CollectivePermute(operand, source_target_pairs, channel_id, inplace)

Arguments	Type	Semantics
`operand`	`XlaOp`	n dimensional input array
`source_target_pairs`	`<int64, int64>` vector	A list of
: : : (source_replica_id, :
: : : target_replica_id) pairs. :
: : : For each pair, the :
: : : operand is sent from :
: : : source replica to target :
: : : replica. :
`channel_id`	optional `ChannelHandle`	Optional channel ID for
: : : cross-module :
: : : communication :
`inplace`	optional `bool`	flag whether permutation
: : : should be done in place :

Note that there are the following restrictions on the source_target_pairs:

Any two pairs should not have the same target replica id, and they should not have the same source replica id.
If a replica id is not a target in any pair, then the output on that replica is a tensor consisting of 0(s) with the same shape as the input.

The API of CollectivePermute operation is internally decomposed into 2 HLO instructions (CollectivePermuteStart and CollectivePermuteDone).

CollectivePermuteStart and CollectivePermuteDone serve as primitives in HLO. These ops may appear in HLO dumps, but they are not intended to be constructed manually by end users.

For StableHLO information see StableHLO - collective_permute.

Note: CollectivePermuteStart and CollectivePermuteDone are only found in HLO. They are not found in StableHLO.

Compare

Eq

Ne

Ge

Gt

Le

Lt

Complex

ConcatInDim (Concatenate)

Conditional

Constant

ConvertElementType

Similar to an element-wise static_cast in C++, ConvertElementType performs an element-wise conversion operation from a data shape to a target shape. The dimensions must match, and the conversion is an element-wise one; e.g. s32 elements become f32 elements via an s32-to-f32 conversion routine.

ConvertElementType(operand, new_element_type)

Arguments	Type	Semantics
`operand`	`XlaOp`	array of type T with dims D
`new_element_type`	`PrimitiveType`	type U

The dimensions of the operand and the target shape must match. The source and destination element types must not be tuples.

A conversion such as T=s32 to U=f32 will perform a normalizing int-to-float conversion routine such as round-to-nearest-even.

Note: The precise float-to-int and visa-versa conversions are currently unspecified, but may become additional arguments to the convert operation in the future. Not all possible conversions have been implemented for all targets.

let a: s32[3] = {0, 1, 2};
let b: f32[3] = convert(a, f32);
then b == f32[3]{0.0, 1.0, 2.0}

For StableHLO information see StableHLO - convert.

Conv (Convolution)

ConvWithGeneralPadding

ConvWithGeneralPadding(lhs, rhs, window_strides, padding, feature_group_count, batch_group_count, precision_config, preferred_element_type)

Same as Conv where padding configuration is explicit.

Arguments	Type	Semantics
`lhs`	`XlaOp`	(n+2)-dimensional array of
: : : inputs :
`rhs`	`XlaOp`	(n+2)-dimensional array of
: : : kernel weights :
`window_strides`	`ArraySlice<int64>`	n-d array of kernel strides
`padding`	`ArraySlice<	n-d array of (low, high)
: : pair<int64,int64>>` : padding :
`feature_group_count`	int64	the number of feature
: : : groups :
`batch_group_count`	int64	the number of batch groups
`precision_config`	optional	enum for level of precision
: : `PrecisionConfig` : :
`preferred_element_type`	optional	enum of scalar element type
: : `PrimitiveType` : :

ConvWithGeneralDimensions

ConvWithGeneralDimensions(lhs, rhs, window_strides, padding, dimension_numbers, feature_group_count, batch_group_count, precision_config, preferred_element_type)

Same as Conv where dimension numbers are explicit.

Arguments	Type	Semantics
`lhs`	`XlaOp`	(n+2)-dimensional
: : : array of inputs :
`rhs`	`XlaOp`	(n+2)-dimensional
: : : array of kernel :
: : : weights :
`window_strides`	`ArraySlice<int64>`	n-d array of
: : : kernel strides :
`padding`	`Padding`	enum of padding
`dimension_numbers`	`ConvolutionDimensionNumbers`	the number of
: : : dimensions :
`feature_group_count`	int64	the number of
: : : feature groups :
`batch_group_count`	int64	the number of
: : : batch groups :
`precision_config`	optional `PrecisionConfig`	enum for level of
: : : precision :
`preferred_element_type`	optional `PrimitiveType`	enum of scalar
: : : element type :

ConvGeneral

ConvGeneralDilated

ConvGeneralDilated(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count, batch_group_count, precision_config, preferred_element_type, window_reversal)

Same as Conv where padding configuration, dilation factors, and dimension numbers are explicit.

Arguments	Type	Semantics
`lhs`	`XlaOp`	(n+2)-dimensional
: : : array of inputs :
`rhs`	`XlaOp`	(n+2)-dimensional
: : : array of kernel :
: : : weights :
`window_strides`	`ArraySlice<int64>`	n-d array of
: : : kernel strides :
`padding`	`ArraySlice<	n-d array of
: : pair<int64,int64>>` : (low, high) :
: : : padding :
`lhs_dilation`	`ArraySlice<int64>`	n-d lhs dilation
: : : factor array :
`rhs_dilation`	`ArraySlice<int64>`	n-d rhs dilation
: : : factor array :
`dimension_numbers`	`ConvolutionDimensionNumbers`	the number of
: : : dimensions :
`feature_group_count`	int64	the number of
: : : feature groups :
`batch_group_count`	int64	the number of
: : : batch groups :
`precision_config`	optional `PrecisionConfig`	enum for level of
: : : precision :
`preferred_element_type`	optional `PrimitiveType`	enum of scalar
: : : element type :
`window_reversal`	optional `vector<bool>`	flag used to
: : : logically reverse :
: : : dimension before :
: : : applying the :
: : : convolution :

Copy

Copy is internally decomposed into 2 HLO instructions CopyStart and CopyDone. Copy along with CopyStart and CopyDone serve as primitives in HLO. These ops may appear in HLO dumps, but they are not intended to be constructed manually by end users.

Note: Copy, CopyStart, CopyDone are only found in HLO. They are not found in StableHLO.

Cos

Cosh

CustomCall

Div

Domain

Domain may appear in HLO dumps, but it is not intended to be constructed manually by end users.

Note: Domain is only found in HLO. It is not found in StableHLO.

Dot

DotGeneral

ScaledDot

RaggedDot

DynamicReshape

DynamicSlice

DynamicUpdateSlice

DynamicUpdateSlice generates a result which is the value of the input array operand, with a slice update overwritten at start_indices. The shape of update determines the shape of the sub-array of the result which is updated. The shape of start_indices must be 1-dimensional, with dimension size equal to the number of dimensions of operand.

DynamicUpdateSlice(operand, update, start_indices)

Arguments	Type	Semantics
`operand`	`XlaOp`	N dimensional array of type T
`update`	`XlaOp`	N dimensional array of type T
: : : containing the slice update. Each :
: : : dimension of update shape must be :
: : : strictly greater than zero, and :
: : : start + update must be less than :
: : : or equal to the operand size for :
: : : each dimension to avoid generating :
: : : out-of-bounds update indices. :
`start_indices`	sequence of N `XlaOp`	List of N scalar integers
: : : containing the starting indices of :
: : : the slice for each dimension. :
: : : Value must be greater than or :
: : : equal to zero. :

The effective slice indices are computed by applying the following transformation for each index i in [1, N) before performing the slice:

start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - update.dimension_size[i])

This ensures that the updated slice is always in-bounds with respect to the operand array. If the slice is in-bounds before the transformation is applied, the transformation has no effect.

1-dimensional example:

let a = {0.0, 1.0, 2.0, 3.0, 4.0}
let u = {5.0, 6.0}
let s = {2}

DynamicUpdateSlice(a, u, s)
// Result: {0.0, 1.0, 5.0, 6.0, 4.0}

2-dimensional example:

let b =
{ {0.0,  1.0,  2.0},
  {3.0,  4.0,  5.0},
  {6.0,  7.0,  8.0},
  {9.0, 10.0, 11.0} }
let u =
{ {12.0, 13.0},
  {14.0, 15.0},
  {16.0, 17.0} }

let s = {1, 1}

DynamicUpdateSlice(b, u, s)
// Result:
// { {0.0,  1.0,  2.0},
//   {3.0, 12.0, 13.0},
//   {6.0, 14.0, 15.0},
//   {9.0, 16.0, 17.0} }

For StableHLO information see StableHLO - dynamic_update_slice.

Erf

Exp

Expm1

Fft

Multidimensional FFT

When more than 1 fft_length is provided, this is equivalent to applying a cascade of FFT operations to each of the innermost axes. Note that for the real->complex and complex->real cases, the innermost axis transform is (effectively) performed first (RFFT; last for IRFFT), which is why the innermost axis is the one which changes size. Other axis transforms will then be complex->complex.

Implementation details

CPU FFT is backed by Eigen's TensorFFT. GPU FFT uses cuFFT.

Floor

Fusion

Fusion operation represents HLO instructions and serves as a primitive in HLO. This op may appear in HLO dumps but is not intended to be constructed manually by end users.

Note: Fusion is only found in HLO. It is not found in StableHLO.

Gather

The XLA gather operation stitches together several slices (each slice at a potentially different runtime offset) of an input array.

For StableHLO information see StableHLO - gather.

General Semantics

See also XlaBuilder::Gather. For a more intuitive description, see the “Informal Description” section below.

gather(operand, start_indices, dimension_numbers, slice_sizes, indices_are_sorted)

Arguments	Type	Semantics
`operand`	`XlaOp`	The array we’re gathering
: : : from. :
`start_indices`	`XlaOp`	Array containing the
: : : starting indices of the :
: : : slices we gather. :
`dimension_numbers`	`GatherDimensionNumbers`	The dimension in
: : : `start_indices` that :
: : : “contains” the starting :
: : : indices. See below for a :
: : : detailed description. :
`slice_sizes`	`ArraySlice<int64>`	`slice_sizes[i]` is the
: : : bounds for the slice on :
: : : dimension `i`. :
`indices_are_sorted`	`bool`	Whether the indices are
: : : guaranteed to be sorted by :
: : : the caller. :

For convenience, we label dimensions in the output array not in offset_dims as batch_dims.

The output is an array with batch_dims.size + offset_dims.size dimensions.

The operand.rank must equal the sum of offset_dims.size and collapsed_slice_dims.size. Also, slice_sizes.size has to be equal to operand.rank.

If index_vector_dim is equal to start_indices.rank we implicitly consider start_indices to have a trailing 1 dimension (i.e. if start_indices was of shape [6,7] and index_vector_dim is 2 then we implicitly consider the shape of start_indices to be [6,7,1]).

The bounds for the output array along dimension i is computed as follows:

If i is present in batch_dims (i.e. is equal to batch_dims[k] for some k) then we pick the corresponding dimension bounds out of start_indices.shape, skipping index_vector_dim (i.e. pick start_indices.shape.dims[k] if k < index_vector_dim and start_indices.shape.dims[k+1] otherwise).
If i is present in offset_dims (i.e. equal to offset_dims[k] for some k) then we pick the corresponding bound out of slice_sizes after accounting for collapsed_slice_dims (i.e. we pick adjusted_slice_sizes[k] where adjusted_slice_sizes is slice_sizes with the bounds at indices collapsed_slice_dims removed).

Formally, the operand index In corresponding to a given output index Out is calculated as follows:

Let G = { Out[k] for k in batch_dims }. Use G to slice out a vector S such that S[i] = start_indices[Combine(G, i)] where Combine(A, b) inserts b at position index_vector_dim into A. Note that this is well defined even if G is empty: If G is empty then S = start_indices.
Create a starting index, Sin, into operand using S by scattering S using start_index_map. More precisely:
1. Sin[start_index_map[k]] = S[k] if k < start_index_map.size.
2. Sin[_] = 0 otherwise.
Create an index Oin into operand by scattering the indices at the offset dimensions in Out according to the collapsed_slice_dims set. More precisely:
1. Oin[remapped_offset_dims(k)] = Out[offset_dims[k]] if k < offset_dims.size (remapped_offset_dims is defined below).
2. Oin[_] = 0 otherwise.
In is Oin + Sin where + is element-wise addition.

remapped_offset_dims is a monotonic function with domain [0, offset_dims.size) and range [0, operand.rank) \ collapsed_slice_dims. So if, e.g., offset_dims.size is 4, operand.rank is 6 and collapsed_slice_dims is {0, 2} then remapped_offset_dims is {0→1, 1→3, 2→4, 3→5}.

If indices_are_sorted is set to true then XLA can assume that start_indices are sorted (in ascending order, after scattering its values according to start_index_map) by the user. If they are not then the semantics are implementation defined.

Informal Description and Examples

Informally, every index Out in the output array corresponds to an element E in the operand array, computed as follows:

We use the batch dimensions in Out to look up a starting index from start_indices.
We use start_index_map to map the starting index (whose size may be less than operand.rank) to a “full” starting index into the operand.
We dynamic-slice out a slice with size slice_sizes using the full starting index.
We reshape the slice by collapsing the collapsed_slice_dims dimensions. Since all collapsed slice dimensions must have a bound of 1, this reshape is always legal.
We use the offset dimensions in Out to index into this slice to get the input element, E, corresponding to output index Out.

index_vector_dim is set to start_indices.rank - 1 in all of the examples that follow. More interesting values for index_vector_dim do not change the operation fundamentally, but make the visual representation more cumbersome.

To get an intuition on how all of the above fits together, let's look at an example that gathers 5 slices of shape [8,6] from a [16,11] array. The position of a slice into the [16,11] array can be represented as an index vector of shape S64[2], so the set of 5 positions can be represented as a S64[5,2] array.

The behavior of the gather operation can then be depicted as an index transformation that takes [G,O0,O1], an index in the output shape, and maps it to an element in the input array in the following way:

We first select an (X,Y) vector from the gather indices array using G. The element in the output array at index [G,O0,O1] is then the element in the input array at index [X+O0,Y+O1].

slice_sizes is [8,6], which decides the range of O0 and O1, and this in turn decides the bounds of the slice.

This gather operation acts as a batch dynamic slice with G as the batch dimension.

The gather indices may be multidimensional. For instance, a more general version of the example above using a “gather indices” array of shape [4,5,2] would translate indices like this:

Again, this acts as a batch dynamic slice G0 and G1 as the batch dimensions. The slice size is still [8,6].

The gather operation in XLA generalizes the informal semantics outlined above in the following ways:

We can configure which dimensions in the output shape are the offset dimensions (dimensions containing O0, O1 in the last example). The output batch dimensions (dimensions containing G0, G1 in the last example) are defined to be the output dimensions that are not offset dimensions.
The number of output offset dimensions explicitly present in the output shape may be smaller than the input number of dimensions. These “missing” dimensions, which are listed explicitly as collapsed_slice_dims, must have a slice size of 1. Since they have a slice size of 1 the only valid index for them is 0 and eliding them does not introduce ambiguity.
The slice extracted from the “Gather Indices” array ((X, Y) in the last example) may have fewer elements than the input array's number of dimensions, and an explicit mapping dictates how the index should be expanded to have the same number of dimensions as the input.

As a final example, we use (2) and (3) to implement tf.gather_nd:

G0 and G1 are used to slice out a starting index from the gather indices array as usual, except the starting index has only one element, X. Similarly, there is only one output offset index with the value O0. However, before being used as indices into the input array, these are expanded in accordance to “Gather Index Mapping” (start_index_map in the formal description) and “Offset Mapping” (remapped_offset_dims in the formal description) into [X,0] and [0,O0] respectively, adding up to [X,O0]. In other words, the output index [G0,G1,O0] maps to the input index [GatherIndices[G0,G1,0],O0] which gives us the semantics for tf.gather_nd.

slice_sizes for this case is [1,11]. Intuitively this means that every index X in the gather indices array picks an entire row and the result is the concatenation of all these rows.

GetDimensionSize

Returns the size of the given dimension of the operand. The operand must be array shaped.

GetDimensionSize(operand, dimension)

Arguments	Type	Semantics
`operand`	`XlaOp`	n dimensional input array
`dimension`	`int64`	A value in the interval `[0, n)` that specifies the
: : : dimension :

For StableHLO information see StableHLO - get_dimension_size.

GetTupleElement

Imag

Infeed

Iota

IsFinite

Log

Log1p

Logistic

Map

Max

Min

Mul

Neg

Not

OptimizationBarrier

Blocks any optimization pass from moving computations across the barrier.

OptimizationBarrier(operand)

Arguments	Type	Semantics
`operand`	`XlaOp`	The operand to the function

Ensures that all inputs are evaluated before any operators that depend on the barrier's outputs.

For StableHLO information see StableHLO - optimization_barrier.

Or

Outfeed

Pad

Parameter

PartitionID

Produces partition_id of the current process.

PartitionID(shape)

Arguments	Type	Semantics
`shape`	`Shape`	Shape of the data

PartitionID may appear in HLO dumps but it is not intended to be constructed manually by end users.

For StableHLO information see StableHLO - partition_id.

PopulationCount

Pow

Real

Recv

RecvDone

Similar to Send, the client API of Recv operation represents synchronous communication. However, the instruction is internally decomposed into 2 HLO instructions (Recv and RecvDone) to enable asynchronous data transfers.

Recv(const Shape& shape, int64 channel_id)

Allocates resources required to receive data from a Send instruction with the same channel_id. Returns a context for the allocated resources, which is used by a following RecvDone instruction to wait for the completion of the data transfer. The context is a tuple of {receive buffer (shape), request identifier (U32)} and it can only be used by a RecvDone instruction.

Given a context created by a Recv instruction, waits for the data transfer to complete and return the received data.

Note: RecvDone is only found in HLO. It is not found in StableHLO.

Reduce

Examples

When reducing across one dimension in a single 1D array with values [10, 11, 12, 13], with reduction function f (this is computation) then that could be computed as

f(10, f(11, f(12, f(init_value, 13)))

but there are also many other possibilities, e.g.

f(init_value, f(f(10, f(init_value, 11)), f(f(init_value, 12), f(init_value, 13))))

The following is a rough pseudo-code example of how reduction could be implemented, using summation as the reduction computation with an initial value of 0.

result_shape <- remove all dims in dimensions from operand_shape

# Iterate over all elements in result_shape. The number of r's here is equal
# to the number of dimensions of the result.
for r0 in range(result_shape[0]), r1 in range(result_shape[1]), ...:
  # Initialize this result element
  result[r0, r1...] <- 0

  # Iterate over all the reduction dimensions
  for d0 in range(dimensions[0]), d1 in range(dimensions[1]), ...:
    # Increment the result element with the value of the operand's element.
    # The index of the operand's element is constructed from all ri's and di's
    # in the right order (by construction ri's and di's together index over the
    # whole operand shape).
    result[r0, r1...] += operand[ri... di]

Here's an example of reducing a 2D array (matrix). The shape has 2 dimensions, dimension 0 of size 2 and dimension 1 of size 3:

Results of reducing dimensions 0 or 1 with an “add” function:

Note that both reduction results are 1D arrays. The diagram shows one as column and another as row just for visual convenience.

For a more complex example, here is a 3D array. Its number of dimensions is 3, dimension 0 of size 4, dimension 1 of size 2 and dimension 2 of size 3. For simplicity, the values 1 to 6 are replicated across dimension 0.

Similarly to the 2D example, we can reduce just one dimension. If we reduce dimension 0, for example, we get a 2-dimensional array where all values across dimension 0 were folded into a scalar:

|  4   8  12 |
| 16  20  24 |

If we reduce dimension 2, we also get a 2-dimensional array where all values across dimension 2 were folded into a scalar:

| 6  15 |
| 6  15 |
| 6  15 |
| 6  15 |

Note that the relative order between the remaining dimensions in the input is preserved in the output, but some dimensions may get assigned new numbers (since the number of dimensions changes).

We can also reduce multiple dimensions. Add-reducing dimensions 0 and 1 produces the 1D array [20, 28, 36].

Reducing the 3D array over all its dimensions produces the scalar 84.

Variadic Reduce

When N > 1, reduce function application is slightly more complex, as it is applied simultaneously to all inputs. The operands are supplied to the computation in the following order:

Running reduced value for the first operand
...
Running reduced value for the N'th operand
Input value for the first operand
...
Input value for the N'th operand

For example, consider the following reduction function, which can be used to compute the max and the argmax of a 1-D array in parallel:

f: (Float, Int, Float, Int) -> Float, Int
f(max, argmax, value, index):
  if value >= max:
    return (value, index)
  else:
    return (max, argmax)

For 1-D Input arrays V = Float[N], K = Int[N], and init values I_V = Float, I_K = Int, the result f_(N-1) of reducing across the only input dimension is equivalent to the following recursive application:

f_0 = f(I_V, I_K, V_0, K_0)
f_1 = f(f_0.first, f_0.second, V_1, K_1)
...
f_(N-1) = f(f_(N-2).first, f_(N-2).second, V_(N-1), K_(N-1))

Applying this reduction to an array of values, and an array of sequential indices (i.e. iota), will co-iterate over the arrays, and return a tuple containing the maximal value and the matching index.

ReducePrecision

ReduceScatter

ReduceScatter - Example 1 - StableHLO

An example of ReduceScatter dataflow for StableHLO

In the above example, there are 2 replicas participating in the ReduceScatter. On each replica, the operand has shape f32[2,4]. An all-reduce (sum) is performed across the replicas, producing a reduced value of shape f32[2,4] on each replica. This reduced value is then split into 2 parts along dimension 1, so each part has shape f32[2,2]. Each replica within the process group receives the part corresponding to its position in the group. As a result, the output on each replica has shape f32[2,2].

ReduceWindow

ReduceWindow - Example 1

Input is a matrix of size [4x6] and both window_dimensions and window_stride_dimensions are [2x3].

// Create a computation for the reduction (maximum).
XlaComputation max;
{
  XlaBuilder builder(client_, "max");
  auto y = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "y");
  auto x = builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "x");
  builder.Max(y, x);
  max = builder.Build().value();
}

// Create a ReduceWindow computation with the max reduction computation.
XlaBuilder builder(client_, "reduce_window_2x3");
auto shape = ShapeUtil::MakeShape(F32, {4, 6});
auto input = builder.Parameter(0, shape, "input");
builder.ReduceWindow(
    input,
    /*init_val=*/builder.ConstantLiteral(LiteralUtil::MinValue(F32)),
    *max,
    /*window_dimensions=*/{2, 3},
    /*window_stride_dimensions=*/{2, 3},
    Padding::kValid);

Stride of 1 in a dimension specifies that the position of a window in the dimension is 1 element away from its adjacent window. In order to specify that no windows overlap with each other, window_stride_dimensions should be equal to window_dimensions. The figure below illustrates the use of two different stride values. Padding is applied to each dimension of the input and the calculations are the same as though the input came in with the dimensions it has after padding.

For a non-trivial padding example, consider computing reduce-window minimum (initial value is MAX_FLOAT) with dimension 3 and stride 2 over the input array [10000, 1000, 100, 10, 1]. Padding kValid computes minimums over two valid windows: [10000, 1000, 100] and [100, 10, 1], resulting in the output [100, 1]. Padding kSame first pads the array so that the shape after the reduce-window would be the same as input for stride one by adding initial elements on both sides, getting [MAX_VALUE, 10000, 1000, 100, 10, 1, MAX_VALUE]. Running reduce-window over the padded array operates on three windows [MAX_VALUE, 10000, 1000], [1000, 100, 10], [10, 1, MAX_VALUE], and yields [1000, 10, 1].

The evaluation order of the reduction function is arbitrary and may be non-deterministic. Therefore, the reduction function should not be overly sensitive to reassociation. See the discussion about associativity in the context of Reduce for more details.

ReduceWindow - Example 2 - StableHLO

An example of ReduceWindow dataflow for StableHLO

In the above example:

Input) The operand has an input shape of S32[3,2]. With a values of [[1,2],[3,4],[5,6]]

Step 1) Base dilation with factor 2 along the row dimension inserts holes between each row of the operand. Padding of 2 rows at the top and 1 row at the bottom is applied after dilation. As a result, the tensor becomes taller.

Step 2) A window of shape [2,1] is defined, with window dilation [3,1]. This means each window selects two elements from the same column, but the second element is taken three rows below the first rather than directly beneath it.

Step 3) The windows are then slid across the operand with stride [4,1]. This causes the window to move down four rows at a time, while shifting one column at a time horizontally. Padding cells are filled with the init_value (in this case init_value = 0). Values ‘falling into’ dilation cells are ignored. Because of the stride and padding, some windows overlap only zeros and holes, while others overlap real input values.

Step 4) Within each window, the elements are combined using the reduction function (a, b) → a + b, starting from an initial value of 0. The top two windows see only padding and holes, so their results are 0. The bottom windows capture the values 3 and 4 from the input and return those as results.

Results) The final output has shape S32[2,2], with values: [[0,0],[3,4]]

Rem

ReplicaId

Reshape

See also XlaBuilder::Reshape. and the Collapse operation.

Reshapes the dimensions of an array into a new configuration.

Reshape(operand, dimensions)

Arguments	Type	Semantics
`operand`	`XlaOp`	array of type T
`dimensions`	`int64` vector	vector of sizes of new dimensions

Conceptually, reshape first flattens an array into a one-dimensional vector of data values, and then refines this vector into a new shape. The input arguments are an arbitrary array of type T, a compile-time-constant vector of dimension indices, and a compile-time-constant vector of dimension sizes for the result. The dimensions vector determines the size of the output array. The value at index 0 in dimensions is the size of dimension 0, the value at index 1 is the size of dimension 1, and so on. The product of the dimensions dimensions must equal the product of the operand's dimension sizes. When refining the collapsed array into the multidimensional array defined by dimensions, the dimensions in dimensions are ordered from slowest varying (most major) and to fastest varying (most minor).

For example, let v be an array of 24 elements:

let v = f32[4x2x3] {{{10, 11, 12}, {15, 16, 17}},
                    {{20, 21, 22}, {25, 26, 27}},
                    {{30, 31, 32}, {35, 36, 37}},
                    {{40, 41, 42}, {45, 46, 47}}};

let v012_24 = Reshape(v, {24});
then v012_24 == f32[24] {10, 11, 12, 15, 16, 17, 20, 21, 22, 25, 26, 27,
                         30, 31, 32, 35, 36, 37, 40, 41, 42, 45, 46, 47};

let v012_83 = Reshape(v, {8,3});
then v012_83 == f32[8x3] {{10, 11, 12}, {15, 16, 17},
                          {20, 21, 22}, {25, 26, 27},
                          {30, 31, 32}, {35, 36, 37},
                          {40, 41, 42}, {45, 46, 47}};

As a special case, reshape can transform a single-element array to a scalar and vice versa. For example,

Reshape(f32[1x1] {{5}}, {}) == 5;
Reshape(5, {1,1}) == f32[1x1] {{5}};

For StableHLO information see StableHLO - reshape.

Reshape (explicit)

Rev (reverse)

RngNormal

RngUniform

RngBitGenerator

RngGetAndUpdateState

The API of the various Rng operations are internally decomposed into HLO instructions including RngGetAndUpdateState.

RngGetAndUpdateState serves as a primitive in HLO. This op may appear in HLO dumps, but it is not intended to be constructed manually by end users.

Note: RngGetAndUpdateState is only found in HLO. It is not found in StableHLO.

Round

RoundNearestAfz

RoundNearestEven

Element-wise rounding, ties to the nearest even.

RoundNearestEven(operand)

Arguments	Type	Semantics
`operand`	`XlaOp`	The operand to the function

For StableHLO information see StableHLO - round_nearest_even.

Rsqrt

Scatter

Scatter - Example 1 - StableHLO

An example of Scatter dataflow for StableHLO

In the above image, each row of the table is an example of one update index example. Let's review stepwise from left(Update Index) to right(Result Index):

Input) input has shape S32[2,3,4,2]. scatter_indices have shape S64[2,2,3,2]. updates have shape S32[2,2,3,1,2].

Update Index) As part of the input we are given update_window_dims:[3,4]. This tell us that updates's dim 3 and dim 4 are window dimensions, highlighted in yellow. This allows us to derive that update_scatter_dims = [0,1,2].

Update Scatter Index) Shows us the extracted updated_scatter_dims for each. (The non-yellow of column Update Index)

Start Index) Looking at the scatter_indices tensor image we can see that our values from the previous step (Update scatter Index), give us the location of the start index. From index_vector_dim we are also told the dimension of the starting_indices that contains the starting indices, which for scatter_indices is dim 3 with a size 2.

Full Start Index) scatter_dims_to_operand_dims = [2,1] tells us the first element of the index vector goes to operand dim 2. The second element of the index vector goes to operand dim 1. The remaining operand dimensions are filled with 0.

Full Batching Index) We can see the purple highlighted area is shown in this column(full batching index), the update scatter index column, and update index column.

Full Window Index) Computed from the update_window_dimensions [3,4].

Result Index) The addition of Full Start Index, Full Batching Index, and Full Window Index in the operand tensor. Notice the green highlighted regions correspond to the operand figure as well. The last row is skipped because it falls outside of operand tensor.

Select

SelectAndScatter

This operation can be considered as a composite operation that first computes ReduceWindow on the operand array to select an element from each window, and then scatters the source array to the indices of the selected elements to construct an output array with the same shape as the operand array. The binary select function is used to select an element from each window by applying it across each window, and it is called with the property that the first parameter‘s index vector is lexicographically less than the second parameter’s index vector. The select function returns true if the first parameter is selected and returns false if the second parameter is selected, and the function must hold transitivity (i.e., if select(a, b) and select(b, c) are true, then select(a, c) is also true) so that the selected element does not depend on the order of the elements traversed for a given window.

The function scatter is applied at each selected index in the output array. It takes two scalar parameters:

Current value at the selected index in the output array
The scatter value from source that applies to the selected index

It combines the two parameters and returns a scalar value that's used to update the value at the selected index in the output array. Initially, all indices of the output array are set to init_value.

The output array has the same shape as the operand array and the source array must have the same shape as the result of applying a ReduceWindow operation on the operand array. SelectAndScatter can be used to backpropagate the gradient values for a pooling layer in a neural network.

SelectAndScatter(operand, select, window_dimensions, window_strides, padding, source, init_value, scatter)

Arguments	Type	Semantics
`operand`	`XlaOp`	array of type T over which the
: : : windows slide :
`select`	`XlaComputation`	binary computation of type `T, T
: : : -> PRED`, to apply to all :
: : : elements in each window; returns :
: : : `true` if the first parameter is :
: : : selected and returns `false` if :
: : : the second parameter is selected :
`window_dimensions`	`ArraySlice<int64>`	array of integers for window
: : : dimension values :
`window_strides`	`ArraySlice<int64>`	array of integers for window
: : : stride values :
`padding`	`Padding`	padding type for window
: : : (Padding::kSame or :
: : : Padding::kValid) :
`source`	`XlaOp`	array of type T with the values
: : : to scatter :
`init_value`	`XlaOp`	scalar value of type T for the
: : : initial value of the output :
: : : array :
`scatter`	`XlaComputation`	binary computation of type `T, T
: : : -> T`, to apply each scatter :
: : : source element with its :
: : : destination element :

The figure below shows examples of using SelectAndScatter, with the select function computing the maximal value among its parameters. Note that when the windows overlap, as in the figure (2) below, an index of the operand array may be selected multiple times by different windows. In the figure, the element of value 9 is selected by both of the top windows (blue and red) and the binary addition scatter function produces the output element of value 8 (2 + 6).

The evaluation order of the scatter function is arbitrary and may be non-deterministic. Therefore, the scatter function should not be overly sensitive to reassociation. See the discussion about associativity in the context of Reduce for more details.

For StableHLO information see StableHLO - select_and_scatter.

Send

SendDone

SendDone(HloInstruction context)

Given a context created by a Send instruction, waits for the data transfer to complete. The instruction does not return any data.

Note: SendDone is only found in HLO. It is not found in StableHLO.

Scheduling of channel instructions

The execution order of the 4 instructions for each channel (Recv, RecvDone, Send, SendDone) is as below.

Recv happens before Send
Send happens before RecvDone
Recv happens before RecvDone
Send happens before SendDone

When the backend compilers generate a linear schedule for each computation that communicates via channel instructions, there must not be cycles across the computations. For example, below schedules lead to deadlocks.

SetDimensionSize

Sets the dynamic size of XlaOp's given dimension. The operand must be array shaped.

SetDimensionSize(operand, val, dimension)

Arguments	Type	Semantics
`operand`	`XlaOp`	n dimensional input array.
`val`	`XlaOp`	int32 representing the runtime dynamic size.
`dimension`	`int64`	A value in the interval `[0, n)` that specifies the
: : : dimension. :

Pass through the operand as result, with dynamic dimension tracked by the compiler.

Padded values will be ignored by downstream reduction ops.

let v: f32[10] = f32[10]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
let five: s32 = 5;
let six: s32 = 6;

// Setting dynamic dimension size doesn't change the upper bound of the static
// shape.
let padded_v_five: f32[10] = set_dimension_size(v, five, /*dimension=*/0);
let padded_v_six: f32[10] = set_dimension_size(v, six, /*dimension=*/0);

// sum == 1 + 2 + 3 + 4 + 5
let sum:f32[] = reduce_sum(padded_v_five);
// product == 1 * 2 * 3 * 4 * 5
let product:f32[] = reduce_product(padded_v_five);

// Changing padding size will yield different result.
// sum == 1 + 2 + 3 + 4 + 5 + 6
let sum:f32[] = reduce_sum(padded_v_six);

ShiftLeft

ShiftRightArithmetic

Performs element-wise arithmetic right-shift operation on lhs by rhs number of bits.

ShiftRightArithmetic(lhs, rhs)

Arguments	Type	Semantics
lhs	XlaOp	Left-hand-side operand: array of type T
rhs	XlaOp	Left-hand-side operand: array of type T

An alternative variant with different-dimensional broadcasting support exists for ShiftRightArithmetic:

ShiftRightArithmetic(lhs,rhs, broadcast_dimensions)

Arguments	Type	Semantics
lhs	XlaOp	Left-hand-side operand: array of
: : : type T :
rhs	XlaOp	Left-hand-side operand: array of
: : : type T :
broadcast_dimension	ArraySlice	Which dimension in the target
: : : shape each dimension of the :
: : : operand shape corresponds to :

This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector).

The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the broadcasting page.

For StableHLO information see StableHLO - shift_right_arithmetic.

ShiftRightLogical

Performs element-wise logical right-shift operation on lhs by rhs number of bits.

ShiftRightLogical(lhs, rhs)

Arguments	Type	Semantics
lhs	XlaOp	Left-hand-side operand: array of type T
rhs	XlaOp	Left-hand-side operand: array of type T

An alternative variant with different-dimensional broadcasting support exists for ShiftRightLogical:

ShiftRightLogical(lhs,rhs, broadcast_dimensions)

Arguments	Type	Semantics
lhs	XlaOp	Left-hand-side operand: array of
: : : type T :
rhs	XlaOp	Left-hand-side operand: array of
: : : type T :
broadcast_dimension	ArraySlice	Which dimension in the target
: : : shape each dimension of the :
: : : operand shape corresponds to :

This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector).

The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the broadcasting page.

For StableHLO information see StableHLO - shift_right_logical.

Sign

Sin

Sin(operand) Element-wise sine x -> sin(x).

Arguments	Type	Semantics
`operand`	`XlaOp`	n dimensional input array
`split_dimension`	`int64`	A value in the interval
: : : `[0,n)` that names the :
: : : dimension along which the :
: : : operand is split :
`concat_dimension`	`int64`	A value in the interval
: : : `[0,n)` that names the :
: : : dimension along which the :
: : : split blocks are :
: : : concatenated :
`split_count`	`int64`	The number of cores that
: : : participate in this :
: : : operation. If :
: : : `replica_groups` is empty, :
: : : this should be the number of :
: : : replicas; otherwise, this :
: : : should be equal to the :
: : : number of replicas in each :
: : : group. :
`replica_groups`	`ReplicaGroup`vector	Each group contains a list
: : : of replica ids. :
`layout`	optional `Layout`	user-specified memory layout
`channel_id`	optional `ChannelHandle`	unique identifier for each
: : : send/recv pair :

Arguments	Type	Semantics
`predicate`	`XlaOp`	Scalar of type `PRED`
`true_operand`	`XlaOp`	Argument of type $T_0$
`true_computation`	`XlaComputation`	XlaComputation of type $T_0 \to S$
`false_operand`	`XlaOp`	Argument of type $T_1$
`false_computation`	`XlaComputation`	XlaComputation of type $T_1 \to S$

Arguments	Type	Semantics
`branch_index`	`XlaOp`	Scalar of type `S32`
`branch_computations`	sequence of N `XlaComputation`	XlaComputations of type $T_0 \to S , T_1 \to S , ..., T_{N-1} \to S$
`branch_operands`	sequence of N `XlaOp`	Arguments of type $T_0 , T_1 , ..., T_{N-1}$

Arguments	Type	Semantics
`operands`	sequence of N `XlaOp`	N arrays of type T with dimensions
: : : [L0, L1, ...]. Requires N >= 1. :
`dimension`	`int64`	A value in the interval `[0, N)` that
: : : names the dimension to be concatenated :
: : : between the `operands`. :

Input	Output	Semantics
vector [n] `dot` vector [n]	scalar	vector dot product
matrix [m x k] `dot` vector	vector [m]	matrix-vector multiplication
: [k] : : :
matrix [m x k] `dot` matrix	matrix [m x n]	matrix-matrix multiplication
: [k x n] : : :

DotDimensionNumbers Fields	Type	Semantics
`lhs_contracting_dimensions`	repeated int64	`lhs` contracting dimension
: : : numbers :
`rhs_contracting_dimensions`	repeated int64	`rhs` contracting dimension
: : : numbers :
`lhs_batch_dimensions`	repeated int64	`lhs` batch dimension
: : : numbers :
`rhs_batch_dimensions`	repeated int64	`rhs` batch dimension
: : : numbers :

Input	Output	Semantics
[b0, m, k] `dot` [b0, k, n]	[b0, m, n]	batch matmul
[b0, b1, m, k] `dot` [b0, b1, k, n]	[b0, b1, m, n]	batch matmul

Arguments	Type	Semantics
`operand`	`XlaOp`	The array we are Fourier
: : : transforming. :
`fft_type`	`FftType`	See the table below.
`fft_length`	`ArraySlice<int64>`	The time-domain lengths
: : : of the axes being :
: : : transformed. This is :
: : : needed in particular for :
: : : IRFFT to right-size the :
: : : innermost axis, since :
: : : `RFFT(fft_length=[16])` :
: : : has the same output :
: : : shape as :
: : : `RFFT(fft_length=[17])`. :

`FftType`	Semantics
`FFT`	Forward complex-to-complex FFT. Shape is unchanged.
`IFFT`	Inverse complex-to-complex FFT. Shape is unchanged.
`RFFT`	Forward real-to-complex FFT. Shape of the innermost axis is
: : reduced to `fft_length[-1] // 2 + 1` if `fft_length[-1]` is a :
: : non-zero value, omitting the reversed conjugate part of the :
: : transformed signal beyond the Nyquist frequency. :
`IRFFT`	Inverse real-to-complex FFT (i.e. takes complex, returns real).
: : Shape of the innermost axis is expanded to `fft_length[-1]` if :
: : `fft_length[-1]` is a non-zero value, inferring the part of the :
: : transformed signal beyond the Nyquist frequency from the reverse :
: : conjugate of the `1` to `fft_length[-1] // 2 + 1` entries. :

Argument	Type	Semantics
`tuple_data`	`XlaOP`	The tuple
`index`	`int64`	Index of tuple shape

Argument	Type	Semantics
`shape`	`Shape`	Shape of the data read from the Infeed
: : : interface. The layout field of the shape must :
: : : be set to match the layout of the data sent :
: : : to the device; otherwise its behavior is :
: : : undefined. :
`config`	optional `string`	Configuration of the op.

Arguments	Type	Semantics
`shape`	`Shape`	Shape of the array created by `Iota()`
`iota_dimension`	`int64`	The dimension to increment along.

Arguments	Type	Semantics
`operands`	sequence of N `XlaOp`s	N arrays of types T_0..T_{N-1}
`computation`	`XlaComputation`	Computation of type `T_0, T_1,
: : : .., T_{N + M -1} -> S` with N :
: : : parameters of type T and M of :
: : : arbitrary type. :
`dimensions`	`int64` array	Array of map dimensions
`static_operands`	sequence of N `XlaOp`s	Static ops for the map
: : : operation :

Arguments	Type	Semantics
`operand`	`XlaOp`	array of type `T`
`shape_with_layout`	`Shape`	Defines the layout of the data transferred
`outfeed_config`	`string`	Constant of config for the Outfeed instruction

Arguments	Type	Semantics
`shape`	`Shape`	shape of the data to receive
`handle`	`ChannelHandle`	unique identifier for each send/recv pair

Arguments	Type	Semantics
`operands`	Sequence of N `XlaOp`	N arrays of types `T_0,...,
: : : T_{N-1}`. :
`init_values`	Sequence of N `XlaOp`	N scalars of types
: : : `T_0,..., T_{N-1}`. :
`computation`	`XlaComputation`	computation of type
: : : `T_0,..., T_{N-1}, T_0, :
: : : ...,T_{N-1} ->` :
: : : `Collate(T_0,..., :
: : : T_{N-1})`. :
`dimensions_to_reduce`	`int64` array	unordered array of
: : : dimensions to reduce. :

Arguments	Type	Semantics
`operand`	`XlaOp`	array of floating-point type `T`.
`exponent_bits`	`int32`	number of exponent bits in lower-precision format
`mantissa_bits`	`int32`	number of mantissa bits in lower-precision format

Arguments	Type	Semantics
`operands`	`N XlaOps`	A sequence of N
: : : multi-dimensional arrays of :
: : : types `T_0,..., T_{N-1}`, each :
: : : representing the base area on :
: : : which the window is placed. :
`init_values`	`N XlaOps`	The N starting values for the
: : : reduction, one for each of the N :
: : : operands. See Reduce :
: : : for details. :
`computation`	`XlaComputation`	Reduction function of type `T_0,
: : : ..., T_{N-1}, T_0, ..., T_{N-1} :
: : : -> Collate(T_0, ..., T_{N-1})`, :
: : : to apply to elements in each :
: : : window of all the input :
: : : operands. :
`window_dimensions`	`ArraySlice<int64>`	array of integers for window
: : : dimension values :
`window_strides`	`ArraySlice<int64>`	array of integers for window
: : : stride values :
`base_dilations`	`ArraySlice<int64>`	array of integers for base
: : : dilation values :
`window_dilations`	`ArraySlice<int64>`	array of integers for window
: : : dilation values :
`padding`	`Padding`	padding type for window
: : : (Padding::kSame, which pads so :
: : : as to have the same output shape :
: : : as input if the stride is 1, or :
: : : Padding::kValid, which uses no :
: : : padding and “stops” the window :
: : : once it no longer fits) :

Arguments	Type	Semantics
`shape`	`Shape`	Output shape of type T
`operand`	`XlaOp`	array of type T

Arguments	Type	Semantics
`mu`	`XlaOp`	Scalar of type T specifying mean of generated numbers
`sigma`	`XlaOp`	Scalar of type T specifying standard deviation of
: : : generated :
`shape`	`Shape`	Output shape of type T

Arguments	Type	Semantics
`a`	`XlaOp`	Scalar of type T specifying lower limit of interval
`b`	`XlaOp`	Scalar of type T specifying upper limit of interval
`shape`	`Shape`	Output shape of type T

Arguments	Type	Semantics
`algorithm`	`RandomAlgorithm`	PRNG algorithm to be used.
`initial_state`	`XlaOp`	Initial state for the PRNG algorithm.
`shape`	`Shape`	Output shape for generated data.

Arguments	Type	Semantics
`operands`	Sequence of N `XlaOp`	N arrays of types `T_0, ..., T_N` to be scattered into.
`scatter_indices`	`XlaOp`	Array containing the starting indices of the slices that must be scattered to.
`updates`	Sequence of N `XlaOp`	N arrays of types `T_0, ..., T_N`. `updates[i]` contains the values that must be used for scattering `operands[i]`.
`update_computation`	`XlaComputation`	Computation to be used for combining the existing values in the input array and the updates during scatter. This computation should be of type `T_0, ..., T_N, T_0, ..., T_N -> Collate(T_0, ..., T_N)`.
`index_vector_dim`	`int64`	The dimension in `scatter_indices` that contains the starting indices.
`update_window_dims`	`ArraySlice<int64>`	The set of dimensions in `updates` shape that are window dimensions.
`inserted_window_dims`	`ArraySlice<int64>`	The set of window dimensions that must be inserted into `updates` shape.
`scatter_dims_to_operand_dims`	`ArraySlice<int64>`	A dimensions map from the scatter indices to the operand index space. This array is interpreted as mapping `i` to `scatter_dims_to_operand_dims[i]` . It has to be one-to-one and total.
`dimension_number`	`ScatterDimensionNumbers`	Dimension numbers for scatter operation
`indices_are_sorted`	`bool`	Whether the indices are guaranteed to be sorted by the caller.
`unique_indices`	`bool`	Whether the indices are guaranteed to be unique by the caller.

Arguments	Type	Semantics
`pred`	`XlaOp`	array of type PRED
`on_true`	`XlaOp`	array of type T
`on_false`	`XlaOp`	array of type T

Arguments	Type	Semantics
`operand`	`XlaOp`	data to send (array of type T)
`handle`	`ChannelHandle`	unique identifier for each send/recv pair

Arguments	Type	Semantics
`operands`	`ArraySlice<XlaOp>`	The operands to sort.
`comparator`	`XlaComputation`	The comparator computation to use.
`dimension`	`int64`	The dimension along which to sort.
`is_stable`	`bool`	Whether stable sorting should be used.

Arguments	Type	Semantics
`operand`	`XlaOp`	The tensor from which to extract the top `k` elements.
: : : The tensor must have greater or equal to one :
: : : dimensions. The size of the last dimension of the :
: : : tensor must be greater or equal to `k`. :
`k`	`int64`	The number of elements to extract.
`largest`	`bool`	Whether to extract the largest or smallest `k`
: : : elements. :

Arguments	Type	Semantics
`operand`	`XlaOp`	The operand to transpose.
`permutation`	`ArraySlice<int64>`	How to permute the dimensions.

Arguments	Type	Semantics
`a`	`XlaOp`	a > 2 dimensional array of a complex or
: : : floating-point type with shape `[..., M, :
: : : M]`. :
`b`	`XlaOp`	a > 2 dimensional array of the same type
: : : with shape `[..., M, K]` if `left_side` is :
: : : true, `[..., K, M]` otherwise. :
`left_side`	`bool`	indicates whether to solve a system of the
: : : form `op(a) * x = b` (`true`) or `x * :
: : : op(a) = b `(`false`). :
`lower`	`bool`	whether to use the upper or lower triangle
: : : of `a`. :
`unit_diagonal`	`bool`	if `true`, the diagonal elements of `a` are
: : : assumed to be `1` and not accessed. :
`transpose_a`	`Transpose`	whether to use `a` as is, transpose it or
: : : take its conjugate transpose. :

Arguments	Type	Semantics
`condition`	`XlaComputation`	XlaComputation of type `T -> PRED` which
: : : defines the termination condition of the :
: : : loop. :
`body`	`XlaComputation`	XlaComputation of type `T -> T` which
: : : defines the body of the loop. :
`init`	`T`	Initial value for the parameter of
: : : `condition` and `body`. :