不规则张量运算符

稳定 API

torch.ops.fbgemm.jagged_to_padded_dense(values, offsets, max_lengths, padding_value=0) → Tensor

将不规则张量转换为密集张量，使用指定的填充值进行填充。

参数：

• values (Tensor) – 不规则张量值

• offsets (Tensor[]) – 不规则偏移张量列表，每个不规则维度一个。

• max_lengths (int[]) – 包含每个不规则维度最大长度的列表。

• padding_value (float) – 设置在密集输出的空区域中的值，在不规则张量覆盖范围之外。

返回值：

填充后的密集张量

返回值类型：

Tensor

示例

>>> values = torch.tensor([[1,1],[2,2],[3,3],[4,4]])
>>> offsets = torch.tensor([0, 1, 3])
>>> torch.ops.fbgemm.jagged_to_padded_dense(values, [offsets], [3], 7)
tensor([[[1, 1],
         [7, 7],
         [7, 7]],
        [[2, 2],
         [3, 3],
         [7, 7]]])

其他 API

torch.ops.fbgemm.jagged_2d_to_dense(values, x_offsets, max_sequence_length) → Tensor

将具有 2D 值数组的不规则张量转换为密集张量，使用零进行填充。

参数：

• values (Tensor) – 包含不规则张量值的 2D 张量。

• x_offsets (Tensor) – 包含 values 张量中每个不规则行的起始点的 1D 张量。

• max_sequence_length (int) – 不规则维度中任何行的最大长度。

返回值：

填充后的密集张量

返回值类型：

Tensor

示例

>>> values = torch.tensor([[1,1],[2,2],[3,3],[4,4]])
>>> x_offsets = torch.tensor([0, 1, 3])
>>> torch.ops.fbgemm.jagged_2d_to_dense(values, x_offsets, 3)
tensor([[[1, 1],
         [0, 0],
         [0, 0]],
        [[2, 2],
         [3, 3],
         [0, 0]]])

torch.ops.fbgemm.jagged_1d_to_dense(values, offsets, max_sequence_length, padding_value) -> Tensor)

将具有 1D 值数组的不规则张量转换为密集张量，使用指定的填充值进行填充。

参数：

• values (Tensor) – 包含不规则张量值的 1D 张量。

• offsets (Tensor) – 包含 values 张量中每个不规则行的起始点的 1D 张量。

• max_sequence_length (int) – 不规则维度中任何行的最大长度。

• padding_value (int) – 设置在密集输出的空区域中的值，在不规则张量覆盖范围之外。

返回值：

填充后的密集张量

返回值类型：

Tensor

示例

>>> values = torch.tensor([1,2,3,4])
>>> offsets = torch.tensor([0, 1, 3])
>>> torch.ops.fbgemm.jagged_1d_to_dense(values, x_offsets, 3, 0)
tensor([[1, 0, 0],
        [2, 3, 0]])

torch.ops.fbgemm.dense_to_jagged(dense, x_offsets, total_L) -> (Tensor, Tensor[])

将密集张量转换为不规则张量，给出所得密集张量的所需偏移量。

参数：

• dense (Tensor) – 要转换的密集输入张量

• x_offsets (Tensor[]) – 不规则偏移张量列表，每个不规则维度一个。

• total_L (int, 可选) – 所得不规则张量中的值总数。

返回值：

所得不规则张量的值和偏移量。偏移量与输入的偏移量相同。

返回值类型：

(Tensor, Tensor[])

示例

>>> dense = torch.tensor([[[1, 1], [0, 0], [0, 0]], [[2, 2], [3, 3], [0, 0]]])
>>> x_offsets = torch.tensor([0, 1, 3])
>>> torch.ops.fbgemm.dense_to_jagged(dense, [x_offsets])
(tensor([[1, 1],
         [2, 2],
         [3, 3]]), [tensor([0, 1, 3])])

torch.ops.fbgemm.jagged_dense_elementwise_add(x_values, x_offsets, y) → Tensor

将不规则张量添加到密集张量，得到密集张量。不规则张量输入将使用零进行填充，以用于添加。

参数：

• x_values (Tensor) – 不规则张量值

• offsets (Tensor[]) – 不规则偏移张量列表，每个不规则维度一个。

• y (Tensor) – 密集张量

返回值：

不规则输入张量 + y 的总和

返回值类型：

Tensor

torch.ops.fbgemm.jagged_dense_elementwise_add_jagged_output(x_values, x_offsets, y) -> (Tensor, Tensor[])

将不规则张量添加到密集张量，得到与输入不规则张量结构相同的结构的不规则张量。

参数：

• x_values (Tensor) – 不规则张量值

• x_offsets (Tensor[]) – 不规则偏移张量列表，每个不规则维度一个。

• y (Tensor) – 密集张量

返回值：

所得不规则张量的值和偏移量。偏移量与输入的偏移量相同。

返回值类型：

(Tensor, Tensor[])

torch.ops.fbgemm.jagged_dense_dense_elementwise_add_jagged_output(x_values, x_offsets, y_0, y_1) -> (Tensor, Tensor[])

将不规则张量添加到两个密集张量的总和，得到与输入不规则张量结构相同的结构的不规则张量。

参数：

• x_values (Tensor) – 不规则张量值

• x_offsets (Tensor[]) – 不规则偏移张量列表，每个不规则维度一个。

• y_0 (Tensor) – 密集张量

• y_1 (Tensor) – 密集张量

返回值：

所得不规则张量的值和偏移量。偏移量与输入的偏移量相同。

返回值类型：

(Tensor, Tensor[])

torch.ops.fbgemm.jagged_dense_elementwise_mul(x_values, x_offsets, y) -> (Tensor, Tensor[])

按元素将不规则张量与密集张量相乘，得到与输入不规则张量结构相同的结构的不规则张量。

参数：

• x_values (Tensor) – 不规则张量值

• x_offsets (Tensor[]) – 不规则偏移张量列表，每个不规则维度一个。

• y (Tensor) – 密集张量

返回值：

所得不规则张量的值和偏移量。偏移量与输入的偏移量相同。

返回值类型：

(Tensor, Tensor[])

torch.ops.fbgemm.batched_dense_vec_jagged_2d_mul(Tensor v, Tensor a_values, Tensor a_offsets) → Tensor

批次稠密向量与不规则张量的向量矩阵乘法，稠密向量大小为 (B * H, max_N)，不规则张量大小为 (B, max_N, H * D)，其中 max_N 是不规则维度的最大大小。B * H 是批次大小，每个乘法是 max_N 与 [max_N, D] 的乘法。

参数：

• v (Tensor) – 稠密向量张量

• a_values (Tensor) – 不规则张量值

• a_offsets (Tensor []) – 不规则偏移张量的列表，每个不规则维度一个。

返回值：

大小为 (B * H, D) 的批次矩阵乘法的输出

返回值类型：

Tensor

torch.ops.fbgemm.stacked_jagged_1d_to_dense(*args, **kwargs)

torch.ops.fbgemm.stacked_jagged_2d_to_dense(*args, **kwargs)