如何无需重塑3D张量执行矩阵乘法？含批量矩阵操作问询

一、无需重塑为1980×1980张量的矩阵乘法方案 #

你不需要把整个[330,330,36]张量reshape成大矩阵再计算，因为你只需要b@b对角线的第1-5个元素，用**爱因斯坦求和（einsum）**可以直接定位计算目标，效率更高：

1. 先把a从[330,330,36]拆分为[330,330,6,6]（这只是拆分36为6×6，不是拉平成1980×1980，完全符合你的需求）；
2. 利用einsum直接计算目标对角线元素的乘积和：

import torch

# 假设a是你的[330,330,36]张量
a_4d = a.reshape(330, 330, 6, 6)
# 计算b@b对角线中对应第一个原子轨道1-5的元素
# 原理：sum_j sum_q a[0,j,p,q] * a[j,0,q,p] （p从0到5对应轨道0-5）
diag_vals = torch.einsum('j p q, j q p -> p', a_4d[0], a_4d[:, 0])
# 取第1到第5个元素求和
final_result = diag_vals[1:6].sum()

这个操作和原逻辑等价：原操作中的b是把原子-轨道对拉平为行/列，b@b的对角线元素本质是轨道p（属于原子i）经过所有中间原子j的轨道q，再回到自身的累积相互作用；用einsum直接针对原子0的轨道1-5计算，跳过了对整个大矩阵的构建，节省内存和计算量。

二、矩阵列表的单条命令乘法方案 #

分两种常见场景处理：

1. 列表中矩阵依次相乘（链式乘法） #

用torch.linalg.multi_dot可以直接对列表中的矩阵按顺序做矩阵乘法：

mat_list = [torch.randn(2,2), torch.randn(2,2), torch.randn(2,2)]
# 等价于 mat_list[0] @ mat_list[1] @ mat_list[2]
chain_result = torch.linalg.multi_dot(mat_list)

2. 批量矩阵对应相乘（每个矩阵独立运算） #

如果列表里是同形状的矩阵，需要每个矩阵和另一个列表的对应矩阵相乘，先把列表转成三维张量，再用torch.bmm做批量矩阵乘法：

# 两个各含5个3×3矩阵的列表
mat_list1 = [torch.randn(3,3) for _ in range(5)]
mat_list2 = [torch.randn(3,3) for _ in range(5)]

# 转成批量张量后做乘法
stacked1 = torch.stack(mat_list1)  # shape [5,3,3]
stacked2 = torch.stack(mat_list2)  # shape [5,3,3]
batch_result = torch.bmm(stacked1, stacked2)  # shape [5,3,3]

如果是列表中每个矩阵和同一个矩阵相乘，一行列表推导式就能完成：

single_mat = torch.randn(3,3)
result_list = [mat @ single_mat for mat in mat_list1]

内容的提问来源于stack exchange，提问作者Shih-Han Wang