不规则张量操作符

高层概述

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

这种输入的两个显著示例是

• 推荐系统中的稀疏特征输入

• 可能输入到自然语言处理系统的批处理标记化句子。

不规则张量格式

不规则张量在 FBGEMM_GPU 中实际上表示为三张量对象。这三个张量是：**Values**（值）、**MaxLengths**（最大长度）和 **Offsets**（偏移量）。

Values（值）

Values 定义为一个 2D 张量，其中包含不规则张量中的所有元素值，即 Values.numel() 是不规则张量中的元素数量。Values 中每一行的大小由不规则张量中最小（最内层）维子张量（不包含大小为 0 的张量）的最大公约数确定。

Offsets（偏移量）

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

例如，Offset[i] = [ 0, 3, 4 ] 意味着当前维度 i 被分成两组，由索引边界 [0 , 3) 和 [3, 4) 表示。对于每个 Offsets[i]（其中 0 <= i < len(Offsets) - 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 的行数）。

Max Lengths（最大长度）

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

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

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

不规则张量示例

下图显示了一个不规则张量示例，其中包含三个 2D 子张量，每个子张量具有不同的维度

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

因为示例不规则张量包含 2D 子张量，所以 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 仅支持不规则张量的元素级加法、乘法和转换操作。

算术操作

不规则张量加法和乘法类似于 Hadamard 乘积，且仅涉及不规则张量的 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]

转换操作

不规则到密集

将不规则张量 \(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] 的分区索引中所示的范围加载到密集张量上（参见表了解组与密集表中的对应行的映射）。

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

密集到不规则

对于从密集张量到不规则张量的转换，密集张量中的值被加载到不规则张量的 Values 中。然而，给定的密集张量可能与 Offsets 参考的形状不同。如果密集张量的相关维度小于预期，这可能导致不规则张量无法读取相应的密集位置。发生这种情况时，我们将填充值赋给对应的 Values（见下文）。

组合算术 + 转换操作

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

dense_tensor + jagged_tensor → dense_tensor (or jagged_tensor)

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

1. 转换操作 - 根据目标张量的所需格式，从不规则转换为密集，或从密集转换为不规则。转换后，操作数张量（无论是密集还是不规则）应具有完全相同的形状。

2. 算术操作 - 照常执行密集或不规则张量的算术操作。