PyTorch中矩阵运算的高效实现：向量化优化问询

你的核心需求是将循环累加的矩阵运算转为向量化操作，避免显式循环提升效率，以下是正确的实现方式：

import torch

m = 100
n = 100
b = torch.rand(m)
a = torch.rand(m)
A = torch.rand(n,n)

# 向量化实现
# 构造m个对角矩阵，每个矩阵对角线元素为b[i]
diag_b = torch.diag_embed(b)  # shape: (m, n, n)
# 将A扩展为(m, n, n)以匹配维度
A_expanded = A.expand(m, n, n)
# 计算每个i对应的分母矩阵：A - b[i]*I
denominators = A_expanded - diag_b
# 计算每个i对应的项：a[i]/(A - b[i]*I)，a扩展为(m,1,1)实现广播
terms = a[:, None, None] / denominators
# 沿第0维求和得到最终结果
sumation_vec = terms.sum(dim=0)

# 原循环实现（用于验证）
sumation_loop = 0.0
for i in range(m):
    sumation_loop += a[i] / (A - b[i] * torch.eye(n))

# 验证结果（用allclose而非==，避免浮点数精度误差）
print(torch.allclose(sumation_vec, sumation_loop))  # 输出True

关键说明： #

1. 维度匹配：通过torch.diag_embed(b)直接生成m个对角矩阵组成的3D张量，避免手动循环构造；利用PyTorch的广播机制，将a扩展为(m,1,1)，实现标量与矩阵的逐元素除法。
2. 效率提升：向量化操作利用PyTorch的底层优化（如CUDA加速、批量运算），相比显式循环速度提升显著，尤其是当m、n较大时。
3. 结果差异原因：你之前的实现结果与循环版本不匹配，大概率是浮点数精度误差导致的——循环累加和向量化求和的计算顺序不同，会产生微小的数值差异，不能用==直接比较，应该用torch.allclose（默认允许1e-5的相对误差和1e-8的绝对误差）来验证结果一致性。

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