稀疏算子¶
稳定 API¶
- torch.ops.fbgemm.permute_2D_sparse_data(permute, lengths, values, weights=None, permuted_lengths_sum=None) Tuple[Tensor, Tensor, Tensor | None]¶
沿第一个维度（dim 0）置换二维稀疏数据。请注意，二维指的是密集维度的数量。输入数据实际上是三维的，其中前两个维度是密集的，最后一个维度是不规则的（稀疏的）。要置换的数据可以更少或更多，并且可以有重复或没有重复。
- 参数：
permute (Tensor) – 一个一维张量，描述数据沿 dim 0 如何置换。permute[i] 表示位置 permute[i] 的数据移动到位置 i。此张量的长度是 dim 0 中要置换的总数据量。permute 中的值必须 >= 0 且 < lengths.shape[0]
lengths (Tensor) – 一个二维张量，包含对应于其他两个密集维度的锯齿形状。例如，在嵌入输入的情况下，三维形状为 (特征数、批大小、包大小)。lengths[t][b] 表示特征 t 和样本 b 的包大小。
values (Tensor) – 要置换的一维输入张量。此张量的长度必须等于 lengths.sum()。此张量可以是任何数据类型。
weights (Optional[Tensor] = None) – 可选的一维浮点张量。它必须与 values 具有相同的长度。它将以与 values 相同的方式置换
permuted_lengths_sum (Optional[int] = None) – 可选值，表示置换数据中的元素总数（输出形状）。如果未提供，运算符将计算此数据，这可能会导致设备-主机同步（如果使用 GPU）。因此，建议提供此值以避免这种同步。
- 返回：
置换后的长度、置换后的索引和置换后的权重的元组
示例
>>> permute = torch.tensor([1, 0, 2], dtype=torch.int32, device="cuda") >>> lengths = torch.tensor([[2, 3, 4, 5], [1, 2, 4, 8], [0, 3, 2, 3]], dtype=torch.int64, device="cuda") >>> values = torch.randint(low=0, high=100, size=(lengths.sum().item(),), dtype=torch.int64, device="cuda") >>> print(values) tensor([29, 12, 61, 98, 56, 94, 5, 89, 65, 48, 71, 54, 40, 33, 78, 68, 42, 21, 60, 51, 15, 47, 48, 68, 52, 19, 38, 30, 38, 97, 97, 98, 18, 40, 42, 89, 66], device='cuda:0') >>> torch.ops.fbgemm.permute_2D_sparse_data(permute, lengths, values) (tensor([[1, 2, 4, 8], [2, 3, 4, 5], [0, 3, 2, 3]], device='cuda:0'), tensor([78, 68, 42, 21, 60, 51, 15, 47, 48, 68, 52, 19, 38, 30, 38, 29, 12, 61, 98, 56, 94, 5, 89, 65, 48, 71, 54, 40, 33, 97, 97, 98, 18, 40, 42, 89, 66], device='cuda:0'), None)
- torch.ops.fbgemm.permute_1D_sparse_data(permute, lengths, values, weights=None, permuted_lengths_sum=None) Tuple[Tensor, Tensor, Tensor | None]¶
置换一维稀疏数据。请注意，一维指的是密集维度的数量。输入数据实际上是二维的，其中第一个维度是密集的，第二个维度是不规则的（稀疏的）。要置换的数据可以更少或更多，并且可以有重复或没有重复。
- 参数：
permute (Tensor) – 一个一维张量，描述数据沿 dim 0 如何置换。permute[i] 表示位置 permute[i] 的数据移动到位置 i。此张量的长度是 dim 0 中要置换的总数据量。permute 中的值必须 >= 0 且 < lengths.numel()
lengths (Tensor) – 一个一维张量，包含对应于其他密集维度的锯齿形状。lengths[i] 表示 dim 0 中位置 i 的数据的锯齿形状
values (Tensor) – 要置换的一维输入张量。此张量的长度必须等于 lengths.sum()。此张量可以是任何数据类型。
weights (Optional[Tensor] = None) – 可选的一维浮点张量。它必须与 values 具有相同的长度。它将以与 values 相同的方式置换
permuted_lengths_sum (Optional[int] = None) – 可选值，表示置换数据中的元素总数（输出形状）。如果未提供，运算符将计算此数据，这可能会导致设备-主机同步（如果使用 GPU）。因此，建议提供此值以避免这种同步。
- 返回：
置换后的长度、置换后的索引和置换后的权重的元组
- 示例
>>> permute = torch.tensor([1, 0, 3, 0], dtype=torch.int32, device="cuda") >>> lengths = torch.tensor([2, 3, 4, 5], dtype=torch.int64, device="cuda") >>> values = torch.randint(low=0, high=100, size=(lengths.sum().item(),), dtype=torch.int64, device="cuda") >>> print(values) tensor([ 1, 76, 24, 84, 94, 25, 15, 23, 31, 46, 9, 23, 34, 3], device='cuda:0') >>> torch.ops.fbgemm.permute_1D_sparse_data(permute, lengths, values) (tensor([3, 2, 5, 2], device='cuda:0'), tensor([24, 84, 94, 1, 76, 46, 9, 23, 34, 3, 1, 76], device='cuda:0'), None)
- torch.ops.fbgemm.expand_into_jagged_permute(permute, input_offset, output_offset, output_size) Tensor¶
将稀疏数据置换索引从特征维度扩展到批维度，用于稀疏特征在不同秩上具有不同批大小的情况。
该运算符通过连续地将每个包与其对应的特征映射到特征置换后批所在的位，将置换从特征级别扩展到批级别。该运算符将自动导出特征和批的偏移数组以计算输出置换。
- 参数：
permute (Tensor) – 特征级别的置换索引。
input_offset (Tensor) – 特征级别长度的独占偏移量。
output_offsets (Tensor) – 特征级别置换长度的独占偏移量。
output_size (int) – 输出张量中的元素数量
- 返回：
输出遵循以下公式
>>> output_permute[feature_offset[permute[feature]] + batch] <- bag_offset[batch]
- torch.ops.fbgemm.asynchronous_complete_cumsum(t_in) Tensor¶
计算完整的累积和。对于 GPU 运算符，运算符是非阻塞异步的。对于 CPU 运算符，它是一个阻塞运算符。
- 参数：
t_in (Tensor) – 输入张量
- 返回：
t_in 的完整累积和。形状为 t_in.numel() + 1
示例
>>> t_in = torch.tensor([7, 8, 2, 1, 0, 9, 4], dtype=torch.int64, device="cuda") >>> torch.ops.fbgemm.asynchronous_complete_cumsum(t_in) tensor([ 0, 7, 15, 17, 18, 18, 27, 31], device='cuda:0')
- torch.ops.fbgemm.offsets_range(offsets, range_size) Tensor¶
为每个 i 生成一个从 0 到 (offsets[i+1] - offsets[i]) 的整数序列，其中 0 <= i < offsets.numel()
- 参数：
offsets (Tensor) – 偏移量（完整的累积和值）
range_size (int) – 输出大小（总和）
- 返回：
包含偏移范围的张量
- 示例
>>> # Generate example inputs >>> lengths = torch.tensor([3, 4, 1, 9, 3, 7], dtype=torch.int64, device="cuda") >>> offsets = torch.ops.fbgemm.asynchronous_complete_cumsum(lengths) >>> range_size = offsets[-1].item() >>> print(range_size) 27 >>> offsets = offsets[:-1] >>> print(offsets) tensor([ 0, 3, 7, 8, 17, 20], device='cuda:0') >>> # Invoke >>> torch.ops.fbgemm.offsets_range(offsets, range_size) tensor([0, 1, 2, 0, 1, 2, 3, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 0, 1, 2, 0, 1, 2, 3, 4, 5, 6], device='cuda:0')
- torch.ops.fbgemm.segment_sum_csr(batch_size, csr_seg, values) Tensor¶
在给定的 CSR 数据上对每个段内的值求和，其中每行具有相同数量的非零元素。
- 参数：
batch_size (int) – 行跨度（每行中非零元素的数量）
csr_seg (Tensor) – 段长度的完整累积和。段长度是每个段内行的数量。 csr_seg 张量的形状为 num_segments + 1，其中 num_segments 是段的数量。
values (Tensor) – 要进行分段求和的值张量。张量中的元素数量必须是 batch_size 的倍数。
- 返回：
包含分段求和结果的张量。形状为段的数量。
示例
>>> batch_size = 2 >>> # Randomize inputs >>> lengths = torch.tensor([3, 4, 1], dtype=torch.int, device="cuda") >>> offsets = torch.ops.fbgemm.asynchronous_complete_cumsum(lengths) >>> print(offsets) tensor([0, 3, 7, 8], device='cuda:0', dtype=torch.int32) >>> values = torch.randn(lengths.sum().item() * batch_size, dtype=torch.float32, device="cuda") >>> print(values) tensor([-2.8642e-01, 1.6451e+00, 1.1322e-01, 1.7335e+00, -8.4700e-02, -1.2756e+00, 1.1206e+00, 9.6385e-01, 6.2122e-02, 1.3104e-03, 2.2667e-01, 2.3113e+00, -1.1948e+00, -1.5463e-01, -1.0031e+00, -3.5531e-01], device='cuda:0') >>> # Invoke >>> torch.ops.fbgemm.segment_sum_csr(batch_size, offsets, values) tensor([ 1.8451, 3.3365, -1.3584], device='cuda:0')
- torch.ops.fbgemm.keyed_jagged_index_select_dim1(values, lengths, offsets, indices, batch_size, weights=None, selected_lengths_sum=None) List[Tensor]¶
对给定键式锯齿张量 (KJT) 输入的批次维度 (dim 1) 执行索引选择操作。每个键的批次中的相同样本将被选择。请注意，每个 KJT 都有 3 个维度：( num_keys，batch_size，锯齿形维度)，其中 num_keys 是键的数量，batch_size 是批次大小。此运算符类似于置换运算符。
- 参数：
values (Tensor) – KJT 值张量，其中包含每个键的连接数据
lengths (Tensor) – KJT 长度张量，其中包含每个键 (dim 0) 和样本 (dim 1) 的锯齿形形状。形状为 num_keys * batch_size
offsets (Tensor) – KJT 偏移张量，它是 lengths 的完整累积和。形状为 num_keys * batch_size + 1
indices (Tensor) – 要选择的索引，即要选择的批次中的样本。 indices 的值必须大于等于 0 且小于 batch_size
batch_size (int) – 批次大小（KJT 的 dim 1）
weights (Optional[Tensor] = None) – 一个可选的浮点张量，它将以与 values 相同的方式进行选择。因此，它必须与 values 具有相同的形状。
selected_lengths_sum (Optional[int] = None) – 一个可选值，表示索引选择数据中的元素总数（输出形状）。如果未提供，运算符将计算此数据，这可能会导致设备-主机同步（如果使用 GPU）。因此，建议提供此值以避免此类同步。
- 返回：
索引选择 KJT 张量（作为值的列表、长度以及如果 weights 不为 None 则为权重）
示例
>>> num_keys = 2 >>> batch_size = 4 >>> output_size = 3 >>> # Randomize inputs >>> lengths = torch.randint(low=0, high=10, size=(batch_size * num_keys,), dtype=torch.int64, device="cuda") >>> print(lengths) tensor([8, 5, 1, 4, 2, 7, 5, 9], device='cuda:0') >>> offsets = torch.ops.fbgemm.asynchronous_complete_cumsum(lengths) >>> print(offsets) tensor([ 0, 8, 13, 14, 18, 20, 27, 32, 41], device='cuda:0') >>> indices = torch.randint(low=0, high=batch_size, size=(output_size,), dtype=torch.int64, device="cuda") >>> print(indices) tensor([3, 3, 1], device='cuda:0') >>> # Use torch.arange instead of torch.randn to simplify the example >>> values = torch.arange(lengths.sum().item(), dtype=torch.float32, device="cuda") >>> print(values) tensor([ 0., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12., 13., 14., 15., 16., 17., 18., 19., 20., 21., 22., 23., 24., 25., 26., 27., 28., 29., 30., 31., 32., 33., 34., 35., 36., 37., 38., 39., 40.], device='cuda:0') >>> # Invoke. Output = (output, lengths) >>> torch.ops.fbgemm.keyed_jagged_index_select_dim1(values, lengths, offsets, indices, batch_size) [tensor([14., 15., 16., 17., 14., 15., 16., 17., 8., 9., 10., 11., 12., 32., 33., 34., 35., 36., 37., 38., 39., 40., 32., 33., 34., 35., 36., 37., 38., 39., 40., 20., 21., 22., 23., 24., 25., 26.], device='cuda:0'), tensor([4, 4, 5, 9, 9, 7], device='cuda:0')]
- torch.ops.fbgemm.block_bucketize_sparse_features(lengths, indices, bucketize_pos, sequence, block_sizes, my_size, weights=None, batch_size_per_feature=None, max_B=-1, block_bucketize_pos=None, keep_orig_idx=False) Tuple[Tensor, Tensor, Tensor | None, Tensor | None, Tensor | None]¶
通过将稀疏特征划分为多个桶来预处理稀疏特征。每个特征都拆分为相同数量的桶，但不同特征的桶大小（宽度）可能不同。此外，每个特征内的桶大小也可能不同。
- 参数：
lengths (Tensor) – 稀疏特征的长度。该张量包含批次中每个样本和每个特征的长度。形状为 B * T，其中 B 是批次大小，T 是特征的数量。
indices (Tensor) – 稀疏数据。仅支持整数类型。形状为 lengths 的总和。
bucketize_pos (bool) – 如果为 True，则返回样本内的原始相对索引。例如， indices = [9, 8, 2, 1, 0, 8, 9] 和 lengths = [3, 4]。样本内索引的原始相对索引为 [0, 1, 2, 0, 1, 2, 3]
sequence (bool) – 如果为 True，则返回原始索引位置中的新索引位置（该张量称为 unbucketize_permute_data）。
block_sizes (Tensor) – 此张量用于特征内桶大小统一的情况（即，当 block_bucketize_pos=None 时）。该张量包含每个特征的桶大小（即桶宽度）。 block_sizes[t] 表示特征 t 的桶大小。形状为特征的数量。
my_size (int) – 每个特征的桶的数量。请注意，每个特征都具有相同数量的桶。
weights (Optional[Tensor] = None) – 一个可选的浮点张量，它将以与 indices 相同的方式进行分桶。此张量必须与 indices 具有相同的形状。
batch_size_per_feature (Optional[Tensor] = None) – 一个可选张量，其中包含不同特征的批次大小。如果非 None，则批次大小在特征之间不统一。否则，运算符将假设批次大小是统一的，并根据 lengths 和 block_sizes 张量推断它。
max_B (int = -1) – 最大批次大小。如果 batch_size_per_feature 不为 None，则必须设置此值。
block_bucketize_pos (Optional[List[Tensor]] = None) – 输入用于特征内非统一桶大小。 block_bucketize_pos 是张量的列表。每个张量包含每个特征的桶的范围偏移量。这些范围偏移量等效于桶大小的完整累积和。例如， [0, 4, 20] 表示两个桶。第一个桶的大小为 (4 - 0) = 4，第二个桶的大小为 (20 - 4) = 16。 block_bucketize_pos 的长度必须等于特征的数量。
keep_orig_idx (bool = False) – 如果为 True，则返回原始索引而不是每个桶内的相对索引。
- 返回：
包含张量元组
分桶长度。形状为 lengths.num() * my_size。
分桶索引。与 indices 形状相同。
分桶权重或如果 weights 为 None 则为 None。与 indices 形状相同。
分桶位置或如果 bucketize_pos=False 则为 None。与 indices 形状相同。
unbucketize_permute 或如果 sequence=False 则为 None。与 indices 形状相同。
示例:
>>> # Generate input example. Batch size = 2. Number of features = 4 >>> lengths = torch.tensor([0, 2, 1, 3, 2, 3, 3, 1], dtype=torch.int, device="cuda") >>> indices = torch.tensor([3, 4, 15, 11, 28, 29, 1, 10, 11, 12, 13, 11, 22, 20, 20], dtype=torch.int, device="cuda") >>> block_sizes = torch.tensor([[5, 15, 10, 20]], dtype=torch.int, device="cuda") >>> my_size = 2 # Number of buckets >>> # Invoke with keep_orig_idx=False, bucketize_pos=False, and >>> # sequence=False >>> torch.ops.fbgemm.block_bucketize_sparse_features( >>> lengths, >>> indices, >>> bucketize_pos=False, >>> sequence=False, >>> block_sizes=block_sizes, >>> my_size=my_size, >>> keep_orig_idx=False) >>> # The first 8 values in the returned lengths are the lengths for bucket >>> # 0 and the rests are the legths for bucket 1 (tensor([0, 2, 0, 1, 1, 0, 1, 0, 0, 0, 1, 2, 1, 3, 2, 1], device='cuda:0', dtype=torch.int32), tensor([ 3, 4, 11, 1, 11, 0, 13, 14, 0, 1, 2, 3, 2, 0, 0], device='cuda:0', dtype=torch.int32), None, None, None) >>> # Invoke with keep_orig_idx=True, bucketize_pos=True, and >>> # sequence=True >>> torch.ops.fbgemm.block_bucketize_sparse_features( >>> lengths, >>> indices, >>> bucketize_pos=True, >>> sequence=True, >>> block_sizes=block_sizes, >>> my_size=my_size, >>> keep_orig_idx=True) (tensor([0, 2, 0, 1, 1, 0, 1, 0, 0, 0, 1, 2, 1, 3, 2, 1], device='cuda:0', dtype=torch.int32), tensor([ 3, 4, 11, 1, 11, 15, 28, 29, 10, 11, 12, 13, 22, 20, 20], device='cuda:0', dtype=torch.int32), None, tensor([0, 1, 0, 0, 0, 0, 1, 2, 1, 0, 1, 2, 1, 2, 0], device='cuda:0', dtype=torch.int32), tensor([ 0, 1, 5, 2, 6, 7, 3, 8, 9, 10, 11, 4, 12, 13, 14], device='cuda:0', dtype=torch.int32)) >>> # Invoke with block_bucketize_pos >>> block_bucketize_pos = [ >>> torch.tensor([0, 2, 8], dtype=torch.int), >>> torch.tensor([0, 5, 10], dtype=torch.int), >>> torch.tensor([0, 7, 12], dtype=torch.int), >>> torch.tensor([0, 2, 16], dtype=torch.int), >>> ] >>> torch.ops.fbgemm.block_bucketize_sparse_features( >>> lengths, >>> indices, >>> bucketize_pos=False, >>> sequence=False, >>> block_sizes=block_sizes, >>> my_size=my_size, >>> block_bucketize_pos=block_bucketize_pos, >>> keep_orig_idx=False) (tensor([0, 0, 0, 1, 1, 1, 2, 1, 0, 2, 1, 2, 1, 2, 1, 0], device='cuda:0', dtype=torch.int32), tensor([14, 1, 6, 11, 10, 10, 1, 2, 7, 5, 14, 3, 4, 6, 9], device='cuda:0', dtype=torch.int32), None, None, None)
其他 API
