操作语义

下面介绍了 XlaBuilder 接口中定义的操作的语义。通常，这些操作会一对一映射到 xla_data.proto 的 RPC 接口中定义的操作。

有关命名法的注意事项：XLA 处理的广义数据类型是一个 N 维数组，其中包含某种统一类型的元素（例如 32 位浮点）。在整个文档中，“数组”都用于表示任意维数组。为方便起见，特殊情况具有更加具体和熟悉的名称；例如，向量是一维数组，矩阵是二维数组。

AfterAll

另请参阅 XlaBuilder::AfterAll。

AfterAll 获取可变数量的词元并生成一个词元。令牌是基元类型，可以在附带效应的操作之间进行线程处理，以强制执行排序。AfterAll 可用作令牌联接，用于在集合操作之后对操作进行排序。

AfterAll(operands)

参数	类型	语义
operands	XlaOp	可变数量的词元

AllGather

另请参阅 XlaBuilder::AllGather。

跨副本执行串联。

AllGather(operand, all_gather_dim, shard_count, replica_group_ids, channel_id)

参数	类型	语义
operand	XlaOp	用于跨副本串联的数组
all_gather_dim	int64	串联维度
replica_groups	int64 的向量	执行串联的组
channel_id	可选 int64	用于跨模块通信的可选通道 ID

• replica_groups 是执行串联的副本组的列表（可以使用 ReplicaId 检索当前副本的副本 ID）。每个组中副本的顺序决定了其输入在结果中的顺序。replica_groups 必须为空（在这种情况下，所有副本都属于一个组，按 0 到 N - 1 排序），或者包含与副本数量相同的元素数量。例如，replica_groups = {0, 2}, {1, 3} 在副本 0 和 2 以及 1 和 3 之间执行串联。
• shard_count 是每个副本组的大小。在 replica_groups 为空的情况下，我们需要使用该属性。
• channel_id 用于跨模块通信：只有具有相同 channel_id 的 all-gather 操作可以相互通信。

输出形状是 all_gather_dim 增大了 shard_count 倍的输入形状。例如，如果有两个副本，并且操作数在两个副本上分别为 [1.0, 2.5] 和 [3.0, 5.25]，则 all_gather_dim 为 0 时，此操作在两个副本上的输出值将为 [1.0, 2.5, 3.0, 5.25]。

AllReduce

另请参阅 XlaBuilder::AllReduce。

跨副本执行自定义计算。

AllReduce(operand, computation, replica_group_ids, channel_id)

参数	类型	语义
operand	XlaOp	减少副本数量，减少数组中的数组或非空数组
computation	XlaComputation	归约计算
replica_groups	int64 的向量	执行缩减的组
channel_id	可选 int64	用于跨模块通信的可选通道 ID

• 当 operand 是数组元组时，对元组的每个元素执行全部归约。
• replica_groups 是执行缩减操作的副本组的列表（可以使用 ReplicaId 检索当前副本的副本 ID）。replica_groups 必须为空（在这种情况下，所有副本都属于一个组），或者包含与副本数相同的元素数量。例如，replica_groups = {0, 2}, {1, 3} 在副本 0 和 2 以及 1 和 3 之间执行缩减。
• channel_id 用于跨模块通信：只有具有相同 channel_id 的 all-reduce 操作可以相互通信。

输出形状与输入形状相同。例如，如果有两个副本，并且操作数在两个副本上分别具有值 [1.0, 2.5] 和 [3.0, 5.25]，则此操作和求和计算在两个副本上的输出值将为 [4.0, 7.75]。如果输入是元组，则输出也是元组。

计算 AllReduce 的结果需要来自每个副本的一个输入，因此，如果一个副本执行 AllReduce 节点的次数多于另一个副本，则前一个副本将一直等待。由于所有副本都在运行同一程序，因此要实现这一目的的方法并不多，但当 when 循环的条件依赖于馈入数据，并且馈入的数据导致 when 循环在一个副本上迭代次数多于另一个副本时，就有可能发生这种情况。

AllToAll

另请参阅 XlaBuilder::AllToAll。

AllToAll 是一个从所有核心向所有核心发送数据的集体操作。它包含两个阶段：

1. 散点阶段。在每个核心上，运算数会沿着 split_dimensions 拆分为 split_count 个块，这些块会分散到所有核心，例如，第 i 个块会发送到第 i 个核心。
2. 收集阶段。每个核心都会沿着 concat_dimension 串联收到的块。

可通过以下方式配置参与的核心：

• replica_groups：每个 ReplicaGroup 包含参与计算的副本 ID 列表（可以使用 ReplicaId 检索当前副本的副本 ID）。AllToAll 将按指定顺序应用于子组。例如，replica_groups = { {1,2,3}, {4,5,0} } 表示 AllToAll 将在副本 {1, 2, 3} 内和收集阶段应用，并且收到的块将按照 1、2、3 的相同顺序串联。然后，系统将在副本 4、5、0 内应用另一个 AllToAll，串联顺序也是 4、5、0。如果 replica_groups 为空，则所有副本按照它们的串联顺序属于一个组。

前提条件：

• split_dimension 上运算数的维度大小可被 split_count 整除。
• 操作数的形状不是元组。

AllToAll(operand, split_dimension, concat_dimension, split_count, replica_groups)

参数	类型	语义
operand	XlaOp	N 维输入数组
split_dimension	int64	间隔 [0, n) 中的值，用于为拆分运算数所依据的维度命名
concat_dimension	int64	间隔 [0, n) 中的一个值，用于命名串联拆分块所依据的维度
split_count	int64	参与此操作的核心数量。如果 replica_groups 为空，则此值应为副本数量；否则，此值应等于每个组中的副本数量。
replica_groups	ReplicaGroup 矢量	每个组都包含一个副本 ID 列表。

下面是一个 Alltoall 示例。

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);

在本例中，有 4 个核心参与 Alltoall。在每个核心上，运算数沿维度 0 拆分为 4 个部分，因此每个部分都具有形状 f32[4,4]。这 4 个部分分布到所有核心。然后，每个核心会按核心 0-4 的顺序沿维度 1 串联接收到的部分。因此，每个核心上的输出的形状为 f32[16,4]。

BatchNormGrad

如需详细了解该算法，另请参阅 XlaBuilder::BatchNormGrad 和原始批量归一化论文。

计算批量范数的梯度。

BatchNormGrad(operand, scale, mean, variance, grad_output, epsilon, feature_index)

参数	类型	语义
operand	XlaOp	要规范化的 n 维数组 (x)
scale	XlaOp	一维数组 (\(\gamma\))
mean	XlaOp	一维数组 (\(\mu\))
variance	XlaOp	一维数组 (\(\sigma^2\))
grad_output	XlaOp	传递到 BatchNormTraining (\(\nabla y\)) 的梯度
epsilon	float	艾普西隆值 (\(\epsilon\))
feature_index	int64	“operand”中的特征维度的索引

对于特征维度中的每个特征（feature_index 是 operand 中特征维度的索引），该操作会计算 operand、offset 和 scale 相对于所有其他维度的梯度。feature_index 必须是 operand 中特征维度的有效索引。

这三种渐变分别使用以下公式进行定义（假设一个四维数组为 operand，并且特征维度索引为 l，批量大小为 m，空间大小为 w 和 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} \]

输入 mean 和 variance 表示跨批量和空间维度的时刻值。

输出类型是包含三个句柄的元组：

输出	类型	语义
grad_operand	XlaOp	相对于输入 operand 的梯度 ($\nabla x$)
grad_scale	XlaOp	相对于输入 scale 的梯度 ($\nabla \gamma$)
grad_offset	XlaOp	相对于输入 offset 的梯度($\nabla \beta$)

BatchNormInference

如需详细了解该算法，另请参阅 XlaBuilder::BatchNormInference 和原始批量归一化论文。

跨批量维度和空间维度对数组进行归一化。

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

参数	类型	语义
operand	XlaOp	要归一化的 n 维数组
scale	XlaOp	一维数组
offset	XlaOp	一维数组
mean	XlaOp	一维数组
variance	XlaOp	一维数组
epsilon	float	艾普西隆值
feature_index	int64	“operand”中的特征维度的索引

对于特征维度中的每个特征（feature_index 是 operand 中特征维度的索引），该操作会计算所有其他维度的平均值和方差，并使用平均值和方差对 operand 中的每个元素进行归一化。feature_index 必须是 operand 中特征维度的有效索引。

BatchNormInference 相当于在不计算每个批次的 mean 和 variance 的情况下调用 BatchNormTraining。它会改用输入 mean 和 variance 作为估算值。此操作的目的是减少推理中的延迟时间，因此被命名为 BatchNormInference。

输出是一个 N 维归一化数组，其形状与输入 operand 相同。

BatchNormTraining

如需详细了解该算法，另请参阅 XlaBuilder::BatchNormTraining 和 the original batch normalization paper。

跨批量维度和空间维度对数组进行归一化。

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

参数	类型	语义
operand	XlaOp	要规范化的 n 维数组 (x)
scale	XlaOp	一维数组 (\(\gamma\))
offset	XlaOp	一维数组 (\(\beta\))
epsilon	float	艾普西隆值 (\(\epsilon\))
feature_index	int64	“operand”中的特征维度的索引

对于特征维度中的每个特征（feature_index 是 operand 中特征维度的索引），该操作会计算所有其他维度的平均值和方差，并使用平均值和方差对 operand 中的每个元素进行归一化。feature_index 必须是 operand 中特征维度的有效索引。

对于 operand \(x\) 中包含 m 元素且空间维度大小为 w 和 h 的每个批次，算法如下（假设 operand 是一个四维数组）：

• 计算特征维度中 \(\mu_l\) 每个特征的l批量平均值： \(\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}\)

• 计算批量方差 \(\sigma^2_l\)： $\sigma^2l=\frac{1}{mwh}\sum{i=1}^m\sum{j=1}^w\sum{k=1}^h (x_{ijkl} - \mu_l)^2$

• 归一化、缩放和偏移：\(y_{ijkl}=\frac{\gamma_l(x_{ijkl}-\mu_l)}{\sqrt[2]{\sigma^2_l+\epsilon} }+\beta_l\)

添加 epsilon 值（通常是一个较小的数字）可以避免除零错误。

输出类型是包含三个 XlaOp 的元组：

输出	类型	语义
output	XlaOp	形状与输入 operand (y) 相同的 n 维数组
batch_mean	XlaOp	一维数组 (\(\mu\))
batch_var	XlaOp	一维数组 (\(\sigma^2\))

batch_mean 和 batch_var 是使用上述公式根据批次和空间维度计算的瞬间。

BitcastConvertType

另请参阅 XlaBuilder::BitcastConvertType。

与 TensorFlow 中的 tf.bitcast 类似，可执行从数据形状到目标形状的元素级位投射操作。输入大小和输出大小必须匹配：例如，s32 元素通过位播例程变为 f32 元素，一个 s32 元素变为四个 s8 元素。Bitcast 以低级别类型转换的形式实现，因此具有不同浮点表示法的机器将给出不同的结果。

BitcastConvertType(operand, new_element_type)

参数	类型	语义
operand	XlaOp	维度为 D 的类型为 T 的数组
new_element_type	PrimitiveType	U 型

运算数和目标形状的尺寸必须匹配，但最后一个维度除外，该维度将根据转换前后原始大小的比率进行更改。

源和目标元素类型不得为元组。

将 Bitcast 转换为不同宽度的基元类型

BitcastConvert HLO 指令支持输出元素类型 T' 的大小不等于输入元素 T 的大小的情况。由于整个操作从概念上来讲是一种位图，不会更改底层字节，因此输出元素的形状必须发生变化。对于 B = sizeof(T), B' = sizeof(T')，有两种可能的情况。

首先，当 B > B' 时，输出形状会获得大小为 B/B' 的新最次要维度。例如：

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

对于有效标量，该规则仍然相同：

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

或者，对于 B' > B，该指令要求输入形状的最后一个逻辑维度必须等于 B'/B，并且此维度在转换期间会被丢弃：

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

请注意，不同位宽之间的转换不是元素级的。

广播

另请参阅 XlaBuilder::Broadcast。

通过复制数组中的数据，向数组添加维度。

Broadcast(operand, broadcast_sizes)

参数	类型	语义
operand	XlaOp	要复制的数组
broadcast_sizes	ArraySlice<int64>	新维度的尺寸

新维度会在左侧插入，也就是说，如果 broadcast_sizes 的值为 {a0, ..., aN}，且运算数形状的尺寸为 {b0, ..., bM}，则输出形状的尺寸为 {a0, ..., aN, b0, ..., bM}。

新维度将编入运算数副本的索引，即

output[i0, ..., iN, j0, ..., jM] = operand[j0, ..., jM]

例如，如果 operand 是值为 2.0f 的标量 f32，且 broadcast_sizes 为 {2, 3}，则结果将是一个形状为 f32[2, 3] 的数组，且结果中的所有值都将为 2.0f。

BroadcastInDim

另请参阅 XlaBuilder::BroadcastInDim。

通过复制数组中的数据来扩展数组的大小和排名。

BroadcastInDim(operand, out_dim_size, broadcast_dimensions)

参数	类型	语义
operand	XlaOp	要复制的数组
out_dim_size	ArraySlice<int64>	目标形状的尺寸大小
broadcast_dimensions	ArraySlice<int64>	目标形状中的哪个维度，运算数形状的每个维度对应

与广播类似，但允许在任何位置添加维度以及扩展尺寸为 1 的现有维度。

operand 会广播到 out_dim_size 所描述的形状。broadcast_dimensions 会将 operand 的尺寸映射到目标形状的维度，即操作数的第 i 个维度会映射到输出形状的 broadcast_dimension[i] 维度。operand 的维度大小必须为 1，或者与其映射到的输出形状中的维度相同。其余尺寸填充为尺寸 1。然后，沿这些退化维度进行广播，以达到输出形状。广播页面详细介绍了这些语义。

Call

另请参阅 XlaBuilder::Call。

使用给定参数调用计算。

Call(computation, args...)

参数	类型	语义
computation	XlaComputation	T_0, T_1, ..., T_{N-1} -> S 类型的计算，有 N 个任意类型的参数
args	N XlaOp 序列	N 个任意类型的参数

args 的奇数和类型必须与 computation 的参数一致。不允许有 args。

乔尔斯基

另请参阅 XlaBuilder::Cholesky。

计算一批对称（赫米氏）正定矩阵的 Cholesky 分解。

Cholesky(a, lower)

参数	类型	语义
a	XlaOp	一个复杂或浮点类型的 rank > 2 数组。
lower	bool	使用 a 的上三角形还是下三角形。

如果 lower 为 true，则计算下三角矩阵 l，以使 $a = l。 l^T$。如果 lower 为 false，则计算上三角矩阵 u，使得\(a = u^T . u\)。

根据 lower 的值，仅从 a 的下三角形/上三角形读取输入数据。来自另一个三角形的值会被忽略。输出数据会在同一个三角形中返回；另一个三角形中的值由实现定义，可以是任何内容。

如果 a 的秩大于 2，则 a 会被视为一批矩阵，其中除了次要的 2 个维度都属于批量维度。

如果 a 不是对称（惠米氏）正定定，则结果由实现定义。

夹式

另请参阅 XlaBuilder::Clamp。

将运算数限制在最小值和最大值之间的范围内。

Clamp(min, operand, max)

参数	类型	语义
min	XlaOp	类型为 T 的数组
operand	XlaOp	类型为 T 的数组
max	XlaOp	类型为 T 的数组

给定一个运算数以及最小值和最大值，当运算数介于最小值和最大值之间时，返回最小值；如果运算数低于此范围，则返回最小值；如果运算数高于此范围，则返回最大值。即 clamp(a, x, b) = min(max(a, x), b)。

所有三个数组的形状必须相同。或者，作为受限形式的广播，min 和/或 max 可以是 T 类型的标量。

带有标量 min 和 max 的示例：

let operand: s32[3] = {-1, 5, 9};
let min: s32 = 0;
let max: s32 = 6;
==>
Clamp(min, operand, max) = s32[3]{0, 5, 6};

收起

另请参阅 XlaBuilder::Collapse 和 tf.reshape 运算。

将数组的维度收起为一个维度。

Collapse(operand, dimensions)

参数	类型	语义
operand	XlaOp	类型为 T 的数组
dimensions	int64矢量	T 维度的有序连续子集。

收起操作会将运算数的给定子集替换为单个维度。输入参数是类型 T 的任意数组和维度索引的编译时常量向量。维度索引必须是 T 维度的按顺序（从低到高）连续子集。因此，{0, 1, 2}、{0, 1} 或 {1, 2} 都是有效的维度集，但 {1, 0} 或 {0, 2} 不是。这些维度会被单个新维度取代，新维度在维度序列中的位置与所替换的维度相同，且新维度大小等于原始维度大小的乘积。dimensions 中的最低维度编号是循环嵌套中最慢的变分维度（最主要），它可以收起这些维度，最高维度编号是最快变分（最次要）。如果需要更笼统的收起排序，请参阅 tf.reshape 运算符。

例如，假设 v 是一个包含 24 个元素的数组：

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

CollectivePermute

另请参阅 XlaBuilder::CollectivePermute。

CollectivePermute 是一个发送和接收跨副本数据的集合操作。

CollectivePermute(operand, source_target_pairs)

参数	类型	语义
operand	XlaOp	N 维输入数组
source_target_pairs	<int64, int64>矢量	一系列 (source_Replica_id, target_copy_id) 对。对于每个键值对，操作数会从源副本发送到目标副本。

请注意，source_target_pair 存在以下限制：

• 任何两个对都不应具有相同的目标副本 ID，且不应具有相同的源副本 ID。
• 如果副本 ID 不是任何对中的目标，则该副本的输出是一个张量，由 0 组成，其形状与输入相同。

串联

另请参阅 XlaBuilder::ConcatInDim。

串联可组合多个数组运算数中的数组。该数组与每个输入数组运算数具有相同的排名（它们必须彼此相同），并且按指定顺序包含参数。

Concatenate(operands..., dimension)

参数	类型	语义
operands	N XlaOp 序列	具有维度 [L0, L1, ...] 的 N 个类型 T 数组。需要 N >= 1。
dimension	int64	间隔 [0, N) 中的值，用于指定要在 operands 之间串联的维度。

除 dimension 外，所有维度都必须相同。这是因为 XLA 不支持“不规则”数组。另请注意，rank-0 值无法串联（因为无法命名发生串联的维度）。

一维示例：

Concat({ {2, 3}, {4, 5}, {6, 7} }, 0)
>>> {2, 3, 4, 5, 6, 7}

二维示例：

let a = {
{1, 2},
{3, 4},
{5, 6},
};
let b = {
{7, 8},
};
Concat({a, b}, 0)
>>> {
{1, 2},
{3, 4},
{5, 6},
{7, 8},
}

图表：

基于条件

另请参阅 XlaBuilder::Conditional。

Conditional(pred, true_operand, true_computation, false_operand, false_computation)

参数	类型	语义
pred	XlaOp	PRED 类型的标量
true_operand	XlaOp	类型的参数 \(T_0\)
true_computation	XlaComputation	\(T_0 \to S\)类型的 XlaComputation
false_operand	XlaOp	类型的参数 \(T_1\)
false_computation	XlaComputation	\(T_1 \to S\)类型的 XlaComputation

如果 pred 为 true，则执行 true_computation；如果 pred 为 false，则执行 false_computation，并返回结果。

true_computation 必须接受类型为 \(T_0\) 的单个参数，并且将使用必须属于同一类型的 true_operand 进行调用。false_computation 必须接受类型为 \(T_1\) 的单个参数，并且将使用必须属于同一类型的 false_operand 进行调用。true_computation 和 false_computation 的返回值类型必须相同。

请注意，系统仅会执行 true_computation 和 false_computation 中的一项，具体取决于 pred 的值。

Conditional(branch_index, branch_computations, branch_operands)

参数	类型	语义
branch_index	XlaOp	S32 类型的标量
branch_computations	N XlaComputation 序列	\(T_0 \to S , T_1 \to S , ..., T_{N-1} \to S\)类型的 XlaComputations
branch_operands	N XlaOp 序列	\(T_0 , T_1 , ..., T_{N-1}\)类型的参数

执行 branch_computations[branch_index] 并返回结果。如果 branch_index 是 < 0 或 >= N 的 S32，则 branch_computations[N-1] 会作为默认分支执行。

每个 branch_computations[b] 都必须接受一个类型为 \(T_b\) 的参数，并且将使用必须属于同一类型的 branch_operands[b] 进行调用。每个 branch_computations[b] 的返回值的类型必须相同。

请注意，根据 branch_index 的值，系统只会执行其中一个 branch_computations。

转化（卷积）

另请参阅 XlaBuilder::Conv。

如 ConvWithGeneralPadding 中，但内边距是以简单方式指定为 SAME 或 VALID。SAME 填充会为输入 (lhs) 填充零，以便在不考虑步长的情况下，输出的形状与输入相同。VALID 填充仅仅表示没有内边距。

ConvWithGeneralPadding（卷积）

另请参阅 XlaBuilder::ConvWithGeneralPadding。

计算神经网络中使用的卷积。在这里，卷积可以视为在 n 维基本区域移动的 n 维窗口，系统会对窗口的每个可能位置执行计算。

参数	类型	语义
lhs	XlaOp	排序 n+2 的输入数组
rhs	XlaOp	n+2 核权重数组
window_strides	ArraySlice<int64>	内核步长的 n-d 数组
padding	ArraySlice< pair<int64,int64>>	（低、高）填充的 N 维数组
lhs_dilation	ArraySlice<int64>	n-d lhs 膨胀因数数组
rhs_dilation	ArraySlice<int64>	n-d rhs 膨胀因数数组
feature_group_count	int64	特征组的数量
batch_group_count	int64	批次组的数量

让 n 是空间维度的数量。lhs 参数是描述底面积的 n+2 秩数组。这称为输入，尽管 rhs 也是一个输入。在神经网络中，这些是输入激活。n+2 个维度按以下顺序排列：

• batch：此维度中的每个坐标表示执行卷积的独立输入。
• z/depth/features：底区中的每个 (y,x) 位置都有一个与之关联的矢量，该矢量属于此维度。
• spatial_dims：描述 n 空间维度，这些空间维度用于定义窗口移动的基本区域。

rhs 参数是描述卷积过滤器/内核/窗口的 n+2 秩数组。这些维度按以下顺序排列：

• output-z：输出的 z 维度。
• input-z：此维度的大小乘以 feature_group_count 应等于 z 维度的大小（以 lhs 为单位）。
• spatial_dims：描述 n 空间维度，这些维度用于定义在基本区域移动的 n-d 窗口。

window_strides 参数指定卷积窗口在空间维度上的步长。例如，如果第一个空间维度的步长为 3，则窗口只能放置在第一个空间索引可被 3 整除的坐标处。

padding 参数指定要应用到基本区域的零内边距大小。内边距可为负值，负内边距的绝对值表示在执行卷积之前要从指定维度中移除的元素数。padding[0] 指定尺寸 y 的内边距，padding[1] 指定维度 x 的内边距。每对中，低内边距作为第一个元素，高内边距作为第二个元素。低内边距会应用于较低索引的方向，而高内边距会应用于较高索引的方向。例如，如果 padding[1] 为 (2,3)，则在第二个空间维度中，左侧填充 2 个零，右侧填充 3 个零。使用填充相当于在执行卷积之前将相同的零值插入输入 (lhs)。

lhs_dilation 和 rhs_dilation 参数分别指定要在每个空间维度中应用于 lhs 和 rhs 的膨胀因数。如果空间维度中的膨胀因数为 d，则系统会在该维度中的每个条目之间隐式放置 d-1 孔，从而增加数组的大小。这些孔用空操作值填充，对于卷积意味着零。

rh 的膨胀也称为室卷积。如需了解详情，请参阅 tf.nn.atrous_conv2d。lhs 的膨胀也称为转置卷积。如需了解详情，请参阅 tf.nn.conv2d_transpose。

feature_group_count 参数（默认值为 1）可用于分组卷积。feature_group_count 需要是输入和输出特征维度的除数。如果 feature_group_count 大于 1，则表示从概念上讲，输入和输出特征维度以及 rhs 输出特征维度被平均拆分为多个 feature_group_count 组，每个组由连续的特征子序列组成。rhs 的输入特征维度需要等于 lhs 输入特征维度除以 feature_group_count（因此它已经具有一组输入特征的大小）。第 i 组用于计算许多单独卷积的 feature_group_count。这些卷积的结果会在输出特征维度中串联在一起。

对于深度卷积，feature_group_count 参数将设置为输入特征维度，并将过滤器的形状从 [filter_height, filter_width, in_channels, channel_multiplier] 重新调整为 [filter_height, filter_width, 1, in_channels * channel_multiplier]。如需了解详情，请参阅 tf.nn.depthwise_conv2d。

batch_group_count（默认值 1）参数可用于反向传播期间的已分组过滤器。batch_group_count 需要是 lhs（输入）批量维度大小的除数。如果 batch_group_count 大于 1，则表示输出批量维度的大小应为 input batch / batch_group_count。batch_group_count 必须是输出特征大小的除数。

输出形状具有以下尺寸：

• batch：此维度的大小乘以 batch_group_count 应与 batch 维度的大小（以 lhs 为单位）相等。
• z：大小与内核上的 output-z (rhs) 相同。
• spatial_dims：卷积窗口的每个有效位置各对应一个值。

上图显示了 batch_group_count 字段的工作原理。实际上，我们将每个 lhs 批次拆分为 batch_group_count 组，并对输出特征执行相同的操作。然后，对于每一组，我们进行成对卷积，并沿输出特征维度串联输出。所有其他维度（特征和空间）的操作语义保持不变。

卷积窗口的有效位置由内边距之后的步长和基础区域的大小决定。

为了描述卷积的作用，请考虑一个 2D 卷积，并在输出中选择一些固定的 batch、z、y、x 坐标。(y,x) 是窗口某个角在基本区域（例如左上角，具体取决于您如何解读空间维度）内的位置。现在，我们得到了一个取自基本区域的 2D 窗口，其中每个 2D 点都与一个一维矢量相关联，因此我们得到一个 3D 方框。在卷积内核中，由于我们固定了输出坐标 z，因此还有一个 3D 方框。这两个框的尺寸相同，因此我们可以取两个框之间的元素级积的总和（类似于点积）。这是输出值。

请注意，如果 output-z 是5，则窗口的每个位置会在输出的 z 维度中生成 5 个值。这些值的不同在于使用的是卷积内核的哪个部分，每个 output-z 坐标都有一个单独的 3D 框。因此您可以将其视为 5 个单独的卷积，并为每个卷积提供不同的过滤器。

以下是含内边距和步长的 2D 卷积的伪代码：

for (b, oz, oy, ox) {  // output coordinates
  value = 0;
  for (iz, ky, kx) {  // kernel coordinates and input z
    iy = oy*stride_y + ky - pad_low_y;
    ix = ox*stride_x + kx - pad_low_x;
    if ((iy, ix) inside the base area considered without padding) {
      value += input(b, iz, iy, ix) * kernel(oz, iz, ky, kx);
    }
  }
  output(b, oz, oy, ox) = value;
}

ConvertElementType

另请参阅 XlaBuilder::ConvertElementType。

与 C++ 中的元素级 static_cast 类似，可执行从数据形状到目标形状的元素级转换操作。维度必须一致，并且转换是元素级转换；例如，通过 s32 到 f32 转换例程，s32 元素会变为 f32 元素。

ConvertElementType(operand, new_element_type)

参数	类型	语义
operand	XlaOp	维度为 D 的类型为 T 的数组
new_element_type	PrimitiveType	U 型

操作数和目标形状的尺寸必须匹配。源和目标元素类型不得为元组。

从 T=s32 到 U=f32 的转换将执行归一化整数到浮点数的转换例程，例如舍入到最近均匀。

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

CrossReplicaSum

通过求和计算执行 AllReduce。

CustomCall

另请参阅 XlaBuilder::CustomCall。

在计算中调用用户提供的函数。

CustomCall(target_name, args..., shape)

参数	类型	语义
target_name	string	函数的名称。系统将发出针对该符号名称的调用指令。
args	N XlaOp 序列	有 N 个任意类型的参数，它们将传递给函数。
shape	Shape	函数的输出形状

无论参数的元数或类型如何，函数签名都相同：

extern "C" void target_name(void* out, void** in);

例如，如果按如下方式使用 CustomCall：

let x = f32[2] {1,2};
let y = f32[2x3] { {10, 20, 30}, {40, 50, 60} };

CustomCall("myfunc", {x, y}, f32[3x3])

下面是 myfunc 的实现的示例：

extern "C" void myfunc(void* out, void** in) {
  float (&x)[2] = *static_cast<float(*)[2]>(in[0]);
  float (&y)[2][3] = *static_cast<float(*)[2][3]>(in[1]);
  EXPECT_EQ(1, x[0]);
  EXPECT_EQ(2, x[1]);
  EXPECT_EQ(10, y[0][0]);
  EXPECT_EQ(20, y[0][1]);
  EXPECT_EQ(30, y[0][2]);
  EXPECT_EQ(40, y[1][0]);
  EXPECT_EQ(50, y[1][1]);
  EXPECT_EQ(60, y[1][2]);
  float (&z)[3][3] = *static_cast<float(*)[3][3]>(out);
  z[0][0] = x[1] + y[1][0];
  // ...
}

用户提供的函数不得产生副作用，并且其执行方式必须是幂等的。

Dot

另请参阅 XlaBuilder::Dot。

Dot(lhs, rhs)

参数	类型	语义
lhs	XlaOp	类型为 T 的数组
rhs	XlaOp	类型为 T 的数组

此操作的确切语义取决于运算数的排名：

输入	输出	语义
矢量 [n] dot 矢量 [n]	标量	矢量点积
矩阵 [m x k] dot 向量 [k]	矢量 [m]	矩阵向量乘法
矩阵 [m x k] dot 矩阵 [k x n]	矩阵 [m x n]	矩阵乘法

运算对 lhs 的第二个维度（如果排名为 1，则为第一个维度）和 rhs 的第一个维度执行乘积之和。这些维度是“简写”维度lhs 和 rhs 的合同规定的尺寸必须具有相同的大小。实际上，它可用于执行向量之间的点积、向量/矩阵乘法或矩阵/矩阵乘法。

DotGeneral

另请参阅 XlaBuilder::DotGeneral。

DotGeneral(lhs, rhs, dimension_numbers)

参数	类型	语义
lhs	XlaOp	类型为 T 的数组
rhs	XlaOp	类型为 T 的数组
dimension_numbers	DotDimensionNumbers	协定和批量维度编号

与 Dot 类似，但允许为 lhs 和 rhs 指定协定和批量维度编号。

DotDimensionNumbers 字段	类型	语义
lhs_contracting_dimensions	重复的 int64	lhs 个合同维度编号
rhs_contracting_dimensions	重复的 int64	rhs 个合同维度编号
lhs_batch_dimensions	重复的 int64	lhs 个批量维度编号
rhs_batch_dimensions	重复的 int64	rhs 个批量维度编号

DotGeneral 会根据 dimension_numbers 中指定的合同维度执行乘积之和。

lhs 和 rhs 中关联的收缩维度编号不必相同，但必须具有相同的维度大小。

具有合同维度编号的示例：

lhs = { {1.0, 2.0, 3.0},
{4.0, 5.0, 6.0} }

rhs = { {1.0, 1.0, 1.0},
{2.0, 2.0, 2.0} }

DotDimensionNumbers dnums;
dnums.add_lhs_contracting_dimensions(1);
dnums.add_rhs_contracting_dimensions(1);

DotGeneral(lhs, rhs, dnums) -> { {6.0, 12.0},
{15.0, 30.0} }

lhs 和 rhs 中关联的批量维度编号必须具有相同的维度大小。

带有批次维度编号（批次大小 2，2x2 矩阵）的示例：

lhs = { { {1.0, 2.0},
{3.0, 4.0} },
{ {5.0, 6.0},
{7.0, 8.0} } }

rhs = { { {1.0, 0.0},
{0.0, 1.0} },
{ {1.0, 0.0},
{0.0, 1.0} } }

DotDimensionNumbers dnums;
dnums.add_lhs_contracting_dimensions(2);
dnums.add_rhs_contracting_dimensions(1);
dnums.add_lhs_batch_dimensions(0);
dnums.add_rhs_batch_dimensions(0);

DotGeneral(lhs, rhs, dnums) -> { { {1.0, 2.0},
{3.0, 4.0} },
{ {5.0, 6.0},
{7.0, 8.0} } }

输入	输出	语义
[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

这样，生成的维度编号从批次维度开始，接着是 lhs 非合约/非批处理维度，最后是 rhs 非合约/非批处理维度。

DynamicSlice

另请参阅 XlaBuilder::DynamicSlice。

DynamicSlice 从位于动态 start_indices 的输入数组中提取子数组。每个维度中的切片的大小在 size_indices 中传递，它指定每个维度中不含切片间隔的终点：[起始值、起始值 + 大小]。start_indices 的形状必须为 == 1，并且尺寸大小等于 operand 的阶。

DynamicSlice(operand, start_indices, size_indices)

参数	类型	语义
operand	XlaOp	T 类型的 N 维数组
start_indices	N XlaOp 序列	包含 N 个标量整数的列表，其中包含每个维度切片的起始索引。值必须大于或等于零。
size_indices	ArraySlice<int64>	包含 N 个整数的列表，其中包含每个维度的切片大小。每个值必须严格大于零，并且起始值 + 大小必须小于或等于维度的大小，以避免模数维度的大小被换行。

在执行切片之前，通过对 [1, N) 中的每个索引 i 应用以下转换来计算有效切片索引：

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

这样可以确保提取的切片始终在运算数数组的边界内。如果切片在应用转换之前在边界内，则转换不会产生任何影响。

一维示例：

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

DynamicSlice(a, s, {2}) produces:
{2.0, 3.0}

二维示例：

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 s = {2, 1}

DynamicSlice(b, s, {2, 2}) produces:
{ { 7.0,  8.0},
{10.0, 11.0} }

DynamicUpdateSlice

另请参阅 XlaBuilder::DynamicUpdateSlice。

DynamicUpdateSlice 会生成一个结果，该值为输入数组 operand 的值，并在 start_indices 处覆盖了 Slice update。update 的形状决定了所更新结果的子数组的形状。 start_indices 的形状必须为 rank == 1，并且维度大小等于 operand 的秩。

DynamicUpdateSlice(operand, update, start_indices)

参数	类型	语义
operand	XlaOp	T 类型的 N 维数组
update	XlaOp	包含切片更新的类型为 T 的 N 维数组。更新形状的每个维度都必须严格大于零，且起始值 + 更新值必须小于或等于每个维度的操作数大小，以免生成出界更新索引。
start_indices	N XlaOp 序列	包含 N 个标量整数的列表，其中包含每个维度切片的起始索引。值必须大于或等于零。

在执行切片之前，通过对 [1, N) 中的每个索引 i 应用以下转换来计算有效切片索引：

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

这样可以确保更新后的切片始终在运算数数组的边界内。如果切片在应用转换之前在边界内，则转换不会产生任何影响。

一维示例：

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) produces:
{0.0, 1.0, 5.0, 6.0, 4.0}

二维示例：

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) produces:
{ {0.0,  1.0,  2.0},
{3.0, 12.0, 13.0},
{6.0, 14.0, 15.0},
{9.0, 16.0, 17.0} }

元素级二元算术运算

另请参阅 XlaBuilder::Add。

支持一组元素级二元算术运算。

Op(lhs, rhs)

其中，Op 是 Add（加法）、Sub（减法）、Mul（乘法）、Div（除法）、Rem（余数）、Max（最大值）、Min（最小）、LogicalAnd（逻辑 AND）或 LogicalOr（逻辑 OR）之一。

参数	类型	语义
lhs	XlaOp	左侧操作数：T 类型的数组
rhs	XlaOp	右侧操作数：T 类型的数组

实参的形状必须相似或兼容。如需了解形状兼容的含义，请参阅广播文档。运算结果具有形状，这是广播两个输入数组的结果。在此变体中，不支持不同秩的数组之间的运算，除非其中一个运算数是标量。

当 Op 为 Rem 时，结果的符号取自被除数，并且结果的绝对值始终小于除数的绝对值。

整数除法溢出（有符号/无符号除数/余数除以 0 或 INT_SMIN 的有符号除数/余数 -1）会生成实现定义的值。

对于以下操作，存在支持不同排名广播的替代变体：

Op(lhs, rhs, broadcast_dimensions)

其中，Op 与上面相同。这种运算的变体应该用于不同秩的数组之间的算术运算（例如，向向量添加矩阵）。

额外的 broadcast_dimensions 运算数是一片整数，用于将低阶运算数的排名扩展至高阶运算数的排名。broadcast_dimensions 会将较低阶形状的尺寸映射到较高阶形状的尺寸。展开形状的未映射尺寸会填充为尺寸 1。然后，退化维度广播会沿着这些退化维度广播形状，使两个运算数的形状相等。广播页面详细介绍了这些语义。

元素级比较操作

另请参阅 XlaBuilder::Eq。

支持一组标准元素级二元比较运算。请注意，比较浮点类型时，标准 IEEE 754 浮点比较语义适用。

Op(lhs, rhs)

其中，Op 是 Eq（等于）、Ne（不等于）、Ge（大于或等于）、Gt（大于）、Le（小于或等于）、Lt（小于）之一。另一组运算符 EqTotalOrder、NeTotalOrder、GeTotalOrder、GtTotalOrder、LeTotalOrder 和 LtTotalOrder 提供相同的功能，只不过它们还通过强制执行 -NaN < -Inf < -Finite < -0 < +FNe < init 来支持对浮点数进行总排序。

参数	类型	语义
lhs	XlaOp	左侧操作数：T 类型的数组
rhs	XlaOp	右侧操作数：T 类型的数组

实参的形状必须相似或兼容。如需了解形状兼容的含义，请参阅广播文档。运算的结果具有形状，这是使用元素类型 PRED 广播两个输入数组的结果。在此变体中，不支持不同秩的数组之间的运算，除非其中一个运算数是标量。

对于以下操作，存在支持不同排名广播的替代变体：

Op(lhs, rhs, broadcast_dimensions)

其中，Op 与上面相同。这种运算的变体应该用于不同秩的数组之间的比较运算（例如，向向量添加矩阵）。

额外的 broadcast_dimensions 运算数是整数切片，指定了用于广播运算数的维度。广播页面中详细介绍了这些语义。

元素级一元函数

XlaBuilder 支持以下元素级一元函数：

Abs(operand) 元素级绝对 x -> |x|。

Ceil(operand) 元素级 ceil x -> ⌈x⌉。

Cos(operand) 元素级余弦 x -> cos(x)。

Exp(operand) 元素级自然指数 x -> e^x。

Floor(operand) 元素对应的楼层 x -> ⌊x⌋。

Imag(operand) 复杂（或实数）形状的元素级虚数部分。x -> imag(x)。如果操作数是浮点类型，则返回 0。

IsFinite(operand) 测试 operand 的每个元素是否有限（即非正无穷或负无穷且非 NaN）。当且仅当相应的输入元素为有限时，返回一个由 PRED 值组成的数组，其形状与输入相同，其中每个元素都是 true。

Log(operand) 元素级自然对数 x -> ln(x)。

LogicalNot(operand) 元素级逻辑，而非 x -> !(x)。

Logistic(operand) 元素级逻辑函数计算 x -> logistic(x)。

PopulationCount(operand) 计算 operand 的每个元素中设置的位数。

Neg(operand) 元素级否定 x -> -x。

Real(operand) 复杂（或实际）形状的元素级实部。 x -> real(x)。如果运算数是浮点类型，则返回相同的值。

Rsqrt(operand) 平方根运算 x -> 1.0 / sqrt(x) 的元素级倒数。

Sign(operand) 元素级符号运算 x -> sgn(x)，其中

\[\text{sgn}(x) = \begin{cases} -1 & x < 0\\ -0 & x = -0\\ NaN & x = NaN\\ +0 & x = +0\\ 1 & x > 0 \end{cases}\]

operand 元素类型的比较运算符。

Sqrt(operand) 元素级平方根运算 x -> sqrt(x)。

Cbrt(operand) 元素级立方根运算 x -> cbrt(x)。

Tanh(operand) 元素级双曲正切 x -> tanh(x)。

Round(operand) 元素级舍入，接近于 0。

RoundNearestEven(operand) 元素级舍入，尽可能接近于最接近的偶数。

参数	类型	语义
operand	XlaOp	函数的运算数

该函数会应用于 operand 数组中的每个元素，以生成形状相同的数组。operand 可以作为标量（0 阶）。

FFF

XLA FFT 运算针对实数和复数输入/输出实现正向和逆逆伏里转换。最多支持 3 个轴上的多维 FFT。

另请参阅 XlaBuilder::Fft。

参数	类型	语义
operand	XlaOp	我们要进行 Fourier 转换的数组。
fft_type	FftType	请参见下表。
fft_length	ArraySlice<int64>	要转换的轴的时间域长度。这对于 IRFFT 而言尤为需要调整最内轴的大小，因为 RFFT(fft_length=[16]) 的输出形状与 RFFT(fft_length=[17]) 相同。

FftType	语义
FFT	将复杂 FFT 转发到复杂 FFT。形状保持不变。
IFFT	将复数转换为复数 FFT。形状保持不变。
RFFT	转发实际到复杂 FFT。如果 fft_length[-1] 为非零值，则最内轴的形状缩减为 fft_length[-1] // 2 + 1，从而省略转换信号中超出奈奎斯特频率的反共产部分。
IRFFT	实数到复数的反函数（即取复数，返回实数）。如果 fft_length[-1] 为非零值，则最内轴的形状会展开为 fft_length[-1]，根据 1 到 fft_length[-1] // 2 + 1 条目的反向共通推断出转换后信号中超出奈奎斯特频率的部分。

多维 FFT

如果提供了多个 fft_length，则相当于将 FFT 运算的级联应用于每个最内层的轴。请注意，对于实际-复杂和复杂->实际情况，最内层的轴转换会首先执行（实际上是）执行（RFFT，最后是 IRFFT），因此最内轴的轴会改变大小。其他轴转换随后将变得复杂->复杂。

实现细节

CPU FFT 由 Eigen 的 TensorFFT 提供支持。GPU FFT 使用 cuFFT。

收集

XLA 收集操作将输入数组的多个切片拼接在一起（每个切片的运行时偏移量可能不同）。

一般语义

另请参阅 XlaBuilder::Gather。 有关更直观的说明，请参阅下面的“非正式说明”部分。

gather(operand, start_indices, offset_dims, collapsed_slice_dims, slice_sizes, start_index_map)

参数	类型	语义
operand	XlaOp	我们从中收集数据的数组。
start_indices	XlaOp	包含我们收集的切片的起始索引的数组。
index_vector_dim	int64	start_indices 中“包含”起始索引的维度。有关详细说明，请参阅下文。
offset_dims	ArraySlice<int64>	输出形状中偏移到从操作数切开的数组中的维度集。
slice_sizes	ArraySlice<int64>	slice_sizes[i] 是维度 i 的切片的边界。
collapsed_slice_dims	ArraySlice<int64>	每个切片中收起的维度集。这些维度的大小必须为 1。
start_index_map	ArraySlice<int64>	描述如何将 start_indices 中的索引映射到操作数的合法索引的映射。
indices_are_sorted	bool	调用方是否保证会对索引进行排序。

为方便起见，我们将输出数组中不在 offset_dims 中的维度标记为 batch_dims。

输出为 batch_dims.size + offset_dims.size 阶数组。

operand.rank 必须等于 offset_dims.size 和 collapsed_slice_dims.size 的总和。此外，slice_sizes.size 必须等于 operand.rank。

如果 index_vector_dim 等于 start_indices.rank，我们隐式地将 start_indices 视为具有尾随 1 维度（即，如果 start_indices 的形状为 [6,7]，index_vector_dim 为 2，那么我们会隐式地将 start_indices 的形状视为 [6,7,1]）。

输出数组沿维度 i 的边界计算方式如下：

1. 如果 batch_dims 中存在 i（即某个 k 等于 batch_dims[k]），我们就会从 start_indices.shape 中选取相应的尺寸边界，跳过 index_vector_dim（即，如果 k < index_vector_dim，则选择 start_indices.shape.dims[k]，否则选择 start_indices.shape.dims[k+1]）。

2. 如果 offset_dims 中存在 i（即某个 k 等于 offset_dims[k]），则在考虑 collapsed_slice_dims 后，我们从 slice_sizes 中选取相应的边界（即，我们选择 adjusted_slice_sizes[k]，其中 adjusted_slice_sizes 为 slice_sizes，索引 collapsed_slice_dims 处的边界已移除）。

正式地，与给定输出索引 Out 对应的运算数索引 In 的计算方法如下：

1. 让 G = { Out[k] for k in batch_dims }。使用 G 分割矢量 S，使 S[i] = start_indices[Combine(G, i)]，其中 Combine(A, b) 将 b 在位置 index_vector_dim 处插入 A。请注意，即使 G 为空，也是如此：如果 G 为空，则 S = start_indices。

2. 通过使用 S 使用 start_index_map 分散 S，在 operand 中创建起始索引 Sin。更确切地说：

   1. 如果 k < start_index_map.size，则 Sin[start_index_map[k]] = S[k]。

   2. 否则，Sin[_] = 0。

3. 创建一个索引，Oin以operand为基准，根据collapsed_slice_dims集合，按偏移量Out中的索引散布索引。更确切地说：

   1. 如果 k < offset_dims.size（remapped_offset_dims 定义如下），则 Oin[remapped_offset_dims(k)] = Out[offset_dims[k]]。

   2. 否则，Oin[_] = 0。

4. In 是 Oin + Sin，其中 + 表示元素级加法。

remapped_offset_dims 是单调函数，域为 [0, offset_dims.size)，范围为 [0, operand.rank) \ collapsed_slice_dims。举例来说，offset_dims.size 为 4，operand.rank 为 6，collapsed_slice_dims 为 {0, 2}，然后 remapped_offset_dims 为 {0→1, 1→3, 2→4, 3→5}。

如果 indices_are_sorted 设置为 true，则 XLA 可以假定用户已对 start_indices 进行排序（start_index_map 升序）。否则，语义是实现定义的。

非正式说明和示例

正式地说，输出数组中的每个索引 Out 对应于运算数数组中的一个元素 E，计算方式如下：

• 我们使用 Out 中的批次维度从 start_indices 查找起始索引。

• 我们使用 start_index_map 将起始索引（其大小可能小于 operand.rank）映射到 operand 中的“完整”起始索引。

• 我们使用完整的起始索引动态切片大小为 slice_sizes 的切片。

• 我们通过收起 collapsed_slice_dims 维度来重塑 Slice。由于所有收起的切片尺寸的边界必须为 1，因此这种重塑始终是合法的。

• 我们使用 Out 中的偏移维度来建立此 Slice 的索引，以获取与输出索引 Out 对应的输入元素 E。

在以下所有示例中，index_vector_dim 均设置为 start_indices.rank - 1。index_vector_dim 更有趣的值并不会从根本上改变操作，但会使视觉表示更繁琐。

为了直观地理解上述所有方法如何组合在一起，让我们看一个示例，该示例从 [16,11] 数组中收集了 5 个形状 [8,6] 切片。切片在 [16,11] 数组中的位置可以表示为形状为 S64[2] 的索引矢量，因此这 5 个位置的集合可以表示为 S64[5,2] 数组。

然后，收集操作的行为可以描述为索引转换，该转换获取 [G,O0,O1]（输出形状中的索引），并按以下方式将其映射到输入数组中的元素：

我们首先使用 G 从收集索引数组中选择一个 (X,Y) 向量。输出数组中位于索引 [G,O0,O1] 处的元素然后是输入数组中位于索引 [X+O0,Y+O1] 处的元素。

slice_sizes 为 [8,6]，用于确定 O₀ 和 O₁ 的范围，而这又决定了 Slice 的边界。

此收集操作充当以 G 作为批量维度的批量动态切片。

收集索引可能是多维的。例如，上述示例的更通用版本使用形状 [4,5,2] 的“gather indices”数组可以转换如下索引：

同样，这充当批量动态切片 G0，G1 作为批量维度。切片大小仍然为 [8,6]。

XLA 中的收集操作通过以下方式泛化上述非正式语义：

1. 我们可以将输出形状中的哪些维度配置为偏移维度（在上一个示例中，包含 O0、O1 的维度）。输出批量维度（上一个示例中包含 G0、G1 的维度）定义为不是偏移维度的输出维度。

2. 输出形状中明确存在的输出偏移维度数可能小于输入秩。这些“缺失”的维度（明确列为 collapsed_slice_dims）的切片大小必须为 1。由于它们的切片大小为 1，因此它们的唯一有效索引是 0，省略它们不会产生歧义。

3. 从“Gather Indices”数组（在上一示例中的 X、Y）中提取的 Slice 包含的元素可能少于输入数组排名，并且显式映射会指示应如何扩展索引，使其与输入具有相同的排名。

在最后一个示例中，我们使用 (2) 和 (3) 来实现 tf.gather_nd：

G0 和 G1 用于照常从收集索引数组中剥离起始索引，但起始索引只有一个元素 X。同样，只有一个输出偏移索引的值为 O0。不过，在将它们用作输入数组的索引之前，这些索引会根据“Gather Index Mapping”（正式说明中的 start_index_map）和“Offset Mapping”（在正式描述中的 remapped_offset_dims）分别扩展为 [X,0] 和 [0,O0]，从而为 [G1] [G] 提供 [G1]。X000000OOOG1GatherIndicestf.gather_nd

本案例的 slice_sizes 为 [1,11]。直观地说，这意味着收集索引数组中的每个索引 X 都会选择整行，并且结果就是所有这些行的串联。

GetDimensionSize

另请参阅 XlaBuilder::GetDimensionSize。

返回操作数指定维度的大小。运算数必须为数组形状。

GetDimensionSize(operand, dimension)

参数	类型	语义
operand	XlaOp	N 维输入数组
dimension	int64	间隔 [0, n) 中的值，用于指定维度

SetDimensionSize

另请参阅 XlaBuilder::SetDimensionSize。

设置 XlaOp 指定维度的动态大小。运算数必须为数组形状。

SetDimensionSize(operand, size, dimension)

参数	类型	语义
operand	XlaOp	n 维输入数组。
size	XlaOp	int32 表示运行时动态大小。
dimension	int64	间隔 [0, n) 中的值，用于指定维度。

传递运算数作为结果，编译器会跟踪动态维度。

下游缩减操作将忽略填充的值。

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);

GetTupleElement

另请参阅 XlaBuilder::GetTupleElement。

编入具有编译时常量值的元组的索引。

该值必须是编译时常量，以便形状推断可以确定所生成值的类型。

这类似于 C++ 中的 std::get<int N>(t)。从概念上来讲：

let v: f32[10] = f32[10]{0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
let s: s32 = 5;
let t: (f32[10], s32) = tuple(v, s);
let element_1: s32 = gettupleelement(t, 1);  // Inferred shape matches s32.

另请参阅 tf.tuple。

信息流广告

另请参阅 XlaBuilder::Infeed。

Infeed(shape)

参数名	类型	语义
shape	Shape	从信息流界面读取的数据的形状。形状的布局字段必须设置为匹配发送到设备的数据的布局；否则其行为将处于未定义状态。

从设备的隐式信息流流式接口中读取单个数据项，将数据解释为给定形状及其布局，然后返回数据的 XlaOp。一次计算中允许执行多个信息流操作，但信息流操作之间必须有一个总顺序。例如，以下代码中的两个 Infeed 具有总顺序，因为 while 循环之间存在依赖关系。

result1 = while (condition, init = init_value) {
  Infeed(shape)
}

result2 = while (condition, init = result1) {
  Infeed(shape)
}

不支持嵌套元组形状。对于空元组形状，馈入操作实际上是一个空操作，并且无需从设备馈送中读取任何数据即可继续进行。

Iota

另请参阅 XlaBuilder::Iota。

Iota(shape, iota_dimension)

在设备上构建常量字面量，而不是潜在的大型主机传输。用于创建具有指定形状的数组，并存储从零开始并按指定维度递增 1 的值。对于浮点类型，生成的数组等同于 ConvertElementType(Iota(...))，其中 Iota 是整数类型，而转换是浮点类型。

参数	类型	语义
shape	Shape	Iota() 创建的数组的形状
iota_dimension	int64	递增的维度。

例如，Iota(s32[4, 8], 0) 会返回

  [[0, 0, 0, 0, 0, 0, 0, 0 ],
   [1, 1, 1, 1, 1, 1, 1, 1 ],
   [2, 2, 2, 2, 2, 2, 2, 2 ],
   [3, 3, 3, 3, 3, 3, 3, 3 ]]

退货金额 Iota(s32[4, 8], 1)

  [[0, 1, 2, 3, 4, 5, 6, 7 ],
   [0, 1, 2, 3, 4, 5, 6, 7 ],
   [0, 1, 2, 3, 4, 5, 6, 7 ],
   [0, 1, 2, 3, 4, 5, 6, 7 ]]

地图

另请参阅 XlaBuilder::Map。

Map(operands..., computation)

参数	类型	语义
operands	N XlaOp 序列	N 个类型为 T0..T{N-1} 的数组
computation	XlaComputation	T_0, T_1, .., T_{N + M -1} -> S 类型的计算，其中有 N 个类型为 T 的参数和任意类型的 M M
dimensions	int64 数组	地图尺寸数组

对给定的 operands 数组应用标量函数，从而生成一个具有相同维度的数组，其中每个元素都是应用于输入数组中相应元素的映射函数的结果。

映射函数是任意计算，但有 N 个标量类型为 T 的输入和单个类型为 S 的输出。除了元素类型 T 被 S 替换之外，输出结果的维度与运算数相同。

例如：Map(op1, op2, op3, computation, par1) 映射输入数组中每个（多维）索引处的 elem_out <- computation(elem1, elem2, elem3, par1)，以生成输出数组。

OptimizationBarrier

阻止任何优化传递跨屏移动计算。

确保先评估所有输入，然后再评估依赖于屏障输出的任何运算符。

棉垫

另请参阅 XlaBuilder::Pad。

Pad(operand, padding_value, padding_config)

参数	类型	语义
operand	XlaOp	类型为 T 的数组
padding_value	XlaOp	T 类型的标量，用于填充添加的内边距
padding_config	PaddingConfig	两侧（低、高）以及各尺寸元素之间的内边距

通过围绕该数组以及具有指定 padding_value 的数组元素之间的内边距来扩展指定的 operand 数组。padding_config 指定每个维度的边缘内边距和内部内边距大小。

PaddingConfig 是 PaddingConfigDimension 的重复字段，其中包含每个维度的三个字段：edge_padding_low、edge_padding_high 和 interior_padding。

edge_padding_low 和 edge_padding_high 分别指定在每个维度的低端（索引 0 旁边）和高端（最高索引旁）处添加的内边距大小。边缘内边距的量可为负；负内边距的绝对值表示要从指定维度移除的元素数量。

interior_padding 指定每个维度中任意两个元素之间添加的内边距；此值不能为负数。内部内边距在逻辑上发生在边缘内边距之前，因此如果是负边缘内边距，元素会从内内边距的运算数中移除。

如果边缘内边距对全部为 (0, 0) 且内部内边距值全部为 0，则此操作为空操作。下图显示了二维数组的不同 edge_padding 和 interior_padding 值的示例。

接收

另请参阅 XlaBuilder::Recv。

Recv(shape, channel_handle)

参数	类型	语义
shape	Shape	要接收的数据的
channel_handle	ChannelHandle	每个发送/接收对的唯一标识符

在共享同一通道句柄的其他计算中，从 Send 指令接收给定形状的数据。针对已接收的数据返回 XlaOp。

Recv 操作的客户端 API 表示同步通信。不过，该指令会在内部分解为 2 条 HLO 指令（Recv 和 RecvDone），以启用异步数据传输。另请参阅 HloInstruction::CreateRecv 和 HloInstruction::CreateRecvDone。

Recv(const Shape& shape, int64 channel_id)

分配从具有相同 channel_id 的 Send 指令接收数据所需的资源。返回已分配资源的上下文，以下 RecvDone 指令使用该上下文来等待数据传输完成。上下文是 {receivebuffer (shape), request identifier (U32)} 元组，只能由 RecvDone 指令使用。

RecvDone(HloInstruction context)

在给定由 Recv 指令创建的上下文的情况下，系统会等待数据传输完成并返回收到的数据。

减少

另请参阅 XlaBuilder::Reduce。

将归约函数并行应用于一个或多个数组。

Reduce(operands..., init_values..., computation, dimensions)

参数	类型	语义
operands	N XlaOp 序列	N 个类型为 T_0, ..., T_{N-1} 的数组。
init_values	N XlaOp 序列	N 个 T_0, ..., T_{N-1} 类型的标量。
computation	XlaComputation	类型为 T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}) 的计算。
dimensions	int64 数组	要缩减的无序维度数组。

其中：

• N 必须大于或等于 1。
• 计算必须是“粗略”关联（见下文）。
• 所有输入数组必须具有相同的维度。
• 所有初始值都必须在 computation 下形成一个恒等式。
• 如果 N = 1，则 Collate(T) 为 T。
• 如果为 N > 1，则 Collate(T_0, ..., T_{N-1}) 是类型为 T 的 N 元素的元组。

此操作会将每个输入数组的一个或多个维度缩减为标量。每个返回数组的排名为 rank(operand) - len(dimensions)。该操作的输出是 Collate(Q_0, ..., Q_N)，其中 Q_i 是 T_i 类型的数组，其维度如下所述。

允许不同的后端重新关联归约计算。这可能会导致数值差异，因为某些归约函数（例如加法）与浮点数无关。 但是，如果数据的范围有限，则浮点加法和大多数实际用途都足够接近关联。

示例

使用值 [10, 11, 12, 13] 跨单个一维数组的一个维度进行缩减，使用归约函数 f（即 computation）时，可以按以下方式进行计算：

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

但也有许多其他可能性，例如

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

以下是一个粗略的伪代码示例，展示如何实现归约，即使用求和作为初始值为 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 rank 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]

下面是一个缩减二维数组（矩阵）的示例。该形状的秩为 2，维度 0 的大小为 2，维度 1 的大小为 3：

使用“add”函数缩减维度 0 或 1 的结果：

请注意，两个归约结果都是一维数组。为方便起见，该图将一个图表显示为列，另一个显示为行。

如需更复杂的示例，请参考 3D 数组。其秩为 3，维度 0 的大小为 4，维度 1 的大小为 2，维度 2 的大小为 3。为简单起见，值 1 到 6 会复制到维度 0 中。

与 2D 示例类似，我们也可以仅缩减一个维度。例如，如果我们缩减维度 0，就会得到一个 2 阶数组，其中维度 0 中的所有值都已折叠成一个标量：

|  4   8  12 |
| 16  20  24 |

如果我们缩减维度 2，则会得到一个 rank-2 数组，其中维度 2 中的所有值都会折叠成一个标量：

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

请注意，输入中其余维度之间的相对顺序会在输出中保留，但某些维度可能会获得新数字（由于排名发生变化）。

我们还可以减少多个维度。将维度 0 和 1 相加减法会生成一维数组 [20, 28, 36]。

缩减 3D 数组的所有维度会生成标量 84。

Variadic 缩减

当为 N > 1 时，Reduce 函数应用稍微复杂一些，因为它同时应用于所有输入。运算数按以下顺序提供给计算：

• 正在运行第一个运算数的降值
• ...
• 正在运行第 N 个操作数的降值
• 第一个运算数的输入值
• ...
• 第 N 个操作数的输入值

以下面的归约函数为例，该函数可用于并行计算一维数组的最大值和 argmax 值：

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

对于一维输入数组 V = Float[N], K = Int[N] 和 init 值 I_V = Float, I_K = Int，在唯一输入维度中缩减的结果 f_(N-1) 相当于以下递归应用：

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

将这种归约应用于值数组和顺序索引数组（即 iota），将共同迭代这些数组，并返回包含最大值和匹配索引的元组。

ReducePrecision

另请参阅 XlaBuilder::ReducePrecision。

对将浮点值转换为低精度格式（例如 IEEE-FP16）并转换回原始格式的效果建模。可以任意指定低精度格式的指数和尾数位数，不过并非所有硬件实现都支持所有位大小。

ReducePrecision(operand, mantissa_bits, exponent_bits)

参数	类型	语义
operand	XlaOp	浮点类型 T 的数组。
exponent_bits	int32	低精度格式的指数位数
mantissa_bits	int32	低精度格式的尾数位数

返回的结果是 T 类型的数组。输入值会四舍五入为可用指定尾数位表示的最接近值（通过“与偶数位计算”语义），任何超过指数位数指定范围的值都会被限制为正无穷大或负无穷大。NaN 值会被保留，但可能会转换为规范的 NaN 值。

低精度格式必须至少有一个指数位（以便区分零值与无穷大，因为两者的尾数为零），并且必须具有非负数的尾数位。指数或尾数位数可能会超过 T 类型的对应值；因此，相应的转换部分只是一个空操作。

ReduceScatter

另请参阅 XlaBuilder::ReduceScatter。

ReduceScatter 是一项集体操作，可有效执行 AllReduce，然后通过 scatter_dimension 将结果拆分为 shard_count 块，而副本组中的副本 i 会收到 ith 分片。

ReduceScatter(operand, computation, scatter_dim, shard_count, replica_group_ids, channel_id)

参数	类型	语义
operand	XlaOp	要跨副本减少的数组或非空数组元组。
computation	XlaComputation	归约计算
scatter_dimension	int64	要进行散布的维度。
shard_count	int64	要拆分的块数 scatter_dimension
replica_groups	int64 的向量	执行缩减的组
channel_id	可选int64	用于跨模块通信的可选通道 ID

• 当 operand 是数组元组时，对元组的每个元素执行归约散点操作。
• replica_groups 是执行缩减操作的副本组的列表（可以使用 ReplicaId 检索当前副本的副本 ID）。每个组中副本的顺序决定了 all-reduce 结果的分散顺序。replica_groups 必须为空（在这种情况下，所有副本都属于一个组），或者包含与副本数量相同的元素数量。如果存在多个副本组，则这些副本组的大小必须相同。例如，replica_groups = {0, 2}, {1, 3} 在副本 0 和 2 以及 1 和 3 之间执行缩减，然后对结果进行散布。
• shard_count 是每个副本组的大小。在 replica_groups 为空的情况下，我们需要使用该属性。如果 replica_groups 不为空，则 shard_count 必须等于每个副本组的大小。
• channel_id 用于跨模块通信：只有具有相同 channel_id 的 reduce-scatter 操作可以相互通信。

输出形状是 scatter_dimension 缩小了 shard_count 倍的输入形状。例如，如果有两个副本，并且操作数在两个副本上分别具有值 [1.0, 2.25] 和 [3.0, 5.25]，则对于第一个副本，此操作的 scatter_dim 为 0 的输出值将为 [4.0]，对于第二个副本，则值为 [7.5]。

ReduceWindow

另请参阅 XlaBuilder::ReduceWindow。

对一系列 N 个多维数组的每个窗口中的所有元素应用归约函数，从而生成 N 个多维数组的单个或一个元组作为输出。每个输出数组具有与窗口的有效位置数量相同数量的元素。池化层可以表示为 ReduceWindow。与 Reduce 类似，所应用的 computation 始终在左侧传递 init_values。

ReduceWindow(operands..., init_values..., computation, window_dimensions, window_strides, padding)

参数	类型	语义
operands	N XlaOps	一系列 T_0,..., T_{N-1} 类型的多维数组，每个数组表示放置窗口的基本区域。
init_values	N XlaOps	N 个归约起始值，N 个运算数各对应一个。如需了解详情，请参阅 Reduce。
computation	XlaComputation	T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}) 类型的归约函数，应用于所有输入运算数的每个窗口中的元素。
window_dimensions	ArraySlice<int64>	窗口维度值的整数数组
window_strides	ArraySlice<int64>	窗口步长值的整数数组
base_dilations	ArraySlice<int64>	基本膨胀值的整数数组
window_dilations	ArraySlice<int64>	窗口膨胀值的整数数组
padding	Padding	窗口的内边距类型（Padding::kSame，如果步长为 1，则会填充，以便输出形状与输入相同；或者 Padding::kValid ：不使用内边距，并在窗口不再适合时“停止”窗口）

其中：

• N 必须大于或等于 1。
• 所有输入数组必须具有相同的维度。
• 如果 N = 1，则 Collate(T) 为 T。
• 如果为 N > 1，则 Collate(T_0, ..., T_{N-1}) 是类型为 (T0,...T{N-1}) 的 N 元素的元组。

下面的代码和图显示了使用 ReduceWindow 的示例。输入大小为 [4x6] 的矩阵，window_dimensions 和 window_stride_dimensions 均为 [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);

某一维度的步长 1 表示窗口在该维度中的位置距离相邻窗口 1 个元素。为了指定没有窗口相互重叠，window_stride_dimensions 应等于 window_dimensions。下图说明了如何使用两个不同的步幅值。内边距会应用于输入的每个维度，计算结果与输入参数在填充后具有的维度相同。

举一个非常重要的填充示例，考虑使用维度 3 和步长 2 对输入数组 [10000, 1000, 100, 10, 1] 计算缩减窗口最小值（初始值为 MAX_FLOAT）。内边距 kValid 会在两个有效窗口（[10000, 1000, 100] 和 [100, 10, 1]）上计算最小值，结果为输出 [100, 1]。先用内边距 kSame 填充数组，使缩小窗口后的形状与步幅 1 的输入相同（通过在两边添加初始元素，得到 [MAX_VALUE, 10000, 1000, 100, 10, 1, MAX_VALUE]）。对含内边距的数组运行缩减窗口操作可对三个窗口（[MAX_VALUE, 10000, 1000]、[1000, 100, 10]、[10, 1, MAX_VALUE] 和收益 [1000, 10, 1]）执行操作。

归约函数的计算顺序是任意的，并且可能是不确定的。因此，归约函数不应对重新关联过于敏感。如需了解详情，请参阅 Reduce 上下文中有关关联的讨论。

ReplicaId

另请参阅 XlaBuilder::ReplicaId。

返回副本的唯一 ID（U32 标量）。

ReplicaId()

每个副本的唯一 ID 是间隔 [0, N) 内的一个无符号整数，其中 N 表示副本数量。由于所有副本都运行同一程序，因此程序中的 ReplicaId() 调用将对每个副本返回不同的值。

调整形状

另请参阅 XlaBuilder::Reshape 和 Collapse 运算。

将数组的维度重塑为新配置。

Reshape(operand, new_sizes) Reshape(operand, dimensions, new_sizes)

参数	类型	语义
operand	XlaOp	类型为 T 的数组
dimensions	int64矢量	维度的收起顺序
new_sizes	int64矢量	新维度大小向量

从概念上讲，首先需要将数组扁平化为一维数据值的一维向量，然后再优化此向量，使其成为新的形状。输入参数是类型为 T 的任意数组、维度索引的编译时常量向量以及结果维度大小的编译时常量向量。dimension 矢量中的值（如果指定）必须是 T 的所有维度的排列，否则默认为 {0, ..., rank - 1}。dimensions 中维度的顺序是从循环嵌套（将输入数组收起为单个维度）中变化最慢的维度（最主要的）到变化最快的维度（最次要的）。new_sizes 矢量决定了输出数组的大小。new_sizes 中索引 0 处的值是维度 0 的大小，索引 1 处的值是维度 1 的大小，以此类推。new_size 维度的乘积必须等于运算数维度大小的乘积。将收起的数组优化为由 new_sizes 定义的多维数组时，new_sizes 中的维度将按照变化最慢（变化最严重）到变化最快（变化最小）的顺序排列。

例如，假设 v 是一个包含 24 个元素的数组：

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

In-order collapse:
let v012_24 = Reshape(v, {0,1,2}, {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, {0,1,2}, {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} };

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

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


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

一种特殊情况是，重塑可将单元素数组转换为标量，反之亦然。例如，

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

还原（反向）

另请参阅 XlaBuilder::Rev。

Rev(operand, dimensions)

参数	类型	语义
operand	XlaOp	类型为 T 的数组
dimensions	ArraySlice<int64>	可反转的维度

反转 operand 数组中元素的顺序沿指定的 dimensions 方向，从而生成形状相同的输出数组。多维索引处的运算数数组的每个元素都存储在经过转换的索引处的输出数组中。我们通过反转每个维度中的索引来转换多维索引（即，如果大小为 N 的维度是反向维度之一，则其索引 i 将转换为 N - 1 - i）。

Rev 运算的一个用途是在神经网络中的梯度计算期间沿两个窗口维度反转卷积权重数组。

RngNormal

另请参阅 XlaBuilder::RngNormal。

使用遵循 \(N(\mu, \sigma)\) 正态分布生成的随机数字构造给定形状的输出。参数 \(\mu\) 、 \(\sigma\)以及输出形状必须具有浮点元素类型。这些参数必须进一步指定标量值。

RngNormal(mu, sigma, shape)

参数	类型	语义
mu	XlaOp	T 类型的标量，指定所生成数字的平均值
sigma	XlaOp	T 类型的标量，用于指定生成的
shape	Shape	T 类型的输出形状

RngUniform

另请参阅 XlaBuilder::RngUniform。

使用间隔 \([a,b)\)内均匀分布生成的随机数字构造给定形状的输出。参数和输出元素类型必须是布尔值类型、整数类型或浮点类型，并且这些类型必须一致。CPU 和 GPU 后端目前仅支持 F64、F32、F16、BF16、S64、U64、S32 和 U32。此外，这些参数需要标量值。如果为 \(b <= a\) ，则结果是由实现定义的。

RngUniform(a, b, shape)

参数	类型	语义
a	XlaOp	类型 T 的标量，用于指定间隔下限
b	XlaOp	类型为 T 的标量，用于指定间隔上限
shape	Shape	T 类型的输出形状

RngBitGenerator

使用指定算法（或后端默认值）生成给定形状（填充了均匀随机位）的输出，并返回更新后的状态（形状与初始状态相同）和生成的随机数据。

初始状态是当前随机数生成的初始状态。形状以及所需的形状和有效值取决于所使用的算法。

输出一定是初始状态的确定性函数，但不保证在后端和不同编译器版本之间是确定性的。

RngBitGenerator(algorithm, key, shape)

参数	类型	语义
algorithm	RandomAlgorithm	要使用的 PRNG 算法。
initial_state	XlaOp	PRNG 算法的初始状态。
shape	Shape	所生成数据的输出形状。

algorithm 的可用值：

• rng_default：具有后端专用形状要求的后端专用算法。

• rng_three_fry：基于 ThreeFry 计数器的 PRNG 算法。initial_state 形状为 u64[2]，具有任意值。Salmon 等，SC 2011。平行随机数字：只需输入 1、2、3 等简单数字即可。

• rng_philox：用于并行生成随机数字的 Philox 算法。initial_state 形状为 u64[3]，具有任意值。Salmon 等，SC 2011。平行随机数字：只需输入 1、2、3 等简单数字即可。

散点图

XLA 散点操作生成一系列结果，这些结果是输入数组 operands 的值，其中几个切片（位于 scatter_indices 指定的索引处）使用 update_computation 使用 updates 中的值序列进行更新。

另请参阅 XlaBuilder::Scatter。

scatter(operands..., scatter_indices, updates..., update_computation, index_vector_dim, update_window_dims, inserted_window_dims, scatter_dims_to_operand_dims)

参数	类型	语义
operands	N XlaOp 序列	N 个类型为 T_0, ..., T_N 的数组。
scatter_indices	XlaOp	包含必须分布的切片的起始索引的数组。
updates	N XlaOp 序列	N 个类型为 T_0, ..., T_N 的数组。updates[i] 包含用于分散 operands[i] 的值。
update_computation	XlaComputation	用于组合输入数组中的现有值和散点期间的更新。此计算的类型应为 T_0, ..., T_N, T_0, ..., T_N -> Collate(T_0, ..., T_N)。
index_vector_dim	int64	scatter_indices 中包含起始索引的维度。
update_window_dims	ArraySlice<int64>	updates 形状中的一组尺寸，即窗口尺寸。
inserted_window_dims	ArraySlice<int64>	必须插入 updates 形状的一组窗口尺寸。
scatter_dims_to_operand_dims	ArraySlice<int64>	从散点索引到运算数索引空间的维度映射。此数组解释为将 i 映射到 scatter_dims_to_operand_dims[i]。必须是一对一的，综合的。
indices_are_sorted	bool	调用方是否保证会对索引进行排序。
unique_indices	bool	调用方是否保证索引是唯一的。

其中：

• N 必须大于或等于 1。
• operands[0]、...、operands[N-1] 必须都具有相同的尺寸。
• updates[0]、...、updates[N-1] 必须都具有相同的尺寸。
• 如果 N = 1，则 Collate(T) 为 T。
• 如果为 N > 1，则 Collate(T_0, ..., T_N) 是 T 类型的 N 元素的元组。

如果 index_vector_dim 等于 scatter_indices.rank，则我们会隐式将 scatter_indices 视为具有尾随 1 维度。

我们将 ArraySlice<int64> 类型的 update_scatter_dims 定义为 updates 形状中不在 update_window_dims 中的一组维度（按升序排列）。

scatter 的参数应遵循以下约束条件：

• 每个 updates 数组的排名必须为 update_window_dims.size + scatter_indices.rank - 1。

• 每个 updates 数组中维度 i 的边界必须符合以下要求：

  • 如果 update_window_dims 中存在 i（即某些 k 等于 update_window_dims[k]），则在考虑 inserted_window_dims 后，updates 维度 i 的边界不得超出 operand 的相应边界（即 adjusted_window_bounds[k]，其中 adjusted_window_bounds 包含 operand 的边界，但索引 inserted_window_dims 处的边界已移除）。
  • 如果 update_scatter_dims 中存在 i（即，对于某些 k，维度 i 等于 update_scatter_dims[k]），则 updates 中维度 i 的边界必须等于 scatter_indices 的相应边界，并跳过 index_vector_dim（即，如果 k < index_vector_dim，则跳过 scatter_indices.shape.dims[k]，否则为 scatter_indices.shape.dims[k+1]）。

• update_window_dims 必须按升序排列，不得包含任何重复的维度数字，且在 [0, updates.rank) 范围内。

• inserted_window_dims 必须按升序排列，不得包含任何重复的维度数字，且在 [0, operand.rank) 范围内。

• operand.rank 必须等于 update_window_dims.size 和 inserted_window_dims.size 的总和。

• scatter_dims_to_operand_dims.size 必须等于 scatter_indices.shape.dims[index_vector_dim]，且其值必须在 [0, operand.rank) 范围内。

对于每个 updates 数组中的给定索引 U，必须应用此更新的相应 operands 数组中的相应索引 I 的计算方式如下：

1. 让 G = { U[k] for k in update_scatter_dims }。使用 G 在 scatter_indices 数组中查找索引向量 S，使 S[i] = scatter_indices[Combine(G, i)]，其中 Combine(A, b) 在 index_vector_dim 位置处将 b 插入 A。
2. 通过使用 S 使用 scatter_dims_to_operand_dims 映射分散 S，在 operand 中创建索引 Sin。更正式地说：
   1. 如果 k < scatter_dims_to_operand_dims.size，则 Sin[scatter_dims_to_operand_dims[k]] = S[k]。
   2. 否则，Sin[_] = 0。
3. 根据 inserted_window_dims，为每个 operands 数组创建一个索引 Win，方法是根据 inserted_window_dims，将索引散布于 U 中的 update_window_dims。更正式地说：
   1. Win[window_dims_to_operand_dims(k)] = U[k]（如果 k 位于 update_window_dims 中，其中 window_dims_to_operand_dims 是单调函数，区间为 [0, update_window_dims.size) 和范围 [0, operand.rank) \ inserted_window_dims。（例如，如果 update_window_dims.size 为 4，operand.rank 为 6，inserted_window_dims 为 {0, 2}，则 window_dims_to_operand_dims 为 {0→1, 1→3, 2→4, 3→5}）。
   2. 否则，Win[_] = 0。
4. I 是 Win + Sin，其中 + 表示元素级加法。

总的来说，散点运算可以定义如下。

• 使用 operands 初始化 output，即对于 operands[J] 数组中的所有索引 O，对所有索引 J 进行初始化：
  output[J][O] = operands[J][O]
• 针对 updates[J] 数组中的每个索引 U 和 operand[J] 数组中的相应索引 O（如果 O 是 output 的有效索引）：
  (output[0][O], ..., output[N-1][O]) =update_computation(output[0][O], ..., ,output[N-1][O],updates[0][U], ...,updates[N-1][U])

应用更新的顺序是不确定的。因此，当 updates 中的多个索引引用 operands 中的同一索引时，output 中的相应值将是不确定的值。

请注意，传递到 update_computation 的第一个参数始终是 output 数组中的当前值，第二个参数始终是 updates 数组中的值。这对于 update_computation 不具有交换性的情况尤其重要。

如果 indices_are_sorted 设置为 true，则 XLA 可以假定用户已对 start_indices 进行排序（start_index_map 升序）。否则，语义是实现定义的。

如果 unique_indices 设置为 true，则 XLA 可以假定所有分散的元素都是唯一的。因此 XLA 可以使用非原子操作。如果将 unique_indices 设置为 true 并且分布到的索引不是唯一的，则表示语义是实现定义的。

通俗来说，散点操作可以视为收集操作的反函数，也就是说，散点操作会更新输入中由相应收集操作提取的元素。

如需查看详细的非正式说明和示例，请参阅 Gather 下的“非正式说明”部分。

选择

另请参阅 XlaBuilder::Select。

根据谓词数组的值，用两个输入数组的元素构造输出数组。

Select(pred, on_true, on_false)

参数	类型	语义
pred	XlaOp	PRED 类型的数组
on_true	XlaOp	类型为 T 的数组
on_false	XlaOp	类型为 T 的数组

数组 on_true 和 on_false 必须具有相同的形状。这也是输出数组的形状。数组 pred 的维度必须与 on_true 和 on_false 相同，但具有 PRED 元素类型。

对于 pred 的每个元素 P，如果 P 的值为 true，则从 on_true 获取输出数组的相应元素；如果 P 的值为 false，则从 on_false 中获取。作为受限形式的广播，pred 可以是 PRED 类型的标量。在这种情况下，如果 pred 为 true，则输出数组全部从 on_true 获取；如果 pred 为 false，则从 on_false 中获取。

使用非标量 pred 的示例：

let pred: PRED[4] = {true, false, false, true};
let v1: s32[4] = {1, 2, 3, 4};
let v2: s32[4] = {100, 200, 300, 400};
==>
Select(pred, v1, v2) = s32[4]{1, 200, 300, 4};

带有标量 pred 的示例：

let pred: PRED = true;
let v1: s32[4] = {1, 2, 3, 4};
let v2: s32[4] = {100, 200, 300, 400};
==>
Select(pred, v1, v2) = s32[4]{1, 2, 3, 4};

支持在元组之间进行选择。为此，元组被视为标量类型。如果 on_true 和 on_false 是元组（必须具有相同的形状！），则 pred 必须是 PRED 类型的标量。

SelectAndScatter

另请参阅 XlaBuilder::SelectAndScatter。

此操作可以视为一种复合操作，首先对 operand 数组计算 ReduceWindow 以从每个窗口中选择一个元素，然后将 source 数组散布到所选元素的索引，以构建形状与运算数数组相同的输出数组。二元 select 函数用于通过将某个元素应用到每个窗口来从每个窗口中选择该元素，并在调用该函数时具有如下属性：在字典顺序上，第一个参数的索引向量小于第二个参数的索引向量。如果选择第一个参数，则 select 函数返回 true；如果选择了第二个参数，则返回 false，并且该函数必须保持传递性（即，如果 select(a, b) 和 select(b, c) 为 true，则 select(a, c) 也为 true），以便所选元素不依赖于针对给定窗口遍历元素的顺序。

函数 scatter 应用于输出数组中的每个选定索引。它接受两个标量参数：

1. 输出数组中所选索引处的当前值
2. source 中应用于所选索引的散点值

它会组合这两个参数并返回一个标量值，该值用于更新输出数组中所选索引处的值。最初，输出数组的所有索引都设置为 init_value。

输出数组的形状必须与 operand 数组相同，而 source 数组的形状必须与对 operand 数组应用 ReduceWindow 操作所得到的形状相同。SelectAndScatter 可用于反向传播神经网络中池化层的梯度值。

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

参数	类型	语义
operand	XlaOp	窗口滑动所在的数组，类型为 T
select	XlaComputation	T, T -> PRED 类型的二元计算，应用于每个窗口中的所有元素；如果选择了第一个参数，则返回 true；如果选择了第二个参数，则返回 false
window_dimensions	ArraySlice<int64>	窗口维度值的整数数组
window_strides	ArraySlice<int64>	窗口步长值的整数数组
padding	Padding	窗口的内边距类型（Padding::kSame 或 Padding::kValid）
source	XlaOp	类型 T 的数组，其中包含要分散的值
init_value	XlaOp	输出数组初始值的、类型为 T 的标量值
scatter	XlaComputation	T, T -> T 类型的二元计算，用于应用每个散点源元素及其目标元素

下图显示了使用 SelectAndScatter 的示例，其中 select 函数计算其参数中的最大值。请注意，当窗口重叠时，如下图所示，operand 数组的索引可能会被不同的窗口多次选择。在图中，两个顶部窗口（蓝色和红色）都选择了值 9 的元素，并且二元加法 scatter 函数会生成值 8 (2 + 6) 的输出元素。

scatter 函数的计算顺序是任意的，并且可能是不确定的。因此，scatter 函数不应对重新关联过于敏感。如需了解详情，请参阅 Reduce 上下文中有关关联的讨论。

发送

另请参阅 XlaBuilder::Send。

Send(operand, channel_handle)

参数	类型	语义
operand	XlaOp	要发送的数据（T 类型的数组）
channel_handle	ChannelHandle	每个发送/接收对的唯一标识符

将给定的运算数数据发送到共享同一通道句柄的其他计算中的 Recv 指令。不返回任何数据。

与 Recv 操作类似，Send 操作的客户端 API 表示同步通信，并在内部分解为 2 条 HLO 指令（Send 和 SendDone），以启用异步数据传输。另请参阅 HloInstruction::CreateSend 和 HloInstruction::CreateSendDone。

Send(HloInstruction operand, int64 channel_id)

开始将运算数异步传输到由具有相同通道 ID 的 Recv 指令分配的资源。返回一个上下文，后续的 SendDone 指令使用该上下文来等待数据传输完成。上下文是一个 {operand (shape), request identifier (U32)} 元组，只能由 SendDone 指令使用。

SendDone(HloInstruction context)

在给定由 Send 指令创建的上下文的情况下，等待数据传输完成。该指令不返回任何数据。

渠道说明时间安排

每个通道（Recv、RecvDone、Send、SendDone）的 4 条指令的执行顺序如下所示。

• Recv发生在 Send之前
• Send发生在 RecvDone之前
• Recv发生在 RecvDone之前
• Send发生在 SendDone之前

当后端编译器为通过通道指令进行通信的每个计算生成线性时间表时，计算之间不得存在循环。例如，以下时间表会导致死锁。

切片

另请参阅 XlaBuilder::Slice。

切片从输入数组中提取子数组。子数组与输入的排名相同，并包含输入数组（其中边界框的维度和索引作为切片操作的参数）内边界框内的值。

Slice(operand, start_indices, limit_indices, strides)

参数	类型	语义
operand	XlaOp	T 类型的 N 维数组
start_indices	ArraySlice<int64>	由 N 个整数组成的列表，其中包含每个维度切片的起始索引。值必须大于或等于零。
limit_indices	ArraySlice<int64>	由 N 个整数组成的列表，其中包含每个维度切片的结束索引（不含索引）。每个值必须大于或等于维度的相应 start_indices 值，且小于或等于维度的大小。
strides	ArraySlice<int64>	决定切片输入步长的 N 个整数的列表。Slice 会选择维度 d 中的每个 strides[d] 元素。

一维示例：

let a = {0.0, 1.0, 2.0, 3.0, 4.0}
Slice(a, {2}, {4}) produces:
  {2.0, 3.0}

二维示例：

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

Slice(b, {2, 1}, {4, 3}) produces:
  { { 7.0,  8.0},
    {10.0, 11.0} }

排序

另请参阅 XlaBuilder::Sort。

Sort(operands, comparator, dimension, is_stable)

参数	类型	语义
operands	ArraySlice<XlaOp>	要排序的操作数。
comparator	XlaComputation	要使用的比较器计算。
dimension	int64	排序所依据的维度。
is_stable	bool	是否应使用稳定排序。

如果只提供一个操作数：

• 如果操作数是 1 阶张量（数组），则结果为有序数组。 如果要将数组按升序排序，比较器应执行小于比较。正式地说，数组排序后，它会存储所有索引位置 i, j，其中 i < j 为 comparator(value[i], value[j]) = comparator(value[j], value[i]) = false 或 comparator(value[i], value[j]) = true。

• 如果运算数具有较高的排名，则系统会根据提供的维度对运算数进行排序。例如，对于 2 阶张量（矩阵），维度值 0 将单独对每个列进行排序，维度值 1 将独立对每一行进行排序。如果未提供维度编号，则默认选择最后一个维度。对于已排序的维度，排序顺序与排名为 1 的情况相同。

如果提供了 n > 1 运算数：

• 所有 n 运算数都必须是具有相同维度的张量。张量的元素类型可能不同。

• 所有运算数会一起排序，而不是单独排序。从概念上讲，运算数被视为元组。在检查是否需要交换索引位置 i 和 j 处的每个运算数的元素时，会使用 2 * n 标量参数调用比较器，其中参数 2 * k 对应于 k-th 运算数中 i 位置的值，参数 2 * k + 1 对应于 k-th 运算数中位置 j 处的值。因此，比较器通常会相互比较 2 * k 和 2 * k + 1 参数，并且可能会使用其他参数对作为约束条件。

• 结果会得到一个元组，其中包含有序的运算数（沿提供的维度，如上所述）。元组的 i-th 运算数对应于 Sort 的 i-th 运算数。

例如，如果有三个运算数 operand0 = [3, 1]、operand1 = [42, 50]、operand2 = [-3.0, 1.1]，并且比较器只比较 operand0 的值与小于号的值，则排序的输出为元组 ([1, 3], [50, 42], [1.1, -3.0])。

如果 is_stable 设置为 true，排序可以保证稳定，也就是说，如果存在比较器认为相等的元素，则相等值的相对顺序会保留。当且仅当 comparator(e1, e2) = comparator(e2, e1) = false 时，e1 和 e2 两个元素相等。默认情况下，is_stable 设置为 false。

Transpose

另请参阅 tf.reshape 操作。

Transpose(operand)

参数	类型	语义
operand	XlaOp	要转置的操作数。
permutation	ArraySlice<int64>	如何排列维度。

使用给定的排列对运算数维度进行排列，因此 ∀ i . 0 ≤ i < rank ⇒ input_dimensions[permutation[i]] = output_dimensions[i]。

这与 Reshape(operand, per Workspace, Permute(per 相匹配, operand.shape.dimensions)) 相同。

TriangularSolve

另请参阅 XlaBuilder::TriangularSolve。

通过前向换元或后换换法求解具有下三角系数矩阵或上三角系数矩阵的线性方程组。该例程沿前导维度广播，在给定 a 和 b 的情况下，为变量 x 求解一个矩阵系统 op(a) * x = b（即 x * op(a) = b），其中 op(a) 为 op(a) = a、op(a) = Transpose(a) 或 op(a) = Conj(Transpose(a))。

TriangularSolve(a, b, left_side, lower, unit_diagonal, transpose_a)

参数	类型	语义
a	XlaOp	形状为 [..., M, M] 的复杂或浮点类型的秩 > 2 数组。
b	XlaOp	如果 left_side 为 true，则为 > 2 的相同类型数组，如果 left_side 为 true，则为 [..., K, M]，否则为 [..., K, M]。[..., M, K]
left_side	bool	指示是求解 op(a) * x = b (true) 还是 x * op(a) = b (false) 形式的系统。
lower	bool	使用 a 的上三角形还是下三角形。
unit_diagonal	bool	如果为 true，则假定 a 的对角线元素为 1，且不会被访问。
transpose_a	Transpose	是按原样使用 a，转置它，还是采用其共置转置。

根据 lower 的值，仅从 a 的下三角形/上三角形读取输入数据。来自另一个三角形的值会被忽略。输出数据会在同一个三角形中返回；另一个三角形中的值由实现定义，可以是任何内容。

如果 a 和 b 的阶大于 2，则将其视为矩阵批次，其中除了次要的 2 个维度外，它们都属于批量维度。a 和 b 的批量维度必须相同。

元组

另请参阅 XlaBuilder::Tuple。

包含可变数量数据句柄的元组，每个句柄都有自己的形状。

这类似于 C++ 中的 std::tuple。从概念上来讲：

let v: f32[10] = f32[10]{0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
let s: s32 = 5;
let t: (f32[10], s32) = tuple(v, s);

元组可以通过 GetTupleElement 操作解构（访问）。

同时

另请参阅 XlaBuilder::While。

While(condition, body, init)

参数	类型	语义
condition	XlaComputation	T -> PRED 类型的 XlaComputation，用于定义循环的终止条件。
body	XlaComputation	T -> T 类型的 XlaComputation，用于定义循环的正文。
init	T	condition 和 body 的参数的初始值。

按顺序执行 body，直到 condition 失败。这与许多其他语言中的典型 while 循环类似，但下面列出的差异和限制除外。

• While 节点返回 T 类型的值，该值是上次执行 body 的结果。
• T 类型的形状是静态确定的，并且在所有迭代中必须相同。

计算的 T 参数在第一次迭代中使用 init 值初始化，并在后续每次迭代中自动更新为 body 中的新结果。

While 节点的一个主要用例是在神经网络中实现重复执行训练。简化的伪代码如下所示，以及表示计算的图表。可在 while_test.cc 中找到该代码。此示例中的类型 T 是 Tuple，由 int32（表示迭代计数）和 vector[10]（表示累加器）组成。对于 1000 次迭代，循环不断向累加器添加一个常量向量。

// Pseudocode for the computation.
init = {0, zero_vector[10]} // Tuple of int32 and float[10].
result = init;
while (result(0) < 1000) {
  iteration = result(0) + 1;
  new_vector = result(1) + constant_vector[10];
  result = {iteration, new_vector};
}