PyTorch中对批量3×3布尔张量执行按位或运算的最优方法

针对你的需求，这里有两种简洁高效的实现方式，都是PyTorch原生优化的操作，速度拉满：

方法1：利用逻辑或等价性（推荐，最简洁） #

因为你的张量元素只有0和1，按位或运算等价于判断每个位置是否存在至少一个1，直接用torch.any()在第0维（对应n个矩阵的维度）上做归约，再转成整数类型即可：

import torch

# 构造示例输入张量
input_tensor = torch.tensor([
    [[0,0,1], [0,0,0], [1,0,0]],
    [[1,0,0], [0,0,0], [1,0,1]],
    [[0,1,1], [0,1,0], [1,0,1]]
])

# 核心操作
result = input_tensor.any(dim=0).int()

方法2：直接按位或归约 #

用PyTorch专门的按位归约函数torch.bitwise_or.reduce()，直接对指定维度的所有矩阵依次执行按位或：

result = torch.bitwise_or.reduce(input_tensor, dim=0)

结果验证 #

两种方法都会得到你期望的输出：

tensor([[1, 1, 1],
        [0, 1, 0],
        [1, 0, 1]], dtype=torch.int32)

新手注意点 #

• dim=0是指定要合并的维度（对应n个矩阵的那一层），如果你的张量维度顺序不同，需要调整这个参数；
• 如果你的张量是float类型（比如元素是0.0/1.0），torch.any()会自动把非0值识别为True，转int后同样得到正确结果；
• 两种方法的效率几乎无差别，都是PyTorch底层优化的操作，适合大规模n的场景。

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