锯齿状张量算子

高层概述

锯齿状张量算子的目的是处理输入数据的某个维度是“锯齿状”的情况，即给定维度中的每个连续行的长度可能不同。这类似于 PyTorch 中的 NestedTensor 实现和 Tensorflow 中的 RaggedTensor 实现。

这类输入的两个著名例子是

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

• 可能输入到自然语言处理系统的批量分词句子。

锯齿状张量格式

在 FBGEMM_GPU 中，锯齿状张量实际上被表示为一个三张量对象。这三个张量是：值 (Values)、最大长度 (MaxLengths) 和 偏移量 (Offsets)。

值 (Values)

Values 定义为一个二维张量，包含锯齿状张量中的所有元素值，即 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(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 的行数）结束。

最大长度 (Max Lengths)

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 仅支持锯齿状张量的逐元素加法、乘法和转换操作。

算术运算

锯齿状张量的加法和乘法类似于哈达玛积 (Hadamard Product)，并且只涉及锯齿状张量的 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. 算术运算 - 像往常一样对密集或锯齿状张量执行算术运算。