本文介绍一种基于 PyTorch scatter_reduce_ 的高效方法，用于在 GPU 上快速构建索引张量：对输入有序整型张量 t，输出张量 a 满足 a[i] 为值 i 在 t 中首次出现的下标，未出现的值对应位置填 -1。
本文介绍一种基于 pytorch `scatter_reduce_` 的高效方法，用于在 gpu 上快速构建索引张量：对输入有序整型张量 `t`，输出张量 `a` 满足 `a[i]` 为值 `i` 在 `t` 中**首次出现的下标**，未出现的值对应位置填 `-1`。
在处理大规模序列数据（如 token ID 序列、类别标签流）时，常需快速获取每个类别首次出现的位置——例如用于注意力掩码构造、动态截断或稀疏索引预处理。若使用 Python 循环或 torch.where 逐值查找，时间复杂度为 O(N×K)（N 为输入长度，K 为类别数），无法充分利用 GPU 并行性。
PyTorch 2.1+ 提供的 scatter_reduce_ 操作可一次性完成“按键聚合”，其中 reduce="amin" 表示对每个目标索引位置取所有映射到该位置的源索引中的最小值，恰好对应“首次出现”的语义。
以下是完整实现：
import torch
def first_occurrence_indices(t):
"""
返回张量 t 中每个值首次出现的索引。
Args:
t (torch.Tensor): 1D int/long 张量，值域应为非负整数
Returns:
torch.Tensor: 长度为 max(t)+1 的 long 张量，
a[i] = t 中 i 首次出现的下标；若 i 未出现，则为 -1
"""
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)
# 构造源索引：t[i] 对应位置 i
indices = torch.arange(t.numel(), device=t.device)
# 关键操作：将 indices 按 t 的值作为目标索引，执行 "amin" 聚合
# 即对每个值 v，收集所有满足 t[j]==v 的 j，并取 min(j)
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（由torch.full初始化保证）
⚠️ 注意事项：
- 输入
t必须为非负整数（torch.long或torch.int64），否则scatter_reduce_可能越界或静默失败；建议添加assert (t >= 0).all()校验。 - 若
t中存在极大值（如1e6），会导致输出张量过大；此时可先用torch.unique(t)获取实际出现的值，再结合torch.searchsorted构建紧凑映射（进阶优化）。 -
scatter_reduce_在 CUDA 上性能优异（O(N) 时间复杂度），但需 PyTorch ≥ 2.1 且 CUDA ≥ 11.8；旧版本可退化为torch.scatter_reduce_（注意 API 差异）或使用torch.bucketize+torch.cummin组合方案。
该方法兼顾简洁性、可读性与极致性能，是 PyTorch 生态中实现“首次出现索引”任务的推荐范式。
