本文介绍一种基于 torch.scatter_reduce_ 的 GPU 加速方法，用于在单次操作中计算有序整数张量中每个值（0 到 max(t)−1）首次出现的索引，避免 Python 循环，显著提升大规模数据处理效率。
本文介绍一种基于 `torch.scatter_reduce_` 的 gpu 加速方法，用于在单次操作中计算有序整数张量中每个值（0 到 max(t)−1）首次出现的索引，避免 python 循环，显著提升大规模数据处理效率。
在深度学习与科学计算中，常需对离散标签或类别索引张量进行“首次出现位置”统计——例如，在序列建模中定位每类 token 的起始位置，或在图神经网络中快速构建节点首次访问映射。若使用 Python 循环或 torch.where 逐值查找，时间复杂度为 O(N×K)（N 为输入长度，K 为类别数），在 GPU 上无法并行化，性能极差。
PyTorch 2.1+ 提供的 scatter_reduce_ 操作为此类问题提供了高效解法：它支持在目标张量上按索引聚合源值，且 reduce="amin" 模式可天然实现“对每个目标槽位保留最小（即最早）索引”的语义。
以下是完整、健壮的实现：
import torch
def first_occurrence_indices(t):
"""
Compute the first occurrence index of each value in [0, max(t)].
Args:
t (torch.Tensor): 1D tensor of non-negative integers (dtype=torch.long).
Returns:
torch.Tensor: 1D tensor `a` of length max(t)+1, where `a[i]` is the
smallest index j such that t[j] == i, or -1 if i not in t.
"""
if t.numel() == 0:
return torch.empty(0, dtype=torch.long, device=t.device)
max_val = torch.max(t).item() + 1
a = torch.full((max_val,), -1, dtype=torch.long, device=t.device)
# Generate indices [0, 1, 2, ..., len(t)-1] on the same device
indices = torch.arange(t.numel(), device=t.device)
# Scatter indices into `a` at positions `t`, keeping minimum index per key
a.scatter_reduce_(0, t, indices, reduce="amin", include_self=False)
return a
✅ 使用示例：
t = torch.tensor([2, 1, 3, 2, 3, 0, 1, 5], device="cuda") result = first_occurrence_indices(t) print(result) # tensor([ 5, 1, 0, 2, -1, 7])
输出解释：
- 值
0首次出现在索引5→result[0] == 5 - 值
1首次出现在索引1→result[1] == 1 - 值
4未出现 →result[4] == -1（初始填充值）
⚠️ 注意事项：
- 输入
t必须为非负整数（否则t中负值将导致索引越界错误）；如需支持负数，应先做偏移归一化（如t_shifted = t - t.min()），并在结果后逆向调整。 -
scatter_reduce_在 CUDA 上表现优异，但对极稀疏场景（如max(t) ≫ len(t)）可能因分配大张量而浪费显存；此时可考虑哈希表方案（如torch.unique(..., return_inverse=True, return_counts=True)配合cumsum优化）。 - 确保 PyTorch ≥ 2.1 且 CUDA 支持（
scatter_reduce_在 CPU 上同样可用，但 GPU 加速优势更明显）。
该方法将算法复杂度降至 O(N)，全程张量化、无分支、完全 GPU 友好，是处理百万级标签张量的首选方案。
