本文介绍如何在 pytorch 中对二维张量每行独立、高效地获取「最大值索引」或「最小非零值索引」，依据每行指定的模式（k=1→最大，k=0→最小非零），全程避免循环，完全向量化。
本文介绍如何在 pytorch 中对二维张量每行独立、高效地获取「最大值索引」或「最小非零值索引」，依据每行指定的模式（k=1→最大，k=0→最小非零），全程避免循环，完全向量化。
在深度学习实践中，常需根据动态策略对 batch 内各行执行不同类型的极值索引检索——例如在强化学习动作选择、稀疏注意力掩码生成或自适应采样中，某些样本需取最大响应位置，另一些则需跳过零值后取最小有效值。若逐行 for 循环调用 torch.argmax() 或手动过滤，将严重损害 GPU 并行性与训练吞吐。本文提供一种纯张量操作、零 Python 循环的解决方案。
核心思想是统一建模为带符号修正的 topk 问题：
- 当 k=1（取最大）：保持原值不变，直接 topk(1)；
- 当 k=0（取最小非零）：将所有 0 替换为 -∞，并对其余值取负，使最小非零值变为最大负值，再 topk(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() # 必须为 float 类型（支持 -inf 填充）
# 每行策略：1 → 取最大值索引；0 → 取最小非零值索引
k = torch.tensor([1, 0, 0]).long()
# 步骤 1：构造符号向量 —— k=1 → +1；k=0 → -1
k_sign = torch.where(k == 1, torch.ones_like(k), -torch.ones_like(k))
# 步骤 2：按行应用符号（广播至列维度）
x_signed = x * k_sign.unsqueeze(1) # shape: (3, 5)
# 步骤 3：将原始张量中的 0 替换为 -inf（确保它们在 argmax 中被忽略）
# 注意：此处用 x == 0 判断原始零值，而非 x_signed，因符号翻转后零仍为零
x_filled = x_signed.masked_fill(x == 0, float('-inf'))
# 步骤 4：每行取 top-1 索引（返回 indices，无需 values）
_, indices = x_filled.topk(1, dim=1) # shape: (3, 1)
output = indices.squeeze(1) # 压缩为一维 tensor: [0, 2, 2]
print("输出索引:", output)
print("对应值:", x[torch.arange(x.size(0)), output])
# 输出：
# 输出索引: tensor([0, 2, 2])
# 对应值: tensor([4., 2., 3.])
✅ 关键设计说明：
- k_sign 使用 torch.where 更清晰安全，替代原文中易出错的 k + (-1 * (k==0).float())（后者在 k 为 long 时会触发隐式类型转换警告）；
- masked_fill(x == 0, -inf) 严格基于原始输入判断零值，避免因符号翻转导致逻辑错误；
- topk(1, dim=1) 天然支持 batch 维度并行，梯度可正常回传（适用于可微场景）；
- 最终 squeeze(1) 确保输出为 (batch_size,) 形状，便于后续索引使用（如 x[batch_idx, output]）。
⚠️ 注意事项：
- 输入张量必须为 float 类型（int 不支持 float('-inf')）；
- 若某行全为零，topk 将返回首个索引（因所有 -inf 相等），建议前置校验 torch.any(x != 0, dim=1)；
- 如需扩展至 k > 1（如 top-3 最大 / top-2 最小非零），只需调整 topk(k_val, dim=1) 并注意 k_val 需为标量——若每行 k 不同，需改用 torch.gather + 排序，但本例中 k 仅控制“方向”而非数量，故保持 topk(1) 即可。
该方法兼具简洁性、可读性与高性能，是 PyTorch 动态极值索引任务的标准实践范式。
