本文介绍如何在 PyTorch 中对二维张量每行向量化地获取 Top-k 索引：k 值按行指定（1 表示取最大值索引，0 表示取非零最小值索引），全程避免循环，充分利用 topk 与掩码操作。
本文介绍如何在 pytorch 中对二维张量每行**向量化地获取 top-k 索引**：k 值按行指定（1 表示取最大值索引，0 表示取非零最小值索引），全程避免循环，充分利用 `topk` 与掩码操作。
在实际深度学习任务中（如动态采样、自适应注意力、稀疏路由等），常需对 batch 维度中每一行独立执行条件化极值索引提取：例如，对每条样本以概率 0.7 选取最大值位置，以 0.3 概率选取非零元素中的最小值位置（零值被显式忽略）。若使用 Python 循环逐行处理，不仅低效，更破坏计算图连续性，无法高效反向传播。因此，必须采用完全向量化（vectorized）方案。
核心思路是统一为 topk 操作建模：
- 当需取“最大值”时，保持原张量符号不变；
- 当需取“非零最小值”时，将原张量符号翻转（即 -x），并把所有 0 替换为 -inf（确保其在 topk 中被彻底排除）；
- 最终对统一处理后的张量调用 .topk(1, dim=1)，即可一次性获得所有行的目标索引。
以下为完整可运行代码：
import torch
# 示例输入：batch_size=3, N=5
x = torch.tensor([[4, 3, 1, 4, 2],
[0, 0, 2, 3, 4],
[4, 4, 3, 0, 3]]).float() # 注意：topk 要求 float 类型
# k 向量：1 → 取最大值索引；0 → 取非零最小值索引
k = torch.tensor([1, 0, 0]).long()
# 步骤 1：构造符号向量 —— k=1 → +1；k=0 → -1
k_sign = 2 * k - 1 # 更简洁写法：[1,0,0] → [1,-1,-1]
# 步骤 2：按行翻转符号（仅对 k=0 的行）
x_signed = x * k_sign.unsqueeze(1) # 自动广播至 (3,5)
# 步骤 3：将原始张量中所有 0 替换为 -inf（k=1 时不影响最大值；k=0 时确保 0 不参与最小值竞争）
x_filled = x_signed.masked_fill(x == 0, float('-inf'))
# 步骤 4：统一执行 topk（k=1），返回索引
_, indices = x_filled.topk(1, dim=1) # shape: (3,1)
output = indices.squeeze(1) # → tensor([0, 2, 2])
print("目标索引:", output)
print("对应值:", x.gather(1, output.unsqueeze(1)).squeeze(1))
# 输出：
# 目标索引: tensor([0, 2, 2])
# 对应值: tensor([4., 2., 3.])✅ 关键要点说明：
- x.float() 是必需的：topk 不支持整数张量（除非是 torch.int64 且设备为 CUDA，但行为不稳定，强烈建议统一转 float）；
- masked_fill(x == 0, -inf) 必须基于原始 x 判断零值，而非 x_signed，否则符号翻转后 -0.0 == 0.0 仍成立，但逻辑更清晰；
- k_sign = 2*k - 1 比原答案中的 k + (-1*(k==0).float()) 更简洁、无类型转换开销；
- .unsqueeze(1) 和自动广播确保行级操作精准生效，是向量化的核心技巧；
- 若需扩展至 k > 1（如每行取前 2 个非零最小值），只需将 topk(1, ...) 改为 topk(k_val, ...) 并适配 k 向量维度（需 torch.nn.functional.topk 配合 dim=1 和 largest=True/False 动态控制）。
该方法时间复杂度为 $O(BN\log K)$（$K=1$ 时近似线性），内存友好，完全兼容 torch.compile 与 DataParallel/DistributedDataParallel，适用于高吞吐训练场景。
