操作语义

下面介绍了 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 必须为空(在这种情况下,所有副本都属于一个组,按照从 0N - 1 的顺序排列),或者包含与副本数量相同的元素数量。例如,replica_groups = {0, 2}, {1, 3} 会在副本 02 以及 13 之间执行串联。
  • shard_count 是每个副本组的大小。在 replica_groups 为空的情况下,我们需要使用该值。
  • channel_id 用于跨模块通信:只有具有相同 channel_idall-gather 操作才能相互通信。

输出形状是 all_gather_dim 使 shard_count 倍变大的输入形状。例如,如果有两个副本,并且运算数在两个副本上分别为 [1.0, 2.5][3.0, 5.25],则此操作的 all_gather_dim0 的输出值在这两个副本上都将是 [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 是数组元组,则会对该元组的每个元素执行 all-reduce。
  • replica_groups 是执行缩减操作的副本组的列表(可以使用 ReplicaId 检索当前副本的副本 ID)。replica_groups 必须为空(在这种情况下,所有副本都属于一个组),或者包含与副本数相同的元素数量。例如,replica_groups = {0, 2}, {1, 3} 会在副本 02 以及 13 之间执行缩减。
  • channel_id 用于跨模块通信:只有具有相同 channel_idall-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 的相同顺序串联。然后,另一个 AllToAll 将在副本 4、5、0 内应用,串联顺序也是 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 个部分分散到所有核心。然后,每个核心按照维度 1 的顺序串联收到的部分(按照核心 0-4 的顺序)。因此,每个核心上的输出的形状为 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 值 (\(\epsilon\))
feature_index int64 operand”中的特征维度索引

对于特征维度中的每个特征(feature_indexoperand 中特征维度的索引),该操作会计算 operandoffsetscale 相对于所有其他维度的梯度。feature_index 必须是 operand 中特征维度的有效索引。

这三次梯度由以下公式定义(假设四维数组为 operand,特征维度索引为 l,批次大小为 m,空间大小为 wh):

\[ \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} \]

输入 meanvariance 表示跨批次和空间维度的时刻值。

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

输出 类型 语义
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 Epsilon 值
feature_index int64 operand”中的特征维度索引

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

BatchNormInference 等同于调用 BatchNormTraining 而不计算每个批次的 meanvariance。它会改用输入 meanvariance 作为估算值。此操作的目的是减少推理过程中的延迟时间,因此称为 BatchNormInference

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

BatchNormTraining

如需详细了解该算法,另请参阅 XlaBuilder::BatchNormTrainingthe original batch normalization paper

跨批次和空间维度归一化数组。

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

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

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

对于 operand \(x\) 中包含 m 元素且空间维度大小为 wh(假设 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_meanbatch_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

计算一批对称 (Hermitian) 正定矩阵的乔尔斯基分解

Cholesky(a, lower)

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

如果 lowertrue,则计算下三角矩阵 l,使 $a = l。 l^T$。如果 lowerfalse,则计算上三角矩阵 u,以便\(a = u^T . u\)。

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

如果 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 类型的标量。

使用标量 minmax 的示例:

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::Collapsetf.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

如果 predtrue,则执行 true_computation;如果 predfalse,则执行 false_computation,并返回结果。

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

请注意,根据 pred 的值,系统仅会执行 true_computationfalse_computation 中的一个。

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 维数组
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 阶数组。尽管 rh 当然也是一个输入,但这称为输入。在神经网络中,这些是输入激活。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_dilationrhs_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_countbatch_group_count 必须是输出特征大小的除数。

输出形状具有以下尺寸:

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

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

卷积窗口的有效位置由填充之后的步长和底区大小决定。

为了描述卷积的作用,请考虑一个 2D 卷积,并在输出中选择一些固定的 batchzyx 坐标。(y,x) 是窗口某个角在基本区域(例如左上角,具体取决于您如何解读空间尺寸)内的位置。现在,我们获得了一个取自基本区域的 2D 窗口,其中每个 2D 点都与一个 1d 矢量相关联,从而得到一个 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 类似,可执行从数据形状到目标形状的元素级转换操作。维度必须匹配,并且转换是元素级转换;例如,通过 s32f32 的转换例程,s32 元素变为 f32 元素。

ConvertElementType(operand, new_element_type)

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

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

T=s32U=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 的第一个维度执行产品求和。这些维度就是“合同规定的”维度lhsrhs 的合同规定的尺寸必须相同。在实践中,它可用于执行向量之间的点积、向量/矩阵乘法或矩阵/矩阵乘法。

DotGeneral

另请参阅 XlaBuilder::DotGeneral

DotGeneral(lhs, rhs, dimension_numbers)

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

与 Dot 类似,但允许为 lhsrhs 指定协定维度和批量维度编号。

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

DotGeneral 会计算产品总和,而不是 dimension_numbers 中指定的合同维度。

lhsrhs 中关联的协定维度编号无需相同,但必须具有相同的维度大小。

签约维度编号示例:

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

lhsrhs 中关联的批次维度编号必须具有相同的维度大小。

带有批次维度编号(批次大小 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 updateupdate 的形状决定了更新结果的子数组的形状。 start_indices 的形状必须为 1,且维度大小等于 operand 的阶。

DynamicUpdateSlice(operand, update, start_indices)

参数 类型 语义
operand XlaOp T 类型的 N 维数组
update XlaOp 包含切片更新的 N 维数组,类型为 T。更新形状的每个维度都必须严格大于零,且起始值 + 更新值必须小于或等于每个维度的操作数大小,以避免生成出界更新索引。
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)

其中,OpAdd(加法)、Sub(减法)、Mul(乘法)、Div(除法)、Rem(余数)、Max(最大值)、Min(最小)、LogicalAnd(逻辑 AND)或 LogicalOr(逻辑 OR)之一。

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

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

OpRem 时,结果的符号取自被除数,并且结果的绝对值始终小于除数的绝对值。

整数除法溢出(有符号/无符号除号/余数除以零,或 INT_SMIN 的有符号除数/余数,其中 -1)会生成实现定义的值。

对于以下操作,存在支持不同排名广播的备用变体:

Op(lhs, rhs, broadcast_dimensions)

其中,Op 同上。这种运算的变体应该用于不同秩数组之间的算术运算(例如向向量添加矩阵)。

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

元素级比较操作

另请参阅 XlaBuilder::Eq

支持一组标准元素级二进制比较操作。请注意,比较浮点类型时,标准 IEEE 754 浮点比较语义适用。

Op(lhs, rhs)

其中,OpEq(等于)、Ne(不等于)、Ge(大于或等于)、Gt(大于)、Le(小于或等于)、Lt(小于)之一。另一组运算符 EqTotalOrder、NeTotalOrder、GeTotalOrder、GtTotalOrder、LeTotalOrder 和 LtTotalOrder 可提供相同的功能,不过它们还支持对浮点数进行总排序,方法是强制执行 -NaN < -Inf < -Finite <-0 + < +0 < 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 阶)。

Fft

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

另请参阅 XlaBuilder::Fft

参数 类型 语义
operand XlaOp 我们要进行傅里叶转换的数组。
fft_type FftType 请参见下表。
fft_length ArraySlice<int64> 要转换的轴的时间域长度。这对于 IRFFT 调整最内轴大小的情况尤其如此,因为 RFFT(fft_length=[16]) 的输出形状与 RFFT(fft_length=[17]) 相同。
FftType 语义
FFT 转发复杂到复杂 FFT。形状保持不变。
IFFT 复杂到复杂 FFT 的逆。形状保持不变。
RFFT 转发实际到复杂 FFT。如果 fft_length[-1] 为非零值,则最内层轴的形状会缩减为 fft_length[-1] // 2 + 1,从而省略转换信号中超出奈奎斯特频率的反共共振部分。
IRFFT 实数转换为复数 FFT(即取复数,返回实数)。如果 fft_length[-1] 为非零值,则最内轴的形状将展开为 fft_length[-1],通过 1fft_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 调用方是否保证会对索引进行排序。
unique_indices bool 调用方是否保证索引是唯一的。

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

输出是排名为 batch_dims.size + offset_dims.size 的数组。

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

如果 index_vector_dim 等于 start_indices.rank,我们会隐式认为 start_indices 具有尾随 1 维度(即,如果 start_indices 的形状为 [6,7]index_vector_dim2,那么我们会隐式将 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_sizesslice_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. 使用 Sstart_index_map 散布,Sin创建起始索引operandS。更确切地说:

    1. Sin[start_index_map[k]] = S[k](如果 k < start_index_map.size)。

    2. Sin[_] = 0,否则为 0

  3. operand 中创建索引 Oin,根据 collapsed_slice_dims 集,在 Out 中的偏移维度分布索引。更确切地说:

    1. 如果 k < offset_dims.sizeremapped_offset_dims 定义见下文),则 Oin[remapped_offset_dims(k)] = Out[offset_dims[k]]。

    2. Oin[_] = 0,否则为 0

  4. InOin + Sin,其中 + 表示元素级加法。

remapped_offset_dims 是一个单调函数,域为 [0, offset_dims.size),范围为 [0, operand.rank) \ collapsed_slice_dims。比如说,offset_dims.size4operand.rank6collapsed_slice_dims 为 {0, 2},则 remapped_offset_dims 为 {01, 13, 24, 35}。

如果 indices_are_sorted 设置为 true,则 XLA 可以假定 start_indices 已按 start_index_map 升序排序。否则,语义就是实现。

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

非正式说明和示例

通俗地说,输出数组中的每个索引 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 - 1index_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],用于确定 O0 和 O1 的范围,而后者又决定了 Slice 的边界。

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

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

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

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

  1. 我们可以配置输出形状中的哪些尺寸是偏移尺寸(在上一个示例中,包含 O0O1 的维度)。输出批量维度(上一个示例中包含 G0G1 的维度)被定义为非偏移维度的输出维度。

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

  3. 从“Gather Indices”数组(在上一示例中为 (X, Y))中提取的切片可能具有少于输入数组排名的元素,并且显式映射会指示应如何扩展索引,以使其与输入的排名相同。

最后,我们使用 (2) 和 (3) 来实现 tf.gather_nd

G0G1 用于照常从收集索引数组中分割起始索引,但起始索引只有一个元素 X同样,只有一个输出偏移索引的值为 O0。不过,在用作输入数组的索引之前,这些索引会根据“Gather Index Mapping”(正式说明中的 start_index_map)和“Offset Mapping”(正式说明中的 remapped_offset_dims)分别扩展到 [X,0] 和 [0,O0],从而为 [X,O0] [G] [G1]{1] [G] [G] [G] [G] 输出索引。00000OO11GatherIndicestf.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 参数
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 指定每个维度的边缘内边距和内部内边距。

PaddingConfigPaddingConfigDimension 的重复字段,其中包含每个维度的三个字段:edge_padding_lowedge_padding_highinterior_padding

edge_padding_lowedge_padding_high 分别指定在每个维度的低端(索引 0 旁边)和高端(最高索引旁边)处添加的内边距。边缘内边距的量可以是负值,负内边距的绝对值表示要从指定维度移除的元素数量。

interior_padding 指定每个维度的任意两个元素之间添加的内边距大小;此值不能为负数。内部内边距在逻辑上发生在边缘内边距之前,因此对于负边缘内边距,系统会从内内边距运算数中移除元素。

如果边缘填充对均为 (0, 0) 且内部填充值全部为 0,则此运算属于空操作。下图显示了二维数组的不同 edge_paddinginterior_padding 值的示例。

接收

另请参阅 XlaBuilder::Recv

Recv(shape, channel_handle)

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

从具有相同通道句柄的另一个计算中的 Send 指令接收指定形状的数据。针对已接收的数据返回 XlaOp。

Recv 操作的客户端 API 表示同步通信。不过,该指令在内部会被分解为 2 条 HLO 指令(RecvRecvDone),以启用异步数据传输。另请参阅 HloInstruction::CreateRecvHloInstruction::CreateRecvDone

Recv(const Shape& shape, int64 channel_id)

分配从具有相同 channel_id 的 Send 指令接收数据所需的资源。返回已分配资源的上下文,以下 RecvDone 指令使用该上下文来等待数据传输完成。上下文是 {receivebuffer (shape)、请求标识符 (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_iT_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 中。

与二维示例类似,我们只需减少一个维度。例如,如果我们减少维度 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 数组的所有维度减小 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] 和初始化值 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} 会在副本 02 以及 13 之间执行缩减,然后散开结果。
  • shard_count 是每个副本组的大小。在 replica_groups 为空的情况下,我们需要使用该值。如果 replica_groups 不为空,则 shard_count 必须等于每个副本组的大小。
  • channel_id 用于跨模块通信:只有具有相同 channel_idreduce-scatter 操作才能相互通信。

输出形状是将 scatter_dimension 缩小了 shard_count 倍的输入形状。例如,如果有两个副本,并且运算数在两个副本上分别为 [1.0, 2.25][3.0, 5.25],则对于第一个副本,此操作的 scatter_dim0 的输出值将为 [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。下图说明了使用两个不同的步幅值的情况。内边距会应用于输入的每个维度,计算过程与输入参数在填充后具有其尺寸一样。

举一个非常重要的填充示例,考虑对输入数组 [10000, 1000, 100, 10, 1] 计算缩减窗口的最小值(初始值为 MAX_FLOAT),维度为 3,步长为 2。内边距 kValid 会计算两个有效窗口([10000, 1000, 100][100, 10, 1])的最小值,得到输出 [100, 1]内边距 kSame 首先会填充数组,以便缩小窗口之后的形状与步幅一相同(方法是在两侧添加初始元素,从而得到[MAX_VALUE, 10000, 1000, 100, 10, 1, MAX_VALUE])。对填充的数组运行 Reduce-window 可对三个窗口执行([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::ReshapeCollapse 操作。

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

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 可将单元素数组转换为标量,反之亦然。例如,

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

Rev(反转)

另请参阅 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 的可用值:

散点图

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 将要分散到的类型为 T_0, ..., T_N 的 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 调用方是否保证会对索引进行排序。

其中:

  • 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 中的一组维度(按升序排列)。

散点的参数应遵循以下限制:

  • 每个 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 等于 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.sizeinserted_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 }。使用 Gscatter_indices 数组中查找索引矢量 S,使 S[i] = scatter_indices[Combine(G, i)],其中 Combine(A, b) 将 b 在 index_vector_dim 的位置插入 A。
  2. 通过使用 S 使用 scatter_dims_to_operand_dims 映射在 S 中散布 S,以在 operand 中创建索引 Sin更正式地说:
    1. Sin[scatter_dims_to_operand_dims[k]] = S[k](如果 k < scatter_dims_to_operand_dims.size)。
    2. Sin[_] = 0,否则为 0
  3. 根据 inserted_window_dims,将索引散布到 U 中的 update_window_dims 处,从而在每个 operands 数组中创建索引 Win。更正式地说:
    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.size4operand.rank6inserted_window_dims 为 {0, 2},则 window_dims_to_operand_dims 为 {01, 13, 24, 35})。
    2. Win[_] = 0,否则为 0
  4. IWin + Sin,其中 + 表示元素级加法。

总的来说,散点操作的定义如下。

  • 使用 operands 初始化 output,即对于 operands[J] 数组中的所有索引 O,使用 operands 初始化 output
    output[J][O] = operands[J][O]J
  • 对于 updates[J] 数组中的每个索引 Uoperand[J] 数组中的相应索引 O,如果 Ooutput 的有效索引:
    (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 升序排序。否则,语义就是实现。

简单来说,可以将散点操作视为收集操作的反转操作,即散点操作会更新输入中由相应收集操作提取的元素。

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

选择

另请参阅 XlaBuilder::Select

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

Select(pred, on_true, on_false)

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

数组 on_trueon_false 必须具有相同的形状。这也是输出数组的形状。数组 pred 必须与 on_trueon_false 具有相同的维度,并且具有 PRED 元素类型。

对于 pred 的每个元素 P,如果 P 的值为 true,则从 on_true 获取输出数组的相应元素;如果 P 的值为 false,则从 on_false 获取。作为受限形式的广播pred 可以是 PRED 类型的标量。在这种情况下,如果 predtrue,则输出数组全部从 on_true 获取;如果 predfalse,则完全从 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_trueon_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 指令(SendSendDone)以启用异步数据传输。另请参阅 HloInstruction::CreateSendHloInstruction::CreateSendDone

Send(HloInstruction operand, int64 channel_id)

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

SendDone(HloInstruction context)

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

频道说明的安排

每个通道(RecvRecvDoneSendSendDone)的 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 < jcomparator(value[i], value[j]) = comparator(value[j], value[i]) = falsecomparator(value[i], value[j]) = true

  • 如果运算数的排名较高,则按提供的维度对运算数进行排序。例如,对于 2 阶张量(矩阵),维度值 0 将独立地对每个列进行排序,维度值 1 将独立对每行进行排序。如果未提供维度编号,系统会默认选择最后一个维度。对于已排序的维度,排序顺序与 rank-1 情况相同。

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

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

  • 所有运算数会一起排序,而不是单独排序。从概念上讲,运算数被视为元组。在检查是否需要交换索引位置 ij 处的每个运算数的元素时,将使用 2 * n 标量参数调用比较器,其中参数 2 * k 对应于 k-th 运算数中位置 i 处的值,参数 2 * k + 1 对应于 k-th 运算数中位置 j 处的值。因此,比较器通常会将 2 * k2 * 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 时,e1e2 两个元素相等。默认情况下,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 新闻, Permute(peruation, operand.shape.dimensions)) 相同。

TriangularSolve

另请参阅 XlaBuilder::TriangularSolve

通过前向或反向替换求解具有下三角系数矩阵或上三角系数矩阵的线性方程组。在给定 ab 的情况下,此例程通过前导维度广播,为变量 x 求解一个矩阵系统 op(a) * x = b(即 x * op(a) = b),其中 op(a)op(a) = aop(a) = Transpose(a)op(a) = Conj(Transpose(a))

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

参数 类型 语义
a XlaOp 形状为 [..., M, M] 的复杂类型或浮点类型的 rank > 2 数组。
b XlaOp 一个 rank > 2 的同类型数组,如果 left_side 为 true,则为 [..., M, K],否则为 [..., K, M]
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、转置它,还是接受其共置转置。

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

如果 ab 的排序大于 2,则将其视为矩阵批次,其中除了次要 2 维度外,其他所有维度都属于批次维度。ab 的批次维度必须相等。

元组

另请参阅 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 conditionbody 的参数的初始值。

按顺序执行 body,直到 condition 失败。这类似于许多其他语言中的典型 when 循环,但存在下面列出的差异和限制。

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