如何获取PyTorch张量每行最后一个非无穷大浮点值的列索引

已知给定张量满足强约束：每行的-inf只会连续出现在行尾，不存在-inf后接有效浮点值的结构，因此可以直接通过统计每行有效值个数快速得到目标索引，全程使用PyTorch向量化操作，无Python循环，执行效率高。

完整实现代码 #

import torch

# 示例输入张量
t = torch.tensor([[2.0000,   -float('inf'), -float('inf')],
                  [5.5000, 5.0000, -float('inf')],
                  [3.0000, 2.0000, 1.0000]])

# 计算目标索引
result = (~torch.isinf(t)).sum(dim=1).sub(1).reshape(-1, 1)

print(result)
# 输出匹配预期：
# tensor([[0],
#         [1],
#         [2]])

逻辑拆解 #

• 用torch.isinf(t)生成布尔掩码，所有无穷大值（含-inf）位置标记为True，对掩码取反~后，所有有效浮点值位置标记为True
• 沿列方向（dim=1）对布尔掩码求和，得到每行从左到右连续排列的有效值总个数
• 由于有效值全部集中在行首前缀，最后一个有效值的列索引等于「该行有效值总个数 - 1」
• 调用reshape(-1, 1)将一维索引结果转为二维列向量格式，和要求的输出形状完全一致

内容的提问来源于stack exchange，提问作者melatonin15