PyTorch中矩阵运算的向量化与高效实现

import torch

m = 100
n = 100
b = torch.rand(m)
a = torch.rand(m)
summation_old = 0.0 # 使用浮点数初始化以避免类型错误
A = torch.rand(n, n)

for i in range(m):
    # 计算 A - b[i] * I
    # torch.eye(n) 创建 n x n 的单位矩阵
    matrix_term = A - b[i] * torch.eye(n)
    # 逐元素除法
    summation_old = summation_old + a[i] / matrix_term

print(f"原始循环计算结果的形状: {summation_old.shape}")

# 尝试使用列表推导式和 torch.stack
# 注意：这里我们假设 A 和 b, a 已经定义如上
# A = torch.rand(n, n)
# b = torch.rand(m)
# a = torch.rand(m)

# 这种方法虽然避免了显式循环求和，但列表推导式本身仍然是Python循环
# 并且在内存上可能需要先构建一个完整的中间张量堆栈
stacked_results = torch.stack([a[i] / (A - b[i] * torch.eye(n)) for i in range(m)], dim=0)
summation_stacked = torch.sum(stacked_results, dim=0)

# 验证结果（注意：由于浮点数精度，直接 == 比较通常会失败）
# print(f"堆叠向量化计算结果的形状: {summation_stacked.shape}")
# print(f"堆叠向量化结果与原始结果是否完全相等: {(summation_stacked == summation_old).all()}")

import torch

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

# 1. 创建批次化的 b_i * I 矩阵
# torch.eye(n) 生成 (n, n) 的单位矩阵
identity_matrix = torch.eye(n) # 形状: (n, n)
# unsqueeze(0) 将 identity_matrix 变为 (1, n, n)，为广播做准备
# b.unsqueeze(1).unsqueeze(2) 将 b 变为 (m, 1, 1)，使其能与 (1, n, n) 广播
# 结果 B 的形状为 (m, n, n)，其中 B[i, :, :] = b[i] * identity_matrix
B_batch = identity_matrix.unsqueeze(0) * b.unsqueeze(1).unsqueeze(2)

# 2. 执行 A - b_i * I 操作
# A.unsqueeze(0) 将 A 变为 (1, n, n)，使其能与 (m, n, n) 的 B_batch 广播
# 结果 A_minus_B 的形状为 (m, n, n)，其中 A_minus_B[i, :, :] = A - b[i] * I
A_minus_B = A.unsqueeze(0) - B_batch

# 3. 执行 a_i / (A - b_i * I) 操作
# a.unsqueeze(1).unsqueeze(2) 将 a 变为 (m, 1, 1)，使其能与 (m, n, n) 的 A_minus_B 广播
# 结果 term_batch 的形状为 (m, n, n)，其中 term_batch[i, :, :] = a[i] / (A - b[i] * I)
term_batch = a.unsqueeze(1).unsqueeze(2) / A_minus_B

# 4. 沿批次维度求和
# torch.sum(..., dim=0) 将 (m, n, n) 的张量沿第一个维度（批次维度）求和
# 最终结果 summation_new 的形状为 (n, n)
summation_new = torch.sum(term_batch, dim=0)

print(f"向量化计算结果的形状: {summation_new.shape}")

# 假设 summation_old 和 summation_new 已经通过上述方法计算得到

# 验证两个结果是否在数值上接近
is_close = torch.allclose(summation_old, summation_new)
print(f"原始循环结果与向量化结果在数值上是否接近: {is_close}")

# 可以通过设置 rtol (相对容忍度) 和 atol (绝对容忍度) 来调整比较的严格性
# is_close_strict = torch.allclose(summation_old, summation_new, rtol=1e-05, atol=1e-08)
# print(f"在更严格的容忍度下是否接近: {is_close_strict}")