不规则张量运算符¶

高级概述¶

不规则张量运算符的目的是处理输入数据中的某些维度“不规则”的情况，即给定维度中每连续行可能具有不同的长度。这类似于 PyTorch 中的 NestedTensor 实现和 Tensorflow 中的 RaggedTensor 实现。

这种类型输入的两个显著例子是

推荐系统中的稀疏特征输入
可能输入自然语言处理系统的标记化句子的批次。

不规则张量格式¶

在 FBGEMm_GPU 中，不规则张量实际上表示为一个三张量对象。这三个张量是：**值**、**最大长度**和**偏移量**。

值¶

Values 定义为一个二维张量，包含不规则张量中的所有元素值，即 Values.numel() 是不规则张量中的元素数量。Values 中每行的尺寸源自不规则张量中最小的（最内层的）维度子张量（不包括尺寸为 0 的张量）的公最大公约数。

偏移量¶

Offsets 是一个张量列表，其中每个张量 Offsets[i] 表示列表中下一个张量 Offsets[i + 1] 的值的划分索引。

例如，Offset[i] = [ 0, 3, 4 ] 意味着当前维度 i 被分成两组，分别由索引边界 [0 , 3) 和 [3, 4) 表示。对于每个 Offsets[i]，其中 0 <= i < len(Offests) - 1，Offsets[i][0] = 0，并且 Offsets[i][-1] = Offsets[i+1].length。

Offsets[-1] 指的是 Values 的外部维度索引（行索引），即 offsets[-1] 将是 Values 本身的划分索引。因此，Offsets[-1]，张量以 0 开头，以 Values.size(0) 结尾（即 Values 的行数）。

最大长度¶

MaxLengths 是一个整数列表，其中每个值 MaxLengths[i] 表示 Offsets[i] 中对应偏移值之间的最大值。

MaxLengths[i] = max( Offsets[i][j] - Offsets[i][j-1]  |  0 < j  < len(Offsets[i]) )

MaxLengths 中的信息用于执行从不规则张量到普通（密集）张量的转换，它将用于确定张量的密集形式的形状。

不规则张量示例¶

下图显示了一个示例不规则张量，它包含三个二维子张量，每个子张量具有不同的维度。

在这个示例中，不规则张量中最内层维度中行的尺寸分别是 8、4 和 0，因此 Values 中每行的元素数量被设置为 4（最大公约数）。这意味着 Values 必须具有 9 x 4 的尺寸才能容纳不规则张量中的所有值。

由于示例不规则张量包含二维子张量，因此 Offsets 列表将需要具有 2 的长度才能创建划分索引。 Offsets[0] 表示维度 0 的划分，而 Offsets[1] 表示维度 1 的划分。

示例不规则张量中的 MaxLengths 值为 [4 , 2]。 MaxLengths[0] 源自 Offsets[0] 范围 [4, 0)，而 MaxLengths[1] 源自 Offsets[1] 范围 [0, 2)（或 [7, 9]、[3,5]）。

下表显示了应用于 Values 张量的划分索引，以构建示例不规则张量的逻辑表示。

`Offsets[0]`	`Offsets[0]` 范围	`Offsets[0]` 组	对应的 `Offsets[1]`	`Offsets[1]` 范围	`Values` 组	对应的 `Values`
`[ 0, 4, 6, 8 ]`	`[0, 4)`	组 1	`[ 0, 2, 3, 3, 5 ]`	`[ 0, 2 )`	组 1	`[ [ 1, 2, 3, 4 ], [ 5, 6, 7, 8 ] ]`
				`[ 2, 3 )`	组 2	`[ [ 1, 2, 3, 4 ] ]`
				`[ 3, 3 )`	组 3	`[ ]`
				`[ 3, 5 )`	组 4	`[ [ 1, 2, 3, 4 ], [ 5, 6, 7, 8 ] ]`
	`[4, 6)`	组 2	`[ 5, 6, 7 ]`	`[ 5, 6 )`	组 5	`[ [ 1, 2, 3, 4 ] ]`
				`[ 6, 7 )`	组 6	`[ [ 1, 2, 7, 9 ] ]`
	`[6, 8)`	组 3	`[ 7, 9 ]`	`[ 7, 9 )`	组 7	`[ [ 1, 2, 3, 4 ], [ 8, 8, 9, 6 ] ]`

不规则张量操作¶

在当前阶段，FBGEMM_GPU 仅支持不规则张量的逐元素加法、乘法和转换操作。

算术运算¶

不规则张量的加法和乘法与哈达玛积类似，只涉及不规则张量的 Values。例如

\[\begin{split}\begin{bmatrix} \begin{bmatrix} 1. & 2. \\ 3. & 4. \\ \end{bmatrix} \\ \begin{bmatrix} 5. & 6. \\ \end{bmatrix} \\ \begin{bmatrix} 7. & 8. \\ 9. & 10. \\ 11. & 12. \\ \end{bmatrix} \\ \end{bmatrix} \times \begin{bmatrix} \begin{bmatrix} 1. & 2. \\ 3. & 4. \\ \end{bmatrix} \\ \begin{bmatrix} 5. & 6. \\ \end{bmatrix} \\ \begin{bmatrix} 7. & 8. \\ 9. & 5. \\ 2. & 3. \\ \end{bmatrix} \\ \end{bmatrix} \rightarrow \begin{bmatrix} \begin{bmatrix} 1. & 4. \\ 9. & 16. \\ \end{bmatrix} \\ \begin{bmatrix} 25. & 36. \\ \end{bmatrix} \\ \begin{bmatrix} 49. & 64. \\ 81. & 50. \\ 22. & 36. \\ \end{bmatrix} \\ \end{bmatrix}\end{split}\]

因此，对不规则张量进行算术运算需要两个操作数具有相同的形状。换句话说，如果我们有不规则张量 \(A\)、\(X\)、\(B\) 和 \(C\)，其中 \(C = AX + B\)，则以下属性成立

// MaxLengths are the same
C.maxlengths == A.maxlengths == X.maxlengths == B.maxlengths

// Offsets are the same
C.offsets == A.offsets == X.offsets == B.offsets

// Values are elementwise equal to the operations applied
C.values[i][j] == A.values[i][j] * X.values[i][j] + B.values[i][j]

转换操作¶

不规则到稠密¶

../../_images/JaggedTensorConversion1.png

将不规则张量 \(J\) 转换为等效的稠密张量 \(D\) 从一个空的稠密张量开始。\(D\) 的形状基于 MaxLengths（Values 的内部维度）、Offsets[0] 的长度。\(D\) 中的维度数为

rank(D) = len(MaxLengths) + 2

对于 \(D\) 中的每个维度，维度大小为

dim(i) = MaxLengths[i-1]  // (0 < i < D.rank-1)

以不规则张量示例中的不规则张量为例，len(MaxLengths) = 2，因此等效的稠密张量的秩（维度数）将为 4。示例不规则张量有两个偏移张量，Offsets[0] 和 Offsets[1]。在转换过程中，Values 中的元素将根据 Offsets[0] 和 Offsets[1] 中的分区索引所表示的范围加载到稠密张量上（有关组到稠密表中相应行的映射，请参阅表格）。

../../_images/JaggedTensorConversion2.png

由于并非 Offsets[i] 中表示的每个分区范围的大小都等于 MaxLengths[i]，因此 \(D\) 的某些部分不会加载来自 \(J\) 的值。在这种情况下，这些部分将用填充值填充。在上面的例子中，填充值为 0。

稠密到不规则¶

对于从稠密到不规则张量的转换，稠密张量中的值将加载到不规则张量的 Values 中。但是，给定的稠密张量可能与 Offsets 所指的形状不同。这可能会导致在稠密张量的相关维度小于预期时，不规则张量无法读取相应的稠密位置。当这种情况发生时，我们会将填充值赋给相应的 Values（见下文）。

../../_images/JaggedTensorConversion3.png

组合的算术 + 转换操作¶

在某些情况下，我们希望执行以下操作

dense_tensor + jagged_tensor → dense_tensor (or jagged_tensor)

我们可以将此类操作分解为两个步骤

转换操作 - 根据目标张量的所需格式，从不规则转换为稠密或从稠密转换为不规则。转换后，操作数张量（无论是稠密的还是不规则的）都应该具有完全相同的形状。
算术运算 - 对稠密或不规则张量执行正常的算术运算。