如何将Python嵌套循环转为列表推导式/Map以优化布尔张量处理

我有一段可正常运行的Python代码，用于从布尔类型的2D张量piou中筛选出所有元素均为True的行的索引，并存入taf列表。当前代码采用嵌套for循环实现，我希望将其转换为列表推导式或map形式以提升运行速度。

我尝试了如下写法，但无法正确引入counter变量并实现计数逻辑：

taf = [i for i in range(len(piou))
       if counter == len(piou[i])
       for j in range(len(piou[i]))
       if piou[i][j] == True]

原嵌套for循环代码如下：

taf = []
for i in range(len(piou)):
    counter = 0
    for j in range(len(piou[i])):
        
        if piou[i][j] == True:
            counter = counter + 1
    if counter == len(piou[i]):
        taf.append(i)

piou张量示例：

tensor([[ True,  True,  True,  True, False, False],
        [ True,  True,  True,  True,  True,  True],
        [False, False,  True,  True,  True,  True],
        [ True,  True, False, False,  True,  True],
        [ True,  True,  True,  True,  True,  True]], device='cuda:0')

解决方案 #

1. PyTorch内置方法（最优选择，速度最快）

既然用的是PyTorch张量，直接用内置的all()方法就能高效完成判断，完全不需要手动循环：

# 沿列方向（dim=1）判断每行所有元素是否为True
all_true_rows = piou.all(dim=1)
# 获取符合条件的行索引并转成Python列表
taf = all_true_rows.nonzero(as_tuple=True)[0].tolist()

这个方法利用PyTorch的底层优化（尤其是GPU张量，能充分利用硬件加速），比任何Python层面的循环/推导式都快。

2. 列表推导式（Python原生写法）

如果一定要用列表推导式，直接用Python内置的all()函数判断整行是否全为True即可：

# 若张量在GPU上，建议先转成CPU列表避免频繁设备间拷贝：piou = piou.cpu().tolist()
taf = [i for i, row in enumerate(piou) if all(row)]

enumerate同时获取行索引和行内容，all(row)直接判断该行所有元素是否为True，逻辑简洁且效率远高于手动计数。

原写法问题说明 #

你之前的列表推导式逻辑有误：嵌套循环的顺序是先执行内层循环，再判断外层条件，没法像原循环那样先完成计数再做判断。而且列表推导式里无法维护counter这类中间变量，用all()直接替代计数判断才是正确思路。

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