如何根据输入矩阵的行和条件生成对应PyTorch向量？

原代码存在的问题 #

• 求和维度错误：torch.sum(self.P, dim=0)是按列求和，要计算行和应该用dim=1。
• 逻辑嵌套错误：第二个if写在第一个if内部，当行和为0时，永远不会触发行和为1的判断。
• 缺少行遍历逻辑：变量i未定义，无法对应矩阵的行索引到向量位置。

正确实现方式 #

PyTorch支持向量化操作，无需手动循环，直接通过条件判断生成目标向量：

# 计算每一行的和，keepdim=True保持(n,1)的维度，和目标向量格式匹配
row_sums = torch.sum(self.P, dim=1, keepdim=True)

# 初始化全0向量
a_mat = torch.zeros_like(row_sums)

# 行和为0的位置设为1
a_mat[row_sums == 0] = 1
# 行和为1的位置保持0（因为初始就是0，这步可省略，若要明确逻辑也可保留）
# a_mat[row_sums == 1] = 0

如果需要兼容行和既不是0也不是1的场景（比如设置默认值），可以用torch.where实现更简洁的分支逻辑：

row_sums = torch.sum(self.P, dim=1, keepdim=True)
# 行和为0设1，行和为1设0，其他情况保持原初始值0
a_mat = torch.where(row_sums == 0, torch.tensor(1.), torch.where(row_sums == 1, torch.tensor(0.), torch.zeros_like(row_sums)))

手动循环实现（不推荐，效率较低） #

如果一定要用循环遍历每一行：

n = self.P.shape[0]
a_mat = torch.zeros(n, 1)

for i in range(n):
    row_sum = torch.sum(self.P[i, :])
    if row_sum == 0:
        a_mat[i] = 1
    elif row_sum == 1:
        a_mat[i] = 0
    # 其他行和情况可按需添加处理逻辑

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